From 1ac52b918752f90dbf5ba710d7a562e83d0176bb Mon Sep 17 00:00:00 2001
From: Shahar Naveh <50263213+ShaharNaveh@users.noreply.github.com>
Date: Fri, 20 Feb 2026 01:19:12 +0200
Subject: [PATCH] Newtype oparg for `StoreFastLoadFast` (#7191)

---
 .../compiler-core/src/bytecode/instruction.rs | 15 ++++---
 crates/compiler-core/src/bytecode/oparg.rs    | 41 +++++++++++++++++++
 crates/vm/src/frame.rs                        | 10 ++---
 3 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs
index fd2d06761..7ac117cb5 100644
--- a/crates/compiler-core/src/bytecode/instruction.rs
+++ b/crates/compiler-core/src/bytecode/instruction.rs
@@ -7,7 +7,7 @@ use crate::{
             BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator,
             ConvertValueOparg, IntrinsicFunction1, IntrinsicFunction2, Invert, Label, LoadAttr,
             LoadSuperAttr, MakeFunctionFlags, NameIdx, OpArg, OpArgByte, OpArgType, RaiseKind,
-            SpecialMethod, UnpackExArgs,
+            SpecialMethod, StoreFastLoadFast, UnpackExArgs,
         },
     },
     marshal::MarshalError,
@@ -241,8 +241,7 @@ pub enum Instruction {
     StoreDeref(Arg<NameIdx>) = 111,
     StoreFast(Arg<NameIdx>) = 112,
     StoreFastLoadFast {
-        store_idx: Arg<NameIdx>,
-        load_idx: Arg<NameIdx>,
+        var_nums: Arg<StoreFastLoadFast>,
     } = 113,
     StoreFastStoreFast {
         arg: Arg<u32>,
@@ -901,12 +900,12 @@ impl InstructionMetadata for Instruction {
             Self::StoreAttr { idx } => w!(STORE_ATTR, name = idx),
             Self::StoreDeref(idx) => w!(STORE_DEREF, cell_name = idx),
             Self::StoreFast(idx) => w!(STORE_FAST, varname = idx),
-            Self::StoreFastLoadFast {
-                store_idx,
-                load_idx,
-            } => {
+            Self::StoreFastLoadFast { var_nums } => {
+                let oparg = var_nums.get(arg);
+                let store_idx = oparg.store_idx();
+                let load_idx = oparg.load_idx();
                 write!(f, "STORE_FAST_LOAD_FAST")?;
-                write!(f, " ({}, {})", store_idx.get(arg), load_idx.get(arg))
+                write!(f, " ({}, {})", store_idx, load_idx)
             }
             Self::StoreGlobal(idx) => w!(STORE_GLOBAL, name = idx),
             Self::StoreName(idx) => w!(STORE_NAME, name = idx),
diff --git a/crates/compiler-core/src/bytecode/oparg.rs b/crates/compiler-core/src/bytecode/oparg.rs
index 849b3897f..729b84db5 100644
--- a/crates/compiler-core/src/bytecode/oparg.rs
+++ b/crates/compiler-core/src/bytecode/oparg.rs
@@ -323,6 +323,47 @@ impl fmt::Display for Label {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
+#[repr(transparent)]
+pub struct StoreFastLoadFast(u32);
+
+impl StoreFastLoadFast {
+    #[must_use]
+    pub const fn new(value: u32) -> Self {
+        Self(value)
+    }
+
+    #[must_use]
+    pub const fn store_idx(self) -> NameIdx {
+        self.0 >> 4
+    }
+
+    #[must_use]
+    pub const fn load_idx(self) -> NameIdx {
+        self.0 & 15
+    }
+}
+
+impl From<u32> for StoreFastLoadFast {
+    fn from(value: u32) -> Self {
+        Self::new(value)
+    }
+}
+
+impl From<StoreFastLoadFast> for u32 {
+    fn from(value: StoreFastLoadFast) -> Self {
+        value.0
+    }
+}
+
+impl OpArgType for StoreFastLoadFast {}
+
+impl fmt::Display for StoreFastLoadFast {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
 oparg_enum!(
     /// The kind of Raise that occurred.
     #[derive(Copy, Clone, Debug, PartialEq, Eq)]
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
index 62df1b298..f03aeeaa0 100644
--- a/crates/vm/src/frame.rs
+++ b/crates/vm/src/frame.rs
@@ -2091,15 +2091,13 @@ impl ExecutingFrame<'_> {
                 self.fastlocals.lock()[idx.get(arg) as usize] = Some(value);
                 Ok(None)
             }
-            Instruction::StoreFastLoadFast {
-                store_idx,
-                load_idx,
-            } => {
+            Instruction::StoreFastLoadFast { var_nums } => {
                 // Store to one slot and load from another (often the same) - for inlined comprehensions
                 let value = self.pop_value();
                 let mut locals = self.fastlocals.lock();
-                locals[store_idx.get(arg) as usize] = Some(value);
-                let load_value = locals[load_idx.get(arg) as usize]
+                let oparg = var_nums.get(arg);
+                locals[oparg.store_idx() as usize] = Some(value);
+                let load_value = locals[oparg.load_idx() as usize]
                     .clone()
                     .expect("StoreFastLoadFast: load slot should have value after store");
                 drop(locals);