Avoid cloning REPL global name maps

Author: tristanang

crates/monty/src/prepare.rs

@@ -1,4 +1,4 @@
-use std::collections::hash_map::Entry;
+use std::{collections::hash_map::Entry, mem};
 
 use ahash::{AHashMap, AHashSet};
 
@@ -30,7 +30,7 @@ pub struct PrepareResult {
     ///
     /// This map is used by:
     /// - ref-count tests for looking up variables by name
-    /// - REPL incremental compilation to preserve stable global slot IDs across snippets
+    /// - initial REPL compilation to seed stable global slot IDs across later snippets
     pub name_map: AHashMap<String, NamespaceId>,
     /// The prepared AST nodes with all names resolved to namespace indices.
     /// Function definitions are inline as `PreparedFunctionDef` variants.
@@ -39,6 +39,20 @@ pub struct PrepareResult {
     pub interner: InternerBuilder,
 }
 
+/// Result of preparing a REPL snippet against an existing mutable global name map.
+///
+/// Unlike [`PrepareResult`], this does not return the module name map because the
+/// caller-owned map is mutated in place during preparation. On failure, any newly
+/// reserved global slots remain in that map so later REPL snippets can reuse them.
+pub struct PrepareExistingNamesResult {
+    /// Number of items in the namespace (at module level, this IS the global namespace)
+    pub namespace_size: usize,
+    /// The prepared AST nodes with all names resolved to namespace indices.
+    pub nodes: Vec<PreparedNode>,
+    /// The string interner containing all interned identifiers and filenames.
+    pub interner: InternerBuilder,
+}
+
 /// Prepares parsed nodes for compilation by resolving names and building the initial namespace.
 ///
 /// The namespace will be converted to runtime Objects when execution begins and the heap is available.
@@ -69,15 +83,22 @@ pub(crate) fn prepare(parse_result: ParseResult, input_names: Vec<String>) -> Re
 
 /// Prepares parsed nodes for REPL-style incremental compilation using an existing global namespace map.
 ///
-/// Existing bindings keep their original namespace slots; any new names are appended with new slots.
-/// This ensures snippets can be compiled independently while sharing one persistent global namespace.
+/// Existing bindings keep their original namespace slots; any new names are appended with new slots
+/// directly in `existing_name_map`. This ensures snippets can be compiled independently while sharing
+/// one persistent global namespace.
+///
+/// If preparation fails, any newly reserved slots stay in `existing_name_map`. That is intentional for
+/// REPL compilation: the corresponding namespace entries remain `Undefined`, so preserving the slots is
+/// harmless and avoids cloning the full map before each feed.
 pub(crate) fn prepare_with_existing_names(
     parse_result: ParseResult,
-    existing_name_map: AHashMap<String, NamespaceId>,
-) -> Result<PrepareResult, ParseError> {
+    existing_name_map: &mut AHashMap<String, NamespaceId>,
+) -> Result<PrepareExistingNamesResult, ParseError> {
     let ParseResult { nodes, interner } = parse_result;
     let mut p = Prepare::new_module_with_name_map(existing_name_map, &interner);
-    let mut prepared_nodes = p.prepare_nodes(nodes)?;
+    let prepare_result = p.prepare_nodes(nodes);
+    p.write_back_module_name_map(existing_name_map);
+    let mut prepared_nodes = prepare_result?;
 
     // In the root frame, the last expression is implicitly returned to match REPL behavior.
     if let Some(Node::Expr(expr_loc)) = prepared_nodes.last()
@@ -88,9 +109,8 @@ pub(crate) fn prepare_with_existing_names(
         prepared_nodes.push(Node::Return(new_expr_loc));
     }
 
-    Ok(PrepareResult {
+    Ok(PrepareExistingNamesResult {
         namespace_size: p.namespace_size,
-        name_map: p.name_map,
         nodes: prepared_nodes,
         interner,
     })
@@ -184,10 +204,13 @@ impl<'i> Prepare<'i> {
         }
     }
 
-    /// Creates a module-scope Prepare instance from an existing global name map.
+    /// Creates a module-scope Prepare instance from an existing mutable global name map.
     ///
     /// Used by incremental REPL compilation to keep stable slot assignments across snippets.
-    fn new_module_with_name_map(name_map: AHashMap<String, NamespaceId>, interner: &'i InternerBuilder) -> Self {
+    /// The caller's map is moved into this preparer and must be written back with
+    /// [`Self::write_back_module_name_map`] before returning.
+    fn new_module_with_name_map(name_map: &mut AHashMap<String, NamespaceId>, interner: &'i InternerBuilder) -> Self {
+        let name_map = mem::take(name_map);
         let namespace_size = name_map
             .values()
             .map(|id| id.index())
@@ -210,6 +233,19 @@ impl<'i> Prepare<'i> {
         }
     }
 
+    /// Writes the mutated module-level name map back to the caller-owned REPL state.
+    ///
+    /// REPL preparation takes ownership of the map so new slots can be appended without cloning.
+    /// Call this before returning from the incremental prepare path, even on errors, so newly
+    /// reserved slots remain visible to later snippets.
+    fn write_back_module_name_map(&mut self, target_name_map: &mut AHashMap<String, NamespaceId>) {
+        debug_assert!(
+            self.is_module_scope,
+            "module name maps can only be written back from module scope"
+        );
+        mem::swap(target_name_map, &mut self.name_map);
+    }
+
     /// Creates a new Prepare instance for function-level code.
     ///
     /// Pre-populates `free_var_map` with nonlocal declarations and implicit captures,

crates/monty/src/repl.rs

@@ -37,11 +37,6 @@ use crate::{
 struct ReplExecutor {
     /// Number of slots needed in the global namespace.
     namespace_size: usize,
-    /// Maps variable names to their indices in the namespace.
-    ///
-    /// Stable slot assignment is required across snippets so previously created
-    /// objects continue to resolve names correctly.
-    name_map: AHashMap<String, NamespaceId>,
     /// Compiled bytecode for the snippet/module.
     module_code: Code,
     /// Interned strings and compiled functions for this snippet/module.
@@ -55,7 +50,11 @@ impl ReplExecutor {
     ///
     /// This is equivalent to normal module compilation but scoped to REPL
     /// infrastructure so `run.rs` can remain unchanged.
-    fn new(code: String, script_name: &str, input_names: Vec<String>) -> Result<Self, MontyException> {
+    fn new(
+        code: String,
+        script_name: &str,
+        input_names: Vec<String>,
+    ) -> Result<(Self, AHashMap<String, NamespaceId>), MontyException> {
         let parse_result = parse(&code, script_name).map_err(|e| e.into_python_exc(script_name, &code))?;
         let prepared = prepare(parse_result, input_names).map_err(|e| e.into_python_exc(script_name, &code))?;
 
@@ -65,13 +64,15 @@ impl ReplExecutor {
             .map_err(|e| e.into_python_exc(script_name, &code))?;
         interns.set_functions(compile_result.functions);
 
-        Ok(Self {
-            namespace_size: prepared.namespace_size,
-            name_map: prepared.name_map,
-            module_code: compile_result.code,
-            interns,
-            code,
-        })
+        Ok((
+            Self {
+                namespace_size: prepared.namespace_size,
+                module_code: compile_result.code,
+                interns,
+                code,
+            },
+            prepared.name_map,
+        ))
     }
 
     /// Compiles one incremental REPL snippet against existing session metadata.
@@ -80,11 +81,15 @@ impl ReplExecutor {
     /// no-replay execution:
     /// - Seeds parsing from `existing_interns` so old `StringId` values stay stable.
     /// - Seeds compilation with existing functions so old `FunctionId` values remain valid.
-    /// - Reuses `existing_name_map` and appends new global names only.
+    /// - Reuses `existing_name_map` in place and appends new global names only.
+    ///
+    /// If snippet preparation fails, newly reserved global slots remain in the
+    /// caller-owned name map. That is harmless because the corresponding runtime
+    /// namespace entries stay `Undefined` until a later successful assignment.
     fn new_repl_snippet(
         code: String,
         script_name: &str,
-        existing_name_map: AHashMap<String, NamespaceId>,
+        existing_name_map: &mut AHashMap<String, NamespaceId>,
         existing_interns: &Interns,
     ) -> Result<Self, MontyException> {
         let seeded_interner = InternerBuilder::from_interns(existing_interns, &code);
@@ -103,7 +108,6 @@ impl ReplExecutor {
 
         Ok(Self {
             namespace_size: prepared.namespace_size,
-            name_map: prepared.name_map,
             module_code: compile_result.code,
             interns,
             code,
@@ -282,7 +286,7 @@ impl<T: ResourceTracker> MontyRepl<T> {
         resource_tracker: T,
         print: &mut PrintWriter<'_>,
     ) -> Result<(Self, MontyObject), MontyException> {
-        let executor = ReplExecutor::new(code, script_name, input_names)?;
+        let (executor, global_name_map) = ReplExecutor::new(code, script_name, input_names)?;
 
         let mut heap = Heap::new(executor.namespace_size, resource_tracker);
         let mut namespaces = executor.prepare_namespaces(inputs, &mut heap)?;
@@ -306,7 +310,7 @@ impl<T: ResourceTracker> MontyRepl<T> {
         let repl = Self {
             script_name: script_name.to_owned(),
             next_input_id: 0,
-            global_name_map: executor.name_map,
+            global_name_map,
             interns: executor.interns,
             heap,
             namespaces,
@@ -345,7 +349,7 @@ impl<T: ResourceTracker> MontyRepl<T> {
         let executor = match ReplExecutor::new_repl_snippet(
             code.to_owned(),
             &input_script_name,
-            this.global_name_map.clone(),
+            &mut this.global_name_map,
             &this.interns,
         ) {
             Ok(exec) => exec,
@@ -377,6 +381,9 @@ impl<T: ResourceTracker> MontyRepl<T> {
     /// partially mutating globals, those mutations remain visible in later feeds,
     /// matching Python REPL semantics.
     ///
+    /// Compile-time failures may still reserve new global slots in the REPL map.
+    /// Those slots remain `Undefined` until a later successful snippet writes them.
+    ///
     /// # Errors
     /// Returns `MontyException` for syntax/compile/runtime failures.
     pub fn feed(&mut self, code: &str, print: &mut PrintWriter<'_>) -> Result<MontyObject, MontyException> {
@@ -388,13 +395,12 @@ impl<T: ResourceTracker> MontyRepl<T> {
         let executor = ReplExecutor::new_repl_snippet(
             code.to_owned(),
             &input_script_name,
-            self.global_name_map.clone(),
+            &mut self.global_name_map,
             &self.interns,
         )?;
 
         let ReplExecutor {
             namespace_size,
-            name_map,
             module_code,
             interns,
             code,
@@ -417,10 +423,8 @@ impl<T: ResourceTracker> MontyRepl<T> {
 
         vm.cleanup();
 
-        // Commit compiler metadata even on runtime errors.
-        // Snippets can mutate globals before raising, and those values may contain
-        // FunctionId/StringId values that must be interpreted with the updated tables.
-        self.global_name_map = name_map;
+        // Commit intern/function metadata even on runtime errors. The REPL global
+        // name map was already updated in place during preparation.
         self.interns = interns;
 
         frame_exit_to_object(frame_exit_result, &mut self.heap, &self.interns)
@@ -976,8 +980,7 @@ fn handle_repl_vm_result<T: ResourceTracker>(
     match result {
         Ok(FrameExit::Return(value)) => {
             let output = MontyObject::new(value, &mut repl.heap, &executor.interns);
-            let ReplExecutor { name_map, interns, .. } = executor;
-            repl.global_name_map = name_map;
+            let ReplExecutor { interns, .. } = executor;
             repl.interns = interns;
             Ok(ReplProgress::Complete { repl, value: output })
         }
@@ -1055,11 +1058,9 @@ fn handle_repl_vm_result<T: ResourceTracker>(
         }
         Err(err) => {
             let error = err.into_python_exception(&executor.interns, &executor.code);
-            // Commit compiler metadata even on runtime errors, matching feed() behavior.
-            // Snippets can create new variables or functions before raising, and those
-            // values may reference FunctionId/StringId values from the new tables.
-            let ReplExecutor { name_map, interns, .. } = executor;
-            repl.global_name_map = name_map;
+            // Commit intern/function metadata even on runtime errors, matching feed() behavior.
+            // The global name map was already updated in place during preparation.
+            let ReplExecutor { interns, .. } = executor;
             repl.interns = interns;
             Err(Box::new(ReplStartError { repl, error }))
         }

crates/monty/tests/repl.rs

@@ -82,6 +82,24 @@ fn repl_runtime_error_keeps_partial_state_consistent() {
     assert_eq!(repl.feed_no_print("x").unwrap(), MontyObject::Int(1));
 }
 
+#[test]
+fn repl_feed_compile_error_keeps_reserved_global_slot() {
+    let (mut repl, init_output) = init_repl("");
+    assert_eq!(init_output, MontyObject::None);
+
+    let error = repl
+        .feed_no_print("reserved_name = 1\nnonlocal missing_name")
+        .expect_err("snippet should raise a prepare-time SyntaxError");
+    assert_eq!(error.exc_type(), monty::ExcType::SyntaxError);
+    assert_eq!(
+        error.message(),
+        Some("nonlocal declaration not allowed at module level")
+    );
+
+    assert_eq!(repl.feed_no_print("reserved_name = 42").unwrap(), MontyObject::None);
+    assert_eq!(repl.feed_no_print("reserved_name").unwrap(), MontyObject::Int(42));
+}
+
 #[test]
 fn repl_heap_mutations_are_not_replayed() {
     let (mut repl, _) = init_repl("items = []");
@@ -255,6 +273,24 @@ fn repl_start_runtime_error_preserves_repl_state() {
     assert_eq!(repl.feed_no_print("x + y + 12").unwrap(), MontyObject::Int(42));
 }
 
+#[test]
+fn repl_start_compile_error_keeps_reserved_global_slot() {
+    let (repl, _) = init_repl("");
+
+    let err = repl
+        .start("reserved_name = 1\nnonlocal missing_name", &mut PrintWriter::Stdout)
+        .expect_err("expected ReplStartError");
+    let ReplStartError { mut repl, error } = *err;
+    assert_eq!(error.exc_type(), monty::ExcType::SyntaxError);
+    assert_eq!(
+        error.message(),
+        Some("nonlocal declaration not allowed at module level")
+    );
+
+    assert_eq!(repl.feed_no_print("reserved_name = 42").unwrap(), MontyObject::None);
+    assert_eq!(repl.feed_no_print("reserved_name").unwrap(), MontyObject::Int(42));
+}
+
 #[test]
 fn repl_start_runtime_error_during_external_call_preserves_repl_state() {
     // An external function returns an error, which should come back as ReplStartError