Force BuildSlice oparg to be either 2 or 3

ShaharNaveh · ShaharNaveh · commit 29ee65688f09 · 2025-11-29T15:56:59.000+02:00
diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs
@@ -35,8 +35,8 @@ use ruff_text_size::{Ranged, TextRange};
 use rustpython_compiler_core::{
     Mode, OneIndexed, PositionEncoding, SourceFile, SourceLocation,
     bytecode::{
-        self, Arg as OpArgMarker, BinaryOperator, CodeObject, ComparisonOperator, ConstantData,
-        Instruction, Invert, OpArg, OpArgType, UnpackExArgs,
+        self, Arg as OpArgMarker, BinaryOperator, BuildSliceArgCount, CodeObject,
+        ComparisonOperator, ConstantData, Instruction, Invert, OpArg, OpArgType, UnpackExArgs,
     },
 };
 use rustpython_wtf8::Wtf8Buf;
@@ -456,12 +456,22 @@ impl Compiler {
             match ctx {
                 ExprContext::Load => {
                     // CPython uses BINARY_SLICE
-                    emit!(self, Instruction::BuildSlice { step: n == 3 });
+                    emit!(
+                        self,
+                        Instruction::BuildSlice {
+                            argc: BuildSliceArgCount::from_op_arg(n).unwrap()
+                        }
+                    );
                     emit!(self, Instruction::Subscript);
                 }
                 ExprContext::Store => {
                     // CPython uses STORE_SLICE
-                    emit!(self, Instruction::BuildSlice { step: n == 3 });
+                    emit!(
+                        self,
+                        Instruction::BuildSlice {
+                            argc: BuildSliceArgCount::from_op_arg(n).unwrap()
+                        }
+                    );
                     emit!(self, Instruction::StoreSubscript);
                 }
                 _ => unreachable!(),
@@ -4587,8 +4597,11 @@ impl Compiler {
                 if let Some(step) = step {
                     self.compile_expression(step)?;
                 }
-                let step = step.is_some();
-                emit!(self, Instruction::BuildSlice { step });
+                let argc = match step {
+                    Some(_) => BuildSliceArgCount::Three,
+                    None => BuildSliceArgCount::Two,
+                };
+                emit!(self, Instruction::BuildSlice { argc });
             }
             Expr::Yield(ExprYield { value, .. }) => {
                 if !self.ctx.in_func() {
diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs
@@ -760,8 +760,7 @@ pub enum Instruction {
         index: Arg<u32>,
     },
     BuildSlice {
-        /// whether build a slice with a third step argument
-        step: Arg<bool>,
+        argc: Arg<BuildSliceArgCount>,
     },
     ListAppend {
         i: Arg<u32>,
@@ -1151,6 +1150,22 @@ op_arg_enum!(
     }
 );
 
+op_arg_enum!(
+    /// Specifies if a slice is built with either 2 or 3 arguments.
+    #[repr(u8)]
+    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
+    pub enum BuildSliceArgCount {
+        /// ```py
+        /// x[5:10]
+        /// ```
+        Two = 2,
+        /// ```py
+        /// x[5:10:2]
+        /// ```
+        Three = 3,
+    }
+);
+
 #[derive(Copy, Clone)]
 pub struct UnpackExArgs {
     pub before: u8,
@@ -1547,7 +1562,11 @@ impl Instruction {
                 -(nargs as i32) + 1
             }
             DictUpdate { .. } => -1,
-            BuildSlice { step } => -2 - (step.get(arg) as i32) + 1,
+            BuildSlice { argc } => {
+                // push 1
+                // pops either 2/3
+                1 - (argc.get(arg) as i32)
+            }
             ListAppend { .. } | SetAdd { .. } => -1,
             MapAdd { .. } => -2,
             PrintExpr => -1,
@@ -1734,7 +1753,7 @@ impl Instruction {
             BuildMap { size } => w!(BuildMap, size),
             BuildMapForCall { size } => w!(BuildMapForCall, size),
             DictUpdate { index } => w!(DictUpdate, index),
-            BuildSlice { step } => w!(BuildSlice, step),
+            BuildSlice { argc } => w!(BuildSlice, ?argc),
             ListAppend { i } => w!(ListAppend, i),
             SetAdd { i } => w!(SetAdd, i),
             MapAdd { i } => w!(MapAdd, i),
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -834,8 +834,8 @@ impl ExecutingFrame<'_> {
                 dict.merge_object(source, vm)?;
                 Ok(None)
             }
-            bytecode::Instruction::BuildSlice { step } => {
-                self.execute_build_slice(vm, step.get(arg))
+            bytecode::Instruction::BuildSlice { argc } => {
+                self.execute_build_slice(vm, argc.get(arg))
             }
             bytecode::Instruction::ListAppend { i } => {
                 let item = self.pop_value();
@@ -1777,8 +1777,15 @@ impl ExecutingFrame<'_> {
         Ok(None)
     }
 
-    fn execute_build_slice(&mut self, vm: &VirtualMachine, step: bool) -> FrameResult {
-        let step = if step { Some(self.pop_value()) } else { None };
+    fn execute_build_slice(
+        &mut self,
+        vm: &VirtualMachine,
+        argc: bytecode::BuildSliceArgCount,
+    ) -> FrameResult {
+        let step = match argc {
+            bytecode::BuildSliceArgCount::Two => None,
+            bytecode::BuildSliceArgCount::Three => Some(self.pop_value()),
+        };
         let stop = self.pop_value();
         let start = self.pop_value();
 
