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

@@ -113,7 +113,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,
             input_names,
         ) {
@@ -154,6 +154,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_run(
@@ -172,7 +175,7 @@ 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,
             input_names,
         )?;
@@ -205,16 +208,9 @@ impl<T: ResourceTracker> MontyRepl<T> {
         self.globals = vm.take_globals();
         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.
-        let ReplExecutor {
-            name_map,
-            interns,
-            code,
-            ..
-        } = executor;
-        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.
+        let ReplExecutor { interns, code, .. } = executor;
         self.interns = interns;
 
         result.map_err(|e| e.into_python_exception(&self.interns, &code))
@@ -803,22 +799,18 @@ pub fn detect_repl_continuation_mode(source: &str) -> ReplContinuationMode {
 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_code: Code,
     /// Interned strings and compiled functions for this snippet.
     interns: Interns,
     /// Source code used for traceback/error rendering.
     code: String,
-    /// Input variable names that were injected for this snippet.
+    /// Global slots assigned to inputs for this snippet.
     ///
-    /// Stored so that `inject_inputs` can look up their namespace slots
-    /// after compilation assigns them.
-    input_names: Vec<String>,
+    /// Inputs are pre-registered in the REPL-global name map before preparation.
+    /// We store only the assigned slots here so execution can inject values without
+    /// needing a cloned copy of the full global name map.
+    input_slots: Vec<usize>,
 }
 
 impl ReplExecutor {
@@ -828,23 +820,29 @@ 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.
     ///
-    /// `input_names` are pre-registered in the name map before preparation so they
+    /// `input_names` are pre-registered in the REPL-global name map before preparation so they
     /// receive stable namespace slots that `inject_inputs` can use to store values.
+    ///
+    /// If snippet preparation fails, newly reserved global slots remain in the
+    /// caller-owned name map. That is harmless because the corresponding runtime
+    /// global entries stay `Undefined` until a later successful assignment.
     fn new_repl_snippet(
         code: String,
         script_name: &str,
-        mut existing_name_map: AHashMap<String, NamespaceId>,
+        existing_name_map: &mut AHashMap<String, NamespaceId>,
         existing_interns: &Interns,
         input_names: Vec<String>,
     ) -> Result<Self, MontyException> {
         // Pre-register input names so they get stable slots before preparation.
+        let mut input_slots = Vec::with_capacity(input_names.len());
         for name in &input_names {
             let next_slot = existing_name_map.len();
-            existing_name_map
+            let slot = *existing_name_map
                 .entry(name.clone())
                 .or_insert_with(|| NamespaceId::new(next_slot));
+            input_slots.push(slot.index());
         }
 
         let seeded_interner = InternerBuilder::from_interns(existing_interns, &code);
@@ -863,11 +861,10 @@ impl ReplExecutor {
 
         Ok(Self {
             namespace_size: prepared.namespace_size,
-            name_map: prepared.name_map,
             module_code: compile_result.code,
             interns,
             code,
-            input_names,
+            input_slots,
         })
     }
 }
@@ -965,12 +962,7 @@ fn inject_inputs_into_vm(
     input_values: Vec<MontyObject>,
     vm: &mut VM<'_, '_, impl ResourceTracker>,
 ) -> Result<(), MontyException> {
-    for (name, obj) in executor.input_names.iter().zip(input_values) {
-        let slot = executor
-            .name_map
-            .get(name)
-            .expect("input name should have a namespace slot")
-            .index();
+    for (slot, obj) in executor.input_slots.iter().copied().zip(input_values) {
         let value = obj
             .to_value(vm)
             .map_err(|e| MontyException::runtime_error(format!("invalid input type: {e}")))?;
@@ -1048,8 +1040,7 @@ fn build_repl_progress<T: ResourceTracker>(
 
     match converted {
         ConvertedExit::Complete(obj) => {
-            let ReplExecutor { name_map, interns, .. } = executor;
-            repl.global_name_map = name_map;
+            let ReplExecutor { interns, .. } = executor;
             repl.interns = interns;
             Ok(ReplProgress::Complete { repl, value: obj })
         }
@@ -1097,11 +1088,9 @@ fn build_repl_progress<T: ResourceTracker>(
         })),
         ConvertedExit::Error(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

@@ -81,6 +81,29 @@ fn repl_runtime_error_keeps_partial_state_consistent() {
     assert_eq!(feed_run_print(&mut repl, "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 = feed_run_print(&mut repl, "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!(
+        feed_run_print(&mut repl, "reserved_name = 42").unwrap(),
+        MontyObject::None
+    );
+    assert_eq!(
+        feed_run_print(&mut repl, "reserved_name").unwrap(),
+        MontyObject::Int(42)
+    );
+}
+
 #[test]
 fn repl_heap_mutations_are_not_replayed() {
     let (mut repl, _) = init_repl("items = []");
@@ -260,6 +283,30 @@ fn repl_start_runtime_error_preserves_repl_state() {
     assert_eq!(feed_run_print(&mut repl, "x + y + 12").unwrap(), MontyObject::Int(42));
 }
 
+#[test]
+fn repl_start_compile_error_keeps_reserved_global_slot() {
+    let (repl, _) = init_repl("");
+
+    let err = repl
+        .feed_start("reserved_name = 1\nnonlocal missing_name", vec![], 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!(
+        feed_run_print(&mut repl, "reserved_name = 42").unwrap(),
+        MontyObject::None
+    );
+    assert_eq!(
+        feed_run_print(&mut repl, "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