Newtype for LoadAttr oparg (#7047)

ShaharNaveh · web-flow · commit c8b4d6308fc3 · 2026-02-08T11:47:14.000Z
diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs
@@ -30,8 +30,8 @@ use rustpython_compiler_core::{
     bytecode::{
         self, AnyInstruction, Arg as OpArgMarker, BinaryOperator, BuildSliceArgCount, CodeObject,
         ComparisonOperator, ConstantData, ConvertValueOparg, Instruction, IntrinsicFunction1,
-        Invert, LoadSuperAttr, OpArg, OpArgType, PseudoInstruction, SpecialMethod, UnpackExArgs,
-        encode_load_attr_arg,
+        Invert, LoadAttr, LoadSuperAttr, OpArg, OpArgType, PseudoInstruction, SpecialMethod,
+        UnpackExArgs,
     },
 };
 use rustpython_wtf8::Wtf8Buf;
@@ -7799,14 +7799,20 @@ impl Compiler {
     /// Emit LOAD_ATTR for attribute access (method=false).
     /// Encodes: (name_idx << 1) | 0
     fn emit_load_attr(&mut self, name_idx: u32) {
-        let encoded = encode_load_attr_arg(name_idx, false);
+        let encoded = LoadAttr::builder()
+            .name_idx(name_idx)
+            .is_method(false)
+            .build();
         self.emit_arg(encoded, |arg| Instruction::LoadAttr { idx: arg })
     }
 
     /// Emit LOAD_ATTR with method flag set (for method calls).
     /// Encodes: (name_idx << 1) | 1
     fn emit_load_attr_method(&mut self, name_idx: u32) {
-        let encoded = encode_load_attr_arg(name_idx, true);
+        let encoded = LoadAttr::builder()
+            .name_idx(name_idx)
+            .is_method(true)
+            .build();
         self.emit_arg(encoded, |arg| Instruction::LoadAttr { idx: arg })
     }
 
diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs
@@ -17,13 +17,12 @@ use rustpython_wtf8::{Wtf8, Wtf8Buf};
 pub use crate::bytecode::{
     instruction::{
         AnyInstruction, Arg, Instruction, InstructionMetadata, PseudoInstruction, StackEffect,
-        decode_load_attr_arg, encode_load_attr_arg,
     },
     oparg::{
         BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator, ConvertValueOparg,
-        IntrinsicFunction1, IntrinsicFunction2, Invert, Label, LoadSuperAttr, MakeFunctionFlags,
-        NameIdx, OpArg, OpArgByte, OpArgState, OpArgType, RaiseKind, ResumeType, SpecialMethod,
-        UnpackExArgs,
+        IntrinsicFunction1, IntrinsicFunction2, Invert, Label, LoadAttr, LoadSuperAttr,
+        MakeFunctionFlags, NameIdx, OpArg, OpArgByte, OpArgState, OpArgType, RaiseKind, ResumeType,
+        SpecialMethod, UnpackExArgs,
     },
 };
 
diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs
@@ -5,7 +5,7 @@ use crate::{
         BorrowedConstant, Constant, InstrDisplayContext,
         oparg::{
             BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator,
-            ConvertValueOparg, IntrinsicFunction1, IntrinsicFunction2, Invert, Label,
+            ConvertValueOparg, IntrinsicFunction1, IntrinsicFunction2, Invert, Label, LoadAttr,
             LoadSuperAttr, MakeFunctionFlags, NameIdx, OpArg, OpArgByte, OpArgType, RaiseKind,
             SpecialMethod, UnpackExArgs,
         },
@@ -168,7 +168,7 @@ pub enum Instruction {
         i: Arg<u32>,
     } = 79,
     LoadAttr {
-        idx: Arg<NameIdx>,
+        idx: Arg<LoadAttr>,
     } = 80,
     LoadCommonConstant {
         idx: Arg<CommonConstant>,
@@ -815,17 +815,17 @@ impl InstructionMetadata for Instruction {
             Self::ListAppend { i } => w!(LIST_APPEND, i),
             Self::ListExtend { i } => w!(LIST_EXTEND, i),
             Self::LoadAttr { idx } => {
-                let encoded = idx.get(arg);
-                let (name_idx, is_method) = decode_load_attr_arg(encoded);
-                let attr_name = name(name_idx);
-                if is_method {
+                let oparg = idx.get(arg);
+                let oparg_u32 = u32::from(oparg);
+                let attr_name = name(oparg.name_idx());
+                if oparg.is_method() {
                     write!(
                         f,
                         "{:pad$}({}, {}, method=true)",
-                        "LOAD_ATTR", encoded, attr_name
+                        "LOAD_ATTR", oparg_u32, attr_name
                     )
                 } else {
-                    write!(f, "{:pad$}({}, {})", "LOAD_ATTR", encoded, attr_name)
+                    write!(f, "{:pad$}({}, {})", "LOAD_ATTR", oparg_u32, attr_name)
                 }
             }
             Self::LoadBuildClass => w!(LOAD_BUILD_CLASS),
@@ -1292,17 +1292,3 @@ impl<T: OpArgType> fmt::Debug for Arg<T> {
         write!(f, "Arg<{}>", core::any::type_name::<T>())
     }
 }
-
-/// Encode LOAD_ATTR oparg: bit 0 = method flag, bits 1+ = name index.
-#[inline]
-pub const fn encode_load_attr_arg(name_idx: u32, is_method: bool) -> u32 {
-    (name_idx << 1) | (is_method as u32)
-}
-
-/// Decode LOAD_ATTR oparg: returns (name_idx, is_method).
-#[inline]
-pub const fn decode_load_attr_arg(oparg: u32) -> (u32, bool) {
-    let is_method = (oparg & 1) == 1;
-    let name_idx = oparg >> 1;
-    (name_idx, is_method)
-}
diff --git a/crates/compiler-core/src/bytecode/oparg.rs b/crates/compiler-core/src/bytecode/oparg.rs
@@ -754,3 +754,68 @@ impl From<LoadSuperAttrBuilder> for LoadSuperAttr {
         builder.build()
     }
 }
+
+#[derive(Clone, Copy)]
+pub struct LoadAttr(u32);
+
+impl LoadAttr {
+    #[must_use]
+    pub const fn new(value: u32) -> Self {
+        Self(value)
+    }
+
+    #[must_use]
+    pub fn builder() -> LoadAttrBuilder {
+        LoadAttrBuilder::default()
+    }
+
+    #[must_use]
+    pub const fn name_idx(self) -> u32 {
+        self.0 >> 1
+    }
+
+    #[must_use]
+    pub const fn is_method(self) -> bool {
+        (self.0 & 1) == 1
+    }
+}
+
+impl OpArgType for LoadAttr {}
+
+impl From<u32> for LoadAttr {
+    fn from(value: u32) -> Self {
+        Self::new(value)
+    }
+}
+
+impl From<LoadAttr> for u32 {
+    fn from(value: LoadAttr) -> Self {
+        value.0
+    }
+}
+
+#[derive(Clone, Copy, Default)]
+pub struct LoadAttrBuilder {
+    name_idx: u32,
+    is_method: bool,
+}
+
+impl LoadAttrBuilder {
+    #[must_use]
+    pub const fn build(self) -> LoadAttr {
+        let value = (self.name_idx << 1) | (self.is_method as u32);
+        LoadAttr::new(value)
+    }
+
+    #[must_use]
+    pub const fn name_idx(mut self, value: u32) -> Self {
+        self.name_idx = value;
+        self
+    }
+
+    #[must_use]
+    pub const fn is_method(mut self, value: bool) -> Self {
+        self.is_method = value;
+        self
+    }
+}
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -10,7 +10,7 @@ use crate::{
         function::{PyCell, PyCellRef, PyFunction},
         tuple::{PyTuple, PyTupleRef},
     },
-    bytecode::{self, Instruction, LoadSuperAttr},
+    bytecode::{self, Instruction, LoadAttr, LoadSuperAttr},
     convert::{IntoObject, ToPyResult},
     coroutine::Coro,
     exceptions::ExceptionCtor,
@@ -2890,12 +2890,11 @@ impl ExecutingFrame<'_> {
         Ok(None)
     }
 
-    fn load_attr(&mut self, vm: &VirtualMachine, oparg: u32) -> FrameResult {
-        let (name_idx, is_method) = bytecode::decode_load_attr_arg(oparg);
-        let attr_name = self.code.names[name_idx as usize];
+    fn load_attr(&mut self, vm: &VirtualMachine, oparg: LoadAttr) -> FrameResult {
+        let attr_name = self.code.names[oparg.name_idx() as usize];
         let parent = self.pop_value();
 
-        if is_method {
+        if oparg.is_method() {
             // Method call: push [method, self_or_null]
             let method = PyMethod::get(parent.clone(), attr_name, vm)?;
             match method {
