Newtype var_nums oparg (#7431)

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

@@ -7,7 +7,7 @@ use crate::{
             self, BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator,
             ConvertValueOparg, IntrinsicFunction1, IntrinsicFunction2, Invert, Label, LoadAttr,
             LoadSuperAttr, MakeFunctionFlag, NameIdx, OpArg, OpArgByte, OpArgType, RaiseKind,
-            SpecialMethod, StoreFastLoadFast, UnpackExArgs,
+            SpecialMethod, UnpackExArgs,
         },
     },
     marshal::MarshalError,
@@ -201,13 +201,13 @@ pub enum Instruction {
         var_num: Arg<oparg::VarNum>,
     } = 86,
     LoadFastBorrowLoadFastBorrow {
-        var_nums: Arg<u32>,
+        var_nums: Arg<oparg::VarNums>,
     } = 87,
     LoadFastCheck {
         var_num: Arg<oparg::VarNum>,
     } = 88,
     LoadFastLoadFast {
-        var_nums: Arg<u32>,
+        var_nums: Arg<oparg::VarNums>,
     } = 89,
     LoadFromDictOrDeref {
         i: Arg<NameIdx>,
@@ -279,10 +279,10 @@ pub enum Instruction {
         var_num: Arg<oparg::VarNum>,
     } = 112,
     StoreFastLoadFast {
-        var_nums: Arg<StoreFastLoadFast>,
+        var_nums: Arg<oparg::VarNums>,
     } = 113,
     StoreFastStoreFast {
-        var_nums: Arg<u32>,
+        var_nums: Arg<oparg::VarNums>,
     } = 114,
     StoreGlobal {
         namei: Arg<NameIdx>,
@@ -1224,18 +1224,16 @@ impl InstructionMetadata for Instruction {
             Self::LoadFastCheck { var_num } => w!(LOAD_FAST_CHECK, varname = var_num),
             Self::LoadFastLoadFast { var_nums } => {
                 let oparg = var_nums.get(arg);
-                let idx1 = oparg >> 4;
-                let idx2 = oparg & 15;
-                let name1 = varname(idx1.into());
-                let name2 = varname(idx2.into());
+                let (idx1, idx2) = oparg.indexes();
+                let name1 = varname(idx1);
+                let name2 = varname(idx2);
                 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.into());
-                let name2 = varname(idx2.into());
+                let (idx1, idx2) = oparg.indexes();
+                let name1 = varname(idx1);
+                let name2 = varname(idx2);
                 write!(
                     f,
                     "{:pad$}({}, {})",
@@ -1349,21 +1347,19 @@ impl InstructionMetadata for Instruction {
             Self::StoreFast { var_num } => w!(STORE_FAST, varname = var_num),
             Self::StoreFastLoadFast { var_nums } => {
                 let oparg = var_nums.get(arg);
-                let store_idx = oparg.store_idx();
-                let load_idx = oparg.load_idx();
+                let (store_idx, load_idx) = oparg.indexes();
                 write!(f, "STORE_FAST_LOAD_FAST")?;
                 write!(f, " ({}, {})", store_idx, load_idx)
             }
             Self::StoreFastStoreFast { var_nums } => {
                 let oparg = var_nums.get(arg);
-                let idx1 = oparg >> 4;
-                let idx2 = oparg & 15;
+                let (idx1, idx2) = oparg.indexes();
                 write!(
                     f,
                     "{:pad$}({}, {})",
                     "STORE_FAST_STORE_FAST",
-                    varname(idx1.into()),
-                    varname(idx2.into())
+                    varname(idx1),
+                    varname(idx2)
                 )
             }
             Self::StoreGlobal { namei } => w!(STORE_GLOBAL, name = namei),

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

@@ -730,6 +730,12 @@ newtype_oparg!(
     pub struct VarNum(u32)
 );
 
+newtype_oparg!(
+    #[derive(Clone, Copy)]
+    #[repr(transparent)]
+    pub struct VarNums(u32)
+);
+
 newtype_oparg!(
     #[derive(Clone, Copy)]
     #[repr(transparent)]
@@ -748,21 +754,20 @@ newtype_oparg!(
     pub struct Label(u32)
 );
 
-newtype_oparg!(
-    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
-    #[repr(transparent)]
-    pub struct StoreFastLoadFast(u32)
-);
-
-impl StoreFastLoadFast {
+impl VarNums {
     #[must_use]
-    pub const fn store_idx(self) -> NameIdx {
-        self.0 >> 4
+    pub const fn idx_1(self) -> VarNum {
+        VarNum::new(self.0 >> 4)
     }
 
     #[must_use]
-    pub const fn load_idx(self) -> NameIdx {
-        self.0 & 15
+    pub const fn idx_2(self) -> VarNum {
+        VarNum::new(self.0 & 15)
+    }
+
+    #[must_use]
+    pub const fn indexes(self) -> (VarNum, VarNum) {
+        (self.idx_1(), self.idx_2())
     }
 }
 

crates/jit/src/instructions.rs

@@ -660,10 +660,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             Instruction::LoadFastLoadFast { var_nums }
             | Instruction::LoadFastBorrowLoadFastBorrow { var_nums } => {
                 let oparg = var_nums.get(arg);
-                let idx1 = oparg >> 4;
-                let idx2 = oparg & 0xF;
+                let (idx1, idx2) = oparg.indexes();
                 for idx in [idx1, idx2] {
-                    let local = self.variables[idx as usize]
+                    let local = self.variables[idx]
                         .as_ref()
                         .ok_or(JitCompileError::BadBytecode)?;
                     self.stack.push(JitValue::from_type_and_value(

crates/vm/src/frame.rs

@@ -2725,12 +2725,11 @@ impl ExecutingFrame<'_> {
                 self.push_value(x);
                 Ok(None)
             }
-            Instruction::LoadFastLoadFast { var_nums: packed } => {
+            Instruction::LoadFastLoadFast { var_nums } => {
                 // Load two local variables at once
                 // oparg encoding: (idx1 << 4) | idx2
-                let oparg = packed.get(arg);
-                let idx1 = (oparg >> 4) as usize;
-                let idx2 = (oparg & 15) as usize;
+                let oparg = var_nums.get(arg);
+                let (idx1, idx2) = oparg.indexes();
                 let fastlocals = self.localsplus.fastlocals();
                 let x1 = fastlocals[idx1].clone().ok_or_else(|| {
                     vm.new_exception_msg(
@@ -2774,10 +2773,9 @@ impl ExecutingFrame<'_> {
                 self.push_value(x);
                 Ok(None)
             }
-            Instruction::LoadFastBorrowLoadFastBorrow { var_nums: packed } => {
-                let oparg = packed.get(arg);
-                let idx1 = (oparg >> 4) as usize;
-                let idx2 = (oparg & 15) as usize;
+            Instruction::LoadFastBorrowLoadFastBorrow { var_nums } => {
+                let oparg = var_nums.get(arg);
+                let (idx1, idx2) = oparg.indexes();
                 let fastlocals = self.localsplus.fastlocals();
                 let x1 = fastlocals[idx1].clone().ok_or_else(|| {
                     vm.new_exception_msg(
@@ -3309,17 +3307,17 @@ impl ExecutingFrame<'_> {
                 let value = self.pop_value();
                 let locals = self.localsplus.fastlocals_mut();
                 let oparg = var_nums.get(arg);
-                locals[oparg.store_idx() as usize] = Some(value);
-                let load_value = locals[oparg.load_idx() as usize]
+                let (store_idx, load_idx) = oparg.indexes();
+                locals[store_idx] = Some(value);
+                let load_value = locals[load_idx]
                     .clone()
                     .expect("StoreFastLoadFast: load slot should have value after store");
                 self.push_value(load_value);
                 Ok(None)
             }
-            Instruction::StoreFastStoreFast { var_nums: packed } => {
-                let oparg = packed.get(arg);
-                let idx1 = (oparg >> 4) as usize;
-                let idx2 = (oparg & 15) as usize;
+            Instruction::StoreFastStoreFast { var_nums } => {
+                let oparg = var_nums.get(arg);
+                let (idx1, idx2) = oparg.indexes();
                 let value1 = self.pop_value();
                 let value2 = self.pop_value();
                 let fastlocals = self.localsplus.fastlocals_mut();