Newtype var_num oparg (#7400)

crates/codegen/src/compile.rs

@@ -715,14 +715,14 @@ impl Compiler {
     }
 
     /// Get the index of a local variable.
-    fn get_local_var_index(&mut self, name: &str) -> CompileResult<u32> {
+    fn get_local_var_index(&mut self, name: &str) -> CompileResult<oparg::VarNum> {
         let info = self.code_stack.last_mut().unwrap();
         let idx = info
             .metadata
             .varnames
             .get_index_of(name)
             .unwrap_or_else(|| info.metadata.varnames.insert_full(name.to_owned()).0);
-        Ok(idx.to_u32())
+        Ok(idx.to_u32().into())
     }
 
     /// Get the index of a global name.
@@ -1283,7 +1283,12 @@ impl Compiler {
     /// if format > VALUE_WITH_FAKE_GLOBALS (2): raise NotImplementedError
     fn emit_format_validation(&mut self) -> CompileResult<()> {
         // Load format parameter (first local variable, index 0)
-        emit!(self, Instruction::LoadFast { var_num: 0 });
+        emit!(
+            self,
+            Instruction::LoadFast {
+                var_num: oparg::VarNum::from_u32(0)
+            }
+        );
 
         // Load VALUE_WITH_FAKE_GLOBALS constant (2)
         self.emit_load_const(ConstantData::Integer { value: 2.into() });
@@ -1562,15 +1567,19 @@ impl Compiler {
     fn name(&mut self, name: &str) -> bytecode::NameIdx {
         self._name_inner(name, |i| &mut i.metadata.names)
     }
-    fn varname(&mut self, name: &str) -> CompileResult<bytecode::NameIdx> {
+
+    fn varname(&mut self, name: &str) -> CompileResult<oparg::VarNum> {
         // Note: __debug__ checks are now handled in symboltable phase
-        Ok(self._name_inner(name, |i| &mut i.metadata.varnames))
+        Ok(oparg::VarNum::from_u32(
+            self._name_inner(name, |i| &mut i.metadata.varnames),
+        ))
     }
+
     fn _name_inner(
         &mut self,
         name: &str,
         cache: impl FnOnce(&mut ir::CodeInfo) -> &mut IndexSet<String>,
-    ) -> bytecode::NameIdx {
+    ) -> u32 {
         let name = self.mangle(name);
         let cache = cache(self.current_code_info());
         cache
@@ -4129,7 +4138,8 @@ impl Compiler {
 
             // Load defaults/kwdefaults with LOAD_FAST
             for i in 0..num_typeparam_args {
-                emit!(self, Instruction::LoadFast { var_num: i as u32 });
+                let var_num = oparg::VarNum::from(i as u32);
+                emit!(self, Instruction::LoadFast { var_num });
             }
         }
 

crates/codegen/src/ir.rs

@@ -1109,8 +1109,8 @@ impl InstrDisplayContext for CodeInfo {
         self.metadata.names[i].as_ref()
     }
 
-    fn get_varname(&self, i: usize) -> &str {
-        self.metadata.varnames[i].as_ref()
+    fn get_varname(&self, var_num: oparg::VarNum) -> &str {
+        self.metadata.varnames[var_num.as_usize()].as_ref()
     }
 
     fn get_cell_name(&self, i: usize) -> &str {

crates/compiler-core/src/bytecode.rs

@@ -11,7 +11,7 @@ use bitflags::bitflags;
 use core::{
     cell::UnsafeCell,
     hash, mem,
-    ops::{Deref, Index},
+    ops::{Deref, Index, IndexMut},
     sync::atomic::{AtomicU8, AtomicU16, AtomicUsize, Ordering},
 };
 use itertools::Itertools;
@@ -318,6 +318,22 @@ impl<C: Constant> FromIterator<C> for Constants<C> {
     }
 }
 
+// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index
+impl<T> Index<oparg::VarNum> for [T] {
+    type Output = T;
+
+    fn index(&self, var_num: oparg::VarNum) -> &Self::Output {
+        &self[var_num.as_usize()]
+    }
+}
+
+// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index
+impl<T> IndexMut<oparg::VarNum> for [T] {
+    fn index_mut(&mut self, var_num: oparg::VarNum) -> &mut Self::Output {
+        &mut self[var_num.as_usize()]
+    }
+}
+
 /// Primary container of a single code object. Each python function has
 /// a code object. Also a module has a code object.
 #[derive(Clone)]
@@ -1037,8 +1053,7 @@ impl<C: Constant> CodeObject<C> {
     pub fn map_bag<Bag: ConstantBag>(self, bag: Bag) -> CodeObject<Bag::Constant> {
         let map_names = |names: Box<[C::Name]>| {
             names
-                .into_vec()
-                .into_iter()
+                .iter()
                 .map(|x| bag.make_name(x.as_ref()))
                 .collect::<Box<[_]>>()
         };
@@ -1123,7 +1138,7 @@ pub trait InstrDisplayContext {
 
     fn get_name(&self, i: usize) -> &str;
 
-    fn get_varname(&self, i: usize) -> &str;
+    fn get_varname(&self, var_num: oparg::VarNum) -> &str;
 
     fn get_cell_name(&self, i: usize) -> &str;
 }
@@ -1139,8 +1154,8 @@ impl<C: Constant> InstrDisplayContext for CodeObject<C> {
         self.names[i].as_ref()
     }
 
-    fn get_varname(&self, i: usize) -> &str {
-        self.varnames[i].as_ref()
+    fn get_varname(&self, var_num: oparg::VarNum) -> &str {
+        self.varnames[var_num].as_ref()
     }
 
     fn get_cell_name(&self, i: usize) -> &str {

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

@@ -133,7 +133,7 @@ pub enum Instruction {
         i: Arg<NameIdx>,
     } = 62,
     DeleteFast {
-        var_num: Arg<NameIdx>,
+        var_num: Arg<oparg::VarNum>,
     } = 63,
     DeleteGlobal {
         namei: Arg<NameIdx>,
@@ -192,19 +192,19 @@ pub enum Instruction {
         i: Arg<NameIdx>,
     } = 83,
     LoadFast {
-        var_num: Arg<NameIdx>,
+        var_num: Arg<oparg::VarNum>,
     } = 84,
     LoadFastAndClear {
-        var_num: Arg<NameIdx>,
+        var_num: Arg<oparg::VarNum>,
     } = 85,
     LoadFastBorrow {
-        var_num: Arg<NameIdx>,
+        var_num: Arg<oparg::VarNum>,
     } = 86,
     LoadFastBorrowLoadFastBorrow {
         var_nums: Arg<u32>,
     } = 87,
     LoadFastCheck {
-        var_num: Arg<NameIdx>,
+        var_num: Arg<oparg::VarNum>,
     } = 88,
     LoadFastLoadFast {
         var_nums: Arg<u32>,
@@ -276,7 +276,7 @@ pub enum Instruction {
         i: Arg<NameIdx>,
     } = 111,
     StoreFast {
-        var_num: Arg<NameIdx>,
+        var_num: Arg<oparg::VarNum>,
     } = 112,
     StoreFastLoadFast {
         var_nums: Arg<StoreFastLoadFast>,
@@ -1105,7 +1105,13 @@ impl InstructionMetadata for Instruction {
             };
             ($variant:ident, $map:ident = $arg_marker:expr) => {{
                 let arg = $arg_marker.get(arg);
-                write!(f, "{:pad$}({}, {})", stringify!($variant), arg, $map(arg))
+                write!(
+                    f,
+                    "{:pad$}({}, {})",
+                    stringify!($variant),
+                    u32::from(arg),
+                    $map(arg)
+                )
             }};
             ($variant:ident, $arg_marker:expr) => {
                 write!(f, "{:pad$}({})", stringify!($variant), $arg_marker.get(arg))
@@ -1120,7 +1126,7 @@ impl InstructionMetadata for Instruction {
             };
         }
 
-        let varname = |i: u32| ctx.get_varname(i as usize);
+        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);
 
@@ -1226,16 +1232,16 @@ impl InstructionMetadata for Instruction {
                 let oparg = var_nums.get(arg);
                 let idx1 = oparg >> 4;
                 let idx2 = oparg & 15;
-                let name1 = varname(idx1);
-                let name2 = varname(idx2);
+                let name1 = varname(idx1.into());
+                let name2 = varname(idx2.into());
                 write!(f, "{:pad$}({}, {})", "LOAD_FAST_LOAD_FAST", name1, name2)
             }
             Self::LoadFastBorrowLoadFastBorrow { var_nums } => {
                 let oparg = var_nums.get(arg);
                 let idx1 = oparg >> 4;
                 let idx2 = oparg & 15;
-                let name1 = varname(idx1);
-                let name2 = varname(idx2);
+                let name1 = varname(idx1.into());
+                let name2 = varname(idx2.into());
                 write!(
                     f,
                     "{:pad$}({}, {})",
@@ -1362,8 +1368,8 @@ impl InstructionMetadata for Instruction {
                     f,
                     "{:pad$}({}, {})",
                     "STORE_FAST_STORE_FAST",
-                    varname(idx1),
-                    varname(idx2)
+                    varname(idx1.into()),
+                    varname(idx2.into())
                 )
             }
             Self::StoreGlobal { namei } => w!(STORE_GLOBAL, name = namei),

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

@@ -873,38 +873,55 @@ impl LoadAttrBuilder {
     }
 }
 
-#[derive(Clone, Copy)]
-pub struct ConstIdx(u32);
+macro_rules! newtype_oparg {
+    (
+      $(#[$oparg_meta:meta])*
+      $vis:vis struct $name:ident(u32)
+    ) => {
+        $(#[$oparg_meta])*
+        $vis struct $name(u32);
 
-impl ConstIdx {
-    #[must_use]
-    pub const fn from_u32(value: u32) -> Self {
-        Self(value)
-    }
+        impl $name {
+            #[must_use]
+            pub const fn from_u32(value: u32) -> Self {
+                Self(value)
+            }
 
-    /// Returns the index as a `u32` value.
-    #[must_use]
-    pub const fn as_u32(self) -> u32 {
-        self.0
-    }
+            /// Returns the oparg as a `u32` value.
+            #[must_use]
+            pub const fn as_u32(self) -> u32 {
+                self.0
+            }
 
-    /// Returns the index as a `usize` value.
-    #[must_use]
-    pub const fn as_usize(self) -> usize {
-        self.0 as usize
+            /// Returns the oparg as a `usize` value.
+            #[must_use]
+            pub const fn as_usize(self) -> usize {
+              self.0 as usize
+            }
+        }
+
+        impl From<u32> for $name {
+            fn from(value: u32) -> Self {
+                Self::from_u32(value)
+            }
+        }
+
+        impl From<$name> for u32 {
+            fn from(value: $name) -> Self {
+                value.as_u32()
+            }
+        }
+
+        impl OpArgType for $name {}
     }
 }
 
-impl From<u32> for ConstIdx {
-    fn from(value: u32) -> Self {
-        Self::from_u32(value)
-    }
-}
+newtype_oparg!(
+    #[derive(Clone, Copy)]
+    pub struct ConstIdx(u32)
+);
 
-impl From<ConstIdx> for u32 {
-    fn from(consti: ConstIdx) -> Self {
-        consti.as_u32()
-    }
-}
-
-impl OpArgType for ConstIdx {}
+newtype_oparg!(
+    #[derive(Clone, Copy)]
+    pub struct VarNum(u32)
+);

crates/jit/src/instructions.rs

@@ -6,7 +6,7 @@ use cranelift::prelude::*;
 use num_traits::cast::ToPrimitive;
 use rustpython_compiler_core::bytecode::{
     self, BinaryOperator, BorrowedConstant, CodeObject, ComparisonOperator, Instruction,
-    IntrinsicFunction1, Label, OpArg, OpArgState,
+    IntrinsicFunction1, Label, OpArg, OpArgState, oparg,
 };
 use std::collections::HashMap;
 
@@ -94,7 +94,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         let params = compiler.builder.func.dfg.block_params(entry_block).to_vec();
         for (i, (ty, val)) in arg_types.iter().zip(params).enumerate() {
             compiler
-                .store_variable(i as u32, JitValue::from_type_and_value(ty.clone(), val))
+                .store_variable(
+                    (i as u32).into(),
+                    JitValue::from_type_and_value(ty.clone(), val),
+                )
                 .unwrap();
         }
         compiler
@@ -105,14 +108,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         self.stack.drain(stack_len - count..).collect()
     }
 
-    fn store_variable(
-        &mut self,
-        idx: bytecode::NameIdx,
-        val: JitValue,
-    ) -> Result<(), JitCompileError> {
+    fn store_variable(&mut self, idx: oparg::VarNum, val: JitValue) -> Result<(), JitCompileError> {
         let builder = &mut self.builder;
         let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?;
-        let local = self.variables[idx as usize].get_or_insert_with(|| {
+        let local = self.variables[idx].get_or_insert_with(|| {
             let var = builder.declare_var(ty.to_cranelift());
             Local {
                 var,
@@ -649,7 +648,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 Ok(())
             }
             Instruction::LoadFast { var_num } | Instruction::LoadFastBorrow { var_num } => {
-                let local = self.variables[var_num.get(arg) as usize]
+                let local = self.variables[var_num.get(arg)]
                     .as_ref()
                     .ok_or(JitCompileError::BadBytecode)?;
                 self.stack.push(JitValue::from_type_and_value(

crates/vm/src/frame.rs

@@ -2162,9 +2162,9 @@ impl ExecutingFrame<'_> {
                 self.cell_ref(i.get(arg) as usize).set(None);
                 Ok(None)
             }
-            Instruction::DeleteFast { var_num: idx } => {
+            Instruction::DeleteFast { var_num } => {
                 let fastlocals = self.localsplus.fastlocals_mut();
-                let idx = idx.get(arg) as usize;
+                let idx = var_num.get(arg);
                 if fastlocals[idx].is_none() {
                     return Err(vm.new_exception_msg(
                         vm.ctx.exceptions.unbound_local_error.to_owned(),
@@ -2680,7 +2680,7 @@ impl ExecutingFrame<'_> {
                 self.push_value(x);
                 Ok(None)
             }
-            Instruction::LoadFast { var_num: idx } => {
+            Instruction::LoadFast { var_num } => {
                 #[cold]
                 fn reference_error(
                     varname: &'static PyStrInterned,
@@ -2691,27 +2691,27 @@ impl ExecutingFrame<'_> {
                         format!("local variable '{varname}' referenced before assignment").into(),
                     )
                 }
-                let idx = idx.get(arg) as usize;
+                let idx = var_num.get(arg);
                 let x = self.localsplus.fastlocals()[idx]
                     .clone()
                     .ok_or_else(|| reference_error(self.code.varnames[idx], vm))?;
                 self.push_value(x);
                 Ok(None)
             }
-            Instruction::LoadFastAndClear { var_num: idx } => {
+            Instruction::LoadFastAndClear { var_num } => {
                 // Load value and clear the slot (for inlined comprehensions)
                 // If slot is empty, push None (not an error - variable may not exist yet)
-                let idx = idx.get(arg) as usize;
+                let idx = var_num.get(arg);
                 let x = self.localsplus.fastlocals_mut()[idx]
                     .take()
                     .unwrap_or_else(|| vm.ctx.none());
                 self.push_value(x);
                 Ok(None)
             }
-            Instruction::LoadFastCheck { var_num: idx } => {
+            Instruction::LoadFastCheck { var_num } => {
                 // Same as LoadFast but explicitly checks for unbound locals
                 // (LoadFast in RustPython already does this check)
-                let idx = idx.get(arg) as usize;
+                let idx = var_num.get(arg);
                 let x = self.localsplus.fastlocals()[idx].clone().ok_or_else(|| {
                     vm.new_exception_msg(
                         vm.ctx.exceptions.unbound_local_error.to_owned(),
@@ -2759,8 +2759,8 @@ impl ExecutingFrame<'_> {
             // Borrow optimization not yet active; falls back to clone.
             // push_borrowed() is available but disabled until stack
             // lifetime issues at yield/exception points are resolved.
-            Instruction::LoadFastBorrow { var_num: idx } => {
-                let idx = idx.get(arg) as usize;
+            Instruction::LoadFastBorrow { var_num } => {
+                let idx = var_num.get(arg);
                 let x = self.localsplus.fastlocals()[idx].clone().ok_or_else(|| {
                     vm.new_exception_msg(
                         vm.ctx.exceptions.unbound_local_error.to_owned(),
@@ -3299,10 +3299,10 @@ impl ExecutingFrame<'_> {
                 self.cell_ref(i.get(arg) as usize).set(Some(value));
                 Ok(None)
             }
-            Instruction::StoreFast { var_num: idx } => {
+            Instruction::StoreFast { var_num } => {
                 let value = self.pop_value();
                 let fastlocals = self.localsplus.fastlocals_mut();
-                fastlocals[idx.get(arg) as usize] = Some(value);
+                fastlocals[var_num.get(arg)] = Some(value);
                 Ok(None)
             }
             Instruction::StoreFastLoadFast { var_nums } => {