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

Author: youknowone

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.
@@ -8004,7 +8006,7 @@ impl Compiler {
         // 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();
+        let mut cell_indices: Vec<Option<oparg::VarNum>> = Vec::new();
         let mut total_stack_items: usize = 0;
         for name in &pushed_locals {
             let var_num = self.varname(name)?;

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},
 };
@@ -255,6 +256,8 @@ impl CodeInfo {
 
         // Convert pseudo ops and remove resulting NOPs (keep line-marker NOPs)
         convert_pseudo_ops(&mut blocks, varname_cache.len() as u32);
+        // Fixup DEREF opargs: cell-relative → localsplus indices
+        fixup_deref_opargs(&mut blocks, varname_cache.len() as u32);
         // Remove redundant NOPs, keeping line-marker NOPs only when
         // they are needed to preserve tracing.
         let mut block_order = Vec::new();
@@ -488,6 +491,15 @@ impl CodeInfo {
         // Generate exception table before moving source_path
         let exceptiontable = generate_exception_table(&blocks, &block_to_index);
 
+        // Build localspluskinds: per-slot kind flags
+        let nlocals = varname_cache.len();
+        let ncells = cellvar_cache.len();
+        let nfrees = freevar_cache.len();
+        let mut localspluskinds = Vec::with_capacity(nlocals + ncells + nfrees);
+        localspluskinds.extend(core::iter::repeat_n(CO_FAST_LOCAL, nlocals));
+        localspluskinds.extend(core::iter::repeat_n(CO_FAST_CELL, ncells));
+        localspluskinds.extend(core::iter::repeat_n(CO_FAST_FREE, nfrees));
+
         Ok(CodeObject {
             flags,
             posonlyarg_count,
@@ -507,6 +519,7 @@ impl CodeInfo {
             cellvars: cellvar_cache.into_iter().collect(),
             freevars: freevar_cache.into_iter().collect(),
             cell2arg,
+            localspluskinds: localspluskinds.into_boxed_slice(),
             linetable,
             exceptiontable,
         })
@@ -1113,12 +1126,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()
+        }
     }
 }
 
@@ -1814,3 +1834,28 @@ pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], varnames_len: u32) {
         }
     }
 }
+
+/// Convert DEREF instruction opargs from cell-relative indices to localsplus indices.
+/// Cell-relative index = cellvar_idx for cells, ncells + freevar_idx for frees.
+/// Localsplus index = nlocals + cell_relative_idx.
+pub(crate) fn fixup_deref_opargs(blocks: &mut [Block], nlocals: 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 { .. }
+                    | Instruction::RestoreCell { .. }
+            );
+            if needs_fixup {
+                info.arg = OpArg::new(nlocals + u32::from(info.arg));
+            }
+        }
+    }
+}

crates/compiler-core/src/bytecode.rs

@@ -334,6 +334,12 @@ impl<T> IndexMut<oparg::VarNum> for [T] {
     }
 }
 
+/// Per-slot kind flags for localsplus (co_localspluskinds).
+pub const CO_FAST_HIDDEN: u8 = 0x10;
+pub const CO_FAST_LOCAL: u8 = 0x20;
+pub const CO_FAST_CELL: u8 = 0x40;
+pub const CO_FAST_FREE: u8 = 0x80;
+
 /// Primary container of a single code object. Each python function has
 /// a code object. Also a module has a code object.
 #[derive(Clone)]
@@ -358,6 +364,9 @@ pub struct CodeObject<C: Constant = ConstantData> {
     pub varnames: Box<[C::Name]>,
     pub cellvars: Box<[C::Name]>,
     pub freevars: Box<[C::Name]>,
+    /// Per-slot kind flags: CO_FAST_LOCAL, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN.
+    /// Length = nlocalsplus (nlocals + ncells + nfrees).
+    pub localspluskinds: Box<[u8]>,
     /// Line number table (CPython 3.11+ format)
     pub linetable: Box<[u8]>,
     /// Exception handling table
@@ -1080,6 +1089,7 @@ impl<C: Constant> CodeObject<C> {
             first_line_number: self.first_line_number,
             max_stackdepth: self.max_stackdepth,
             cell2arg: self.cell2arg,
+            localspluskinds: self.localspluskinds,
             linetable: self.linetable,
             exceptiontable: self.exceptiontable,
         }
@@ -1112,6 +1122,7 @@ impl<C: Constant> CodeObject<C> {
             first_line_number: self.first_line_number,
             max_stackdepth: self.max_stackdepth,
             cell2arg: self.cell2arg.clone(),
+            localspluskinds: self.localspluskinds.clone(),
             linetable: self.linetable.clone(),
             exceptiontable: self.exceptiontable.clone(),
         }
@@ -1140,7 +1151,8 @@ pub trait InstrDisplayContext {
 
     fn get_varname(&self, var_num: oparg::VarNum) -> &str;
 
-    fn get_cell_name(&self, i: usize) -> &str;
+    /// Get name for a localsplus index (used by DEREF instructions).
+    fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str;
 }
 
 impl<C: Constant> InstrDisplayContext for CodeObject<C> {
@@ -1158,11 +1170,18 @@ impl<C: Constant> InstrDisplayContext for CodeObject<C> {
         self.varnames[var_num].as_ref()
     }
 
-    fn get_cell_name(&self, i: usize) -> &str {
-        self.cellvars
-            .get(i)
-            .unwrap_or_else(|| &self.freevars[i - self.cellvars.len()])
-            .as_ref()
+    fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str {
+        let idx = var_num.as_usize();
+        let nlocals = self.varnames.len();
+        if idx < nlocals {
+            self.varnames[idx].as_ref()
+        } else {
+            let cell_idx = idx - nlocals;
+            self.cellvars
+                .get(cell_idx)
+                .unwrap_or_else(|| &self.freevars[cell_idx - self.cellvars.len()])
+                .as_ref()
+        }
     }
 }
 

crates/compiler-core/src/bytecode/instruction.rs

@@ -130,7 +130,7 @@ pub enum Instruction {
         namei: Arg<NameIdx>,
     } = 61,
     DeleteDeref {
-        i: Arg<NameIdx>,
+        i: Arg<oparg::VarNum>,
     } = 62,
     DeleteFast {
         var_num: Arg<oparg::VarNum>,
@@ -189,7 +189,7 @@ pub enum Instruction {
         consti: Arg<oparg::ConstIdx>,
     } = 82,
     LoadDeref {
-        i: Arg<NameIdx>,
+        i: Arg<oparg::VarNum>,
     } = 83,
     LoadFast {
         var_num: Arg<oparg::VarNum>,
@@ -210,7 +210,7 @@ pub enum Instruction {
         var_nums: Arg<oparg::VarNums>,
     } = 89,
     LoadFromDictOrDeref {
-        i: Arg<NameIdx>,
+        i: Arg<oparg::VarNum>,
     } = 90,
     LoadFromDictOrGlobals {
         i: Arg<NameIdx>,
@@ -231,7 +231,7 @@ pub enum Instruction {
         namei: Arg<LoadSuperAttr>,
     } = 96,
     MakeCell {
-        i: Arg<NameIdx>,
+        i: Arg<oparg::VarNum>,
     } = 97,
     MapAdd {
         i: Arg<u32>,
@@ -273,7 +273,7 @@ pub enum Instruction {
         namei: Arg<NameIdx>,
     } = 110,
     StoreDeref {
-        i: Arg<NameIdx>,
+        i: Arg<oparg::VarNum>,
     } = 111,
     StoreFast {
         var_num: Arg<oparg::VarNum>,
@@ -307,7 +307,7 @@ pub enum Instruction {
     /// cell slot at `fastlocals[nlocals + i]`, replacing the temporary cell
     /// created by `MakeCell`.
     RestoreCell {
-        i: Arg<NameIdx>,
+        i: Arg<oparg::VarNum>,
     } = 121,
     // CPython 3.14 RESUME (128)
     Resume {
@@ -1130,7 +1130,7 @@ impl InstructionMetadata for Instruction {
 
         let varname = |var_num: oparg::VarNum| ctx.get_varname(var_num);
         let name = |i: u32| ctx.get_name(i as usize);
-        let cell_name = |i: u32| ctx.get_cell_name(i as usize);
+        let cell_name = |i: oparg::VarNum| ctx.get_localsplus_name(i);
 
         let fmt_const = |op: &str,
                          arg: OpArg,

crates/compiler-core/src/marshal.rs

@@ -267,6 +267,19 @@ pub fn deserialize_code<R: Read, Bag: ConstantBag>(
         .to_vec()
         .into_boxed_slice();
 
+    // Build localspluskinds from varnames/cellvars/freevars
+    let localspluskinds = {
+        use crate::bytecode::*;
+        let nlocals = varnames.len();
+        let ncells = cellvars.len();
+        let nfrees = freevars.len();
+        let mut kinds = Vec::with_capacity(nlocals + ncells + nfrees);
+        kinds.extend(core::iter::repeat_n(CO_FAST_LOCAL, nlocals));
+        kinds.extend(core::iter::repeat_n(CO_FAST_CELL, ncells));
+        kinds.extend(core::iter::repeat_n(CO_FAST_FREE, nfrees));
+        kinds.into_boxed_slice()
+    };
+
     Ok(CodeObject {
         instructions,
         locations,
@@ -285,6 +298,7 @@ pub fn deserialize_code<R: Read, Bag: ConstantBag>(
         varnames,
         cellvars,
         freevars,
+        localspluskinds,
         linetable,
         exceptiontable,
     })

crates/vm/src/builtins/code.rs

@@ -564,6 +564,19 @@ impl Constructor for PyCode {
             )],
         > = vec![(loc, loc); instructions.len()].into_boxed_slice();
 
+        // Build localspluskinds before moving varnames/cellvars/freevars
+        let localspluskinds = {
+            use rustpython_compiler_core::bytecode::*;
+            let nlocals = varnames.len();
+            let ncells = cellvars.len();
+            let nfrees = freevars.len();
+            let mut kinds = Vec::with_capacity(nlocals + ncells + nfrees);
+            kinds.extend(core::iter::repeat_n(CO_FAST_LOCAL, nlocals));
+            kinds.extend(core::iter::repeat_n(CO_FAST_CELL, ncells));
+            kinds.extend(core::iter::repeat_n(CO_FAST_FREE, nfrees));
+            kinds.into_boxed_slice()
+        };
+
         // Build the CodeObject
         let code = CodeObject {
             instructions,
@@ -587,6 +600,7 @@ impl Constructor for PyCode {
             varnames,
             cellvars,
             freevars,
+            localspluskinds,
             linetable: args.linetable.as_bytes().to_vec().into_boxed_slice(),
             exceptiontable: args.exceptiontable.as_bytes().to_vec().into_boxed_slice(),
         };
@@ -1156,6 +1170,7 @@ impl PyCode {
             cellvars,
             freevars,
             cell2arg: self.code.cell2arg.clone(),
+            localspluskinds: self.code.localspluskinds.clone(),
             linetable,
             exceptiontable,
         };

crates/vm/src/builtins/function.rs

@@ -442,9 +442,10 @@ impl PyFunction {
         }
 
         if let Some(cell2arg) = code.cell2arg.as_deref() {
+            let nlocals = code.varnames.len();
             for (cell_idx, arg_idx) in cell2arg.iter().enumerate().filter(|(_, i)| **i != -1) {
                 let x = fastlocals[*arg_idx as usize].take();
-                frame.set_cell_contents(cell_idx, x);
+                frame.set_cell_contents(nlocals + cell_idx, x);
             }
         }
 
@@ -724,10 +725,11 @@ impl Py<PyFunction> {
         }
 
         if let Some(cell2arg) = code.cell2arg.as_deref() {
+            let nlocals = code.varnames.len();
             let fastlocals = unsafe { frame.fastlocals_mut() };
             for (cell_idx, arg_idx) in cell2arg.iter().enumerate().filter(|(_, i)| **i != -1) {
                 let x = fastlocals[*arg_idx as usize].take();
-                frame.set_cell_contents(cell_idx, x);
+                frame.set_cell_contents(nlocals + cell_idx, x);
             }
         }
 

crates/vm/src/builtins/super.rs

@@ -86,6 +86,7 @@ impl Initializer for PySuper {
                 return Err(vm.new_runtime_error("super(): no arguments"));
             }
             // SAFETY: Frame is current and not concurrently mutated.
+            let nlocals = frame.code.varnames.len();
             let obj = unsafe { frame.fastlocals() }[0]
                 .clone()
                 .or_else(|| {
@@ -94,7 +95,7 @@ impl Initializer for PySuper {
                             .iter()
                             .enumerate()
                             .find(|(_, arg_idx)| **arg_idx == 0)
-                            .and_then(|(cell_idx, _)| frame.get_cell_contents(cell_idx))
+                            .and_then(|(cell_idx, _)| frame.get_cell_contents(nlocals + cell_idx))
                     } else {
                         None
                     }
@@ -104,7 +105,7 @@ impl Initializer for PySuper {
             let mut typ = None;
             for (i, var) in frame.code.freevars.iter().enumerate() {
                 if var.as_bytes() == b"__class__" {
-                    let i = frame.code.cellvars.len() + i;
+                    let i = nlocals + frame.code.cellvars.len() + i;
                     let class = frame
                         .get_cell_contents(i)
                         .ok_or_else(|| vm.new_runtime_error("super(): empty __class__ cell"))?;