Add localspluskinds, unify DEREF to localsplus index

- Add CO_FAST_LOCAL/CELL/FREE/HIDDEN constants and localspluskinds field to CodeObject for per-slot metadata - Change DEREF instruction opargs from cell-relative indices (NameIdx) to localsplus absolute indices (oparg::VarNum) - Add fixup_deref_opargs pass in ir.rs to convert cell-relative indices to localsplus indices after finalization - Replace get_cell_name with get_localsplus_name in InstrDisplayContext trait - Update VM cell_ref/get_cell_contents/set_cell_contents to use localsplus indices directly (no nlocals offset) - Update function.rs cell2arg, super.rs __class__ lookup with explicit nlocals offsets
RustPython · youknowone · Mar 21, 2026 · Mar 12, 2026 · Mar 17, 2026 · Mar 17, 2026
commit 6723f2fa060ed2c95452241c5021e8a56e87f90f
diff --git a/Lib/_opcode_metadata.py b/Lib/_opcode_metadata.py
diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs
@@ -692,26 +692,28 @@ impl Compiler {
             .expect("symbol_table_stack is empty! This is a compiler bug.")
     }
 
-    /// Get the index of a free variable.
-    fn get_free_var_index(&mut self, name: &str) -> CompileResult<u32> {
+    /// Get the cell-relative index of a free variable.
+    /// Returns ncells + freevar_idx. Fixed up to localsplus index during finalize.
+    fn get_free_var_index(&mut self, name: &str) -> CompileResult<oparg::VarNum> {
         let info = self.code_stack.last_mut().unwrap();
         let idx = info
             .metadata
             .freevars
             .get_index_of(name)
             .unwrap_or_else(|| info.metadata.freevars.insert_full(name.to_owned()).0);
-        Ok((idx + info.metadata.cellvars.len()).to_u32())
+        Ok((idx + info.metadata.cellvars.len()).to_u32().into())
     }
 
-    /// Get the index of a cell variable.
-    fn get_cell_var_index(&mut self, name: &str) -> CompileResult<u32> {
+    /// Get the cell-relative index of a cell variable.
+    /// Returns cellvar_idx. Fixed up to localsplus index during finalize.
+    fn get_cell_var_index(&mut self, name: &str) -> CompileResult<oparg::VarNum> {
         let info = self.code_stack.last_mut().unwrap();
         let idx = info
             .metadata
             .cellvars
             .get_index_of(name)
             .unwrap_or_else(|| info.metadata.cellvars.insert_full(name.to_owned()).0);
-        Ok(idx.to_u32())
+        Ok(idx.to_u32().into())
     }
 
     /// Get the index of a local variable.
@@ -1149,6 +1151,19 @@ impl Compiler {
             self.set_qualname();
         }
 
+        // Emit COPY_FREE_VARS and MAKE_CELL prolog before RESUME
+        {
+            let nfrees = self.code_stack.last().unwrap().metadata.freevars.len();
+            if nfrees > 0 {
+                emit!(self, Instruction::CopyFreeVars { n: nfrees as u32 });
+            }
+            let ncells = self.code_stack.last().unwrap().metadata.cellvars.len();
+            for i in 0..ncells {
+                let i_varnum: oparg::VarNum = (i as u32).into();
+                emit!(self, Instruction::MakeCell { i: i_varnum });
+            }
+        }
+
         // Emit RESUME (handles async preamble and module lineno 0)
         // CPython: LOCATION(lineno, lineno, 0, 0), then loc.lineno = 0 for module
         self.emit_resume_for_scope(scope_type, lineno);
@@ -8035,16 +8050,15 @@ impl Compiler {
 
         // Step 2: Save local variables that will be shadowed by the comprehension.
         // For each variable, we push the fast local value via LoadFastAndClear.
-        // For CELL variables, we additionally push the cell content via MakeCell
-        // (which saves the old cell value and clears the cell for the comprehension).
-        // Track cell indices to restore them later.
-        let mut cell_indices: Vec<Option<u32>> = Vec::new();
+        // For merged CELL variables, LoadFastAndClear saves the cell object from
+        // the merged slot, and MAKE_CELL creates a new empty cell in-place.
+        // MAKE_CELL has no stack effect (operates only on fastlocals).
         let mut total_stack_items: usize = 0;
         for name in &pushed_locals {
             let var_num = self.varname(name)?;
             emit!(self, Instruction::LoadFastAndClear { var_num });
             total_stack_items += 1;
-            // If the comp symbol is CELL, emit MAKE_CELL to save cell value
+            // If the comp symbol is CELL, emit MAKE_CELL to create fresh cell
             if let Some(comp_sym) = comp_table.symbols.get(name) {
                 if comp_sym.scope == SymbolScope::Cell {
                     let i = if self
@@ -8058,13 +8072,7 @@ impl Compiler {
                         self.get_cell_var_index(name)?
                     };
                     emit!(self, Instruction::MakeCell { i });
-                    cell_indices.push(Some(i));
-                    total_stack_items += 1;
-                } else {
-                    cell_indices.push(None);
                 }
-            } else {
-                cell_indices.push(None);
             }
         }
 
@@ -8192,10 +8200,7 @@ impl Compiler {
                     i: u32::try_from(total_stack_items + 1).unwrap()
                 }
             );
-            for (name, cell_idx) in pushed_locals.iter().rev().zip(cell_indices.iter().rev()) {
-                if let Some(i) = cell_idx {
-                    emit!(self, Instruction::RestoreCell { i: *i });
-                }
+            for name in pushed_locals.iter().rev() {
                 let var_num = self.varname(name)?;
                 emit!(self, Instruction::StoreFast { var_num });
             }
@@ -8216,11 +8221,8 @@ impl Compiler {
             );
         }
 
-        // Restore saved locals and cell values
-        for (name, cell_idx) in pushed_locals.iter().rev().zip(cell_indices.iter().rev()) {
-            if let Some(i) = cell_idx {
-                emit!(self, Instruction::RestoreCell { i: *i });
-            }
+        // Restore saved locals (StoreFast restores the saved cell object for merged cells)
+        for name in pushed_locals.iter().rev() {
             let var_num = self.varname(name)?;
             emit!(self, Instruction::StoreFast { var_num });
         }

diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs
@@ -8,9 +8,10 @@ use num_traits::ToPrimitive;
 use rustpython_compiler_core::{
     OneIndexed, SourceLocation,
     bytecode::{
-        AnyInstruction, Arg, CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData,
-        ExceptionTableEntry, InstrDisplayContext, Instruction, InstructionMetadata, Label, OpArg,
-        PseudoInstruction, PyCodeLocationInfoKind, encode_exception_table, oparg,
+        AnyInstruction, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_LOCAL, CodeFlags, CodeObject,
+        CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, InstrDisplayContext, Instruction,
+        InstructionMetadata, Label, OpArg, PseudoInstruction, PyCodeLocationInfoKind,
+        encode_exception_table, oparg,
     },
     varint::{write_signed_varint, write_varint},
 };
@@ -210,7 +211,6 @@ impl CodeInfo {
         self.optimize_load_global_push_null();
 
         let max_stackdepth = self.max_stackdepth()?;
-        let cell2arg = self.cell2arg();
 
         let Self {
             flags,
@@ -247,8 +247,12 @@ impl CodeInfo {
         let mut locations = Vec::new();
         let mut linetable_locations: Vec<LineTableLocation> = Vec::new();
 
-        // Convert pseudo ops and remove resulting NOPs (keep line-marker NOPs)
-        convert_pseudo_ops(&mut blocks, varname_cache.len() as u32);
+        // Build cellfixedoffsets for cell-local merging
+        let cellfixedoffsets =
+            build_cellfixedoffsets(&varname_cache, &cellvar_cache, &freevar_cache);
+        // Convert pseudo ops (LoadClosure uses cellfixedoffsets) and fixup DEREF opargs
+        convert_pseudo_ops(&mut blocks, &cellfixedoffsets);
+        fixup_deref_opargs(&mut blocks, &cellfixedoffsets);
         // Remove redundant NOPs, keeping line-marker NOPs only when
         // they are needed to preserve tracing.
         let mut block_order = Vec::new();
@@ -482,6 +486,35 @@ impl CodeInfo {
         // Generate exception table before moving source_path
         let exceptiontable = generate_exception_table(&blocks, &block_to_index);
 
+        // Build localspluskinds with cell-local merging
+        let nlocals = varname_cache.len();
+        let ncells = cellvar_cache.len();
+        let nfrees = freevar_cache.len();
+        let numdropped = cellvar_cache
+            .iter()
+            .filter(|cv| varname_cache.contains(cv.as_str()))
+            .count();
+        let nlocalsplus = nlocals + ncells - numdropped + nfrees;
+        let mut localspluskinds = vec![0u8; nlocalsplus];
+        // Mark locals
+        for kind in localspluskinds.iter_mut().take(nlocals) {
+            *kind = CO_FAST_LOCAL;
+        }
+        // Mark cells (merged and non-merged)
+        for (i, cellvar) in cellvar_cache.iter().enumerate() {
+            let idx = cellfixedoffsets[i] as usize;
+            if varname_cache.contains(cellvar.as_str()) {
+                localspluskinds[idx] |= CO_FAST_CELL; // merged: LOCAL | CELL
+            } else {
+                localspluskinds[idx] = CO_FAST_CELL;
+            }
+        }
+        // Mark frees
+        for i in 0..nfrees {
+            let idx = cellfixedoffsets[ncells + i] as usize;
+            localspluskinds[idx] = CO_FAST_FREE;
+        }
+
         Ok(CodeObject {
             flags,
             posonlyarg_count,
@@ -500,43 +533,12 @@ impl CodeInfo {
             varnames: varname_cache.into_iter().collect(),
             cellvars: cellvar_cache.into_iter().collect(),
             freevars: freevar_cache.into_iter().collect(),
-            cell2arg,
+            localspluskinds: localspluskinds.into_boxed_slice(),
             linetable,
             exceptiontable,
         })
     }
 
-    fn cell2arg(&self) -> Option<Box<[i32]>> {
-        if self.metadata.cellvars.is_empty() {
-            return None;
-        }
-
-        let total_args = self.metadata.argcount
-            + self.metadata.kwonlyargcount
-            + self.flags.contains(CodeFlags::VARARGS) as u32
-            + self.flags.contains(CodeFlags::VARKEYWORDS) as u32;
-
-        let mut found_cellarg = false;
-        let cell2arg = self
-            .metadata
-            .cellvars
-            .iter()
-            .map(|var| {
-                self.metadata
-                    .varnames
-                    .get_index_of(var)
-                    // check that it's actually an arg
-                    .filter(|i| *i < total_args as usize)
-                    .map_or(-1, |i| {
-                        found_cellarg = true;
-                        i as i32
-                    })
-            })
-            .collect::<Box<[_]>>();
-
-        if found_cellarg { Some(cell2arg) } else { None }
-    }
-
     fn dce(&mut self) {
         for block in &mut self.blocks {
             let mut last_instr = None;
@@ -1107,12 +1109,19 @@ impl InstrDisplayContext for CodeInfo {
         self.metadata.varnames[var_num.as_usize()].as_ref()
     }
 
-    fn get_cell_name(&self, i: usize) -> &str {
-        self.metadata
-            .cellvars
-            .get_index(i)
-            .unwrap_or_else(|| &self.metadata.freevars[i - self.metadata.cellvars.len()])
-            .as_ref()
+    fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str {
+        let idx = var_num.as_usize();
+        let nlocals = self.metadata.varnames.len();
+        if idx < nlocals {
+            self.metadata.varnames[idx].as_ref()
+        } else {
+            let cell_idx = idx - nlocals;
+            self.metadata
+                .cellvars
+                .get_index(cell_idx)
+                .unwrap_or_else(|| &self.metadata.freevars[cell_idx - self.metadata.cellvars.len()])
+                .as_ref()
+        }
     }
 }
 
@@ -1768,7 +1777,7 @@ pub(crate) fn label_exception_targets(blocks: &mut [Block]) {
 
 /// Convert remaining pseudo ops to real instructions or NOP.
 /// flowgraph.c convert_pseudo_ops
-pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], varnames_len: u32) {
+pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32]) {
     for block in blocks.iter_mut() {
         for info in &mut block.instructions {
             let Some(pseudo) = info.instr.pseudo() else {
@@ -1786,9 +1795,10 @@ pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], varnames_len: u32) {
                 PseudoInstruction::PopBlock => {
                     info.instr = Instruction::Nop.into();
                 }
-                // LOAD_CLOSURE → LOAD_FAST (with varnames offset)
+                // LOAD_CLOSURE → LOAD_FAST (using cellfixedoffsets for merged layout)
                 PseudoInstruction::LoadClosure { i } => {
-                    let new_idx = varnames_len + i.get(info.arg);
+                    let cell_relative = i.get(info.arg) as usize;
+                    let new_idx = cellfixedoffsets[cell_relative];
                     info.arg = OpArg::new(new_idx);
                     info.instr = Instruction::LoadFast {
                         var_num: Arg::marker(),
@@ -1808,3 +1818,54 @@ pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], varnames_len: u32) {
         }
     }
 }
+
+/// Build cellfixedoffsets mapping: cell/free index -> localsplus index.
+/// Merged cells (cellvar also in varnames) get the local slot index.
+/// Non-merged cells get slots after nlocals. Free vars follow.
+pub(crate) fn build_cellfixedoffsets(
+    varnames: &IndexSet<String>,
+    cellvars: &IndexSet<String>,
+    freevars: &IndexSet<String>,
+) -> Vec<u32> {
+    let nlocals = varnames.len();
+    let ncells = cellvars.len();
+    let nfrees = freevars.len();
+    let mut fixed = Vec::with_capacity(ncells + nfrees);
+    let mut numdropped = 0usize;
+    for (i, cellvar) in cellvars.iter().enumerate() {
+        if let Some(local_idx) = varnames.get_index_of(cellvar) {
+            fixed.push(local_idx as u32);
+            numdropped += 1;
+        } else {
+            fixed.push((nlocals + i - numdropped) as u32);
+        }
+    }
+    for i in 0..nfrees {
+        fixed.push((nlocals + ncells - numdropped + i) as u32);
+    }
+    fixed
+}
+
+/// Convert DEREF instruction opargs from cell-relative indices to localsplus indices
+/// using the cellfixedoffsets mapping.
+pub(crate) fn fixup_deref_opargs(blocks: &mut [Block], cellfixedoffsets: &[u32]) {
+    for block in blocks.iter_mut() {
+        for info in &mut block.instructions {
+            let Some(instr) = info.instr.real() else {
+                continue;
+            };
+            let needs_fixup = matches!(
+                instr,
+                Instruction::LoadDeref { .. }
+                    | Instruction::StoreDeref { .. }
+                    | Instruction::DeleteDeref { .. }
+                    | Instruction::LoadFromDictOrDeref { .. }
+                    | Instruction::MakeCell { .. }
+            );
+            if needs_fixup {
+                let cell_relative = u32::from(info.arg) as usize;
+                info.arg = OpArg::new(cellfixedoffsets[cell_relative]);
+            }
+        }
+    }
+}