Ensure BuildSlice oparg to be either 2 or 3 (RustPython#6313)

* Force `BuildSlice` oparg to be either 2 or 3

* `compile_slice` to return `BuildSliceArgCount`

Author: ShaharNaveh

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;
@@ -401,7 +401,7 @@ impl Compiler {
 
     /// Compile a slice expression
     // = compiler_slice
-    fn compile_slice(&mut self, s: &ExprSlice) -> CompileResult<u32> {
+    fn compile_slice(&mut self, s: &ExprSlice) -> CompileResult<BuildSliceArgCount> {
         // Compile lower
         if let Some(lower) = &s.lower {
             self.compile_expression(lower)?;
@@ -416,13 +416,14 @@ impl Compiler {
             self.emit_load_const(ConstantData::None);
         }
 
-        // Compile step if present
-        if let Some(step) = &s.step {
-            self.compile_expression(step)?;
-            Ok(3) // Three values on stack
-        } else {
-            Ok(2) // Two values on stack
-        }
+        Ok(match &s.step {
+            Some(step) => {
+                // Compile step if present
+                self.compile_expression(step)?;
+                BuildSliceArgCount::Three
+            }
+            None => BuildSliceArgCount::Two,
+        })
     }
 
     /// Compile a subscript expression
@@ -449,19 +450,19 @@ impl Compiler {
 
         // Handle two-element slice (for Load/Store, not Del)
         if Self::is_two_element_slice(slice) && !matches!(ctx, ExprContext::Del) {
-            let n = match slice {
+            let argc = match slice {
                 Expr::Slice(s) => self.compile_slice(s)?,
                 _ => unreachable!("is_two_element_slice should only return true for Expr::Slice"),
             };
             match ctx {
                 ExprContext::Load => {
                     // CPython uses BINARY_SLICE
-                    emit!(self, Instruction::BuildSlice { step: n == 3 });
+                    emit!(self, Instruction::BuildSlice { argc });
                     emit!(self, Instruction::Subscript);
                 }
                 ExprContext::Store => {
                     // CPython uses STORE_SLICE
-                    emit!(self, Instruction::BuildSlice { step: n == 3 });
+                    emit!(self, Instruction::BuildSlice { argc });
                     emit!(self, Instruction::StoreSubscript);
                 }
                 _ => unreachable!(),
@@ -4587,8 +4588,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() {

crates/compiler-core/src/bytecode.rs

@@ -10,7 +10,7 @@ use itertools::Itertools;
 use malachite_bigint::BigInt;
 use num_complex::Complex64;
 use rustpython_wtf8::{Wtf8, Wtf8Buf};
-use std::{collections::BTreeSet, fmt, hash, marker::PhantomData, mem, ops::Deref};
+use std::{collections::BTreeSet, fmt, hash, marker::PhantomData, mem, num::NonZeroU8, ops::Deref};
 
 #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
 #[repr(i8)]
@@ -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,48 @@ op_arg_enum!(
     }
 );
 
+/// Specifies if a slice is built with either 2 or 3 arguments.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum BuildSliceArgCount {
+    /// ```py
+    /// x[5:10]
+    /// ```
+    Two,
+    /// ```py
+    /// x[5:10:2]
+    /// ```
+    Three,
+}
+
+impl OpArgType for BuildSliceArgCount {
+    #[inline(always)]
+    fn from_op_arg(x: u32) -> Option<Self> {
+        Some(match x {
+            2 => Self::Two,
+            3 => Self::Three,
+            _ => return None,
+        })
+    }
+
+    #[inline(always)]
+    fn to_op_arg(self) -> u32 {
+        u32::from(self.argc().get())
+    }
+}
+
+impl BuildSliceArgCount {
+    /// Get the numeric value of `Self`.
+    #[must_use]
+    pub const fn argc(self) -> NonZeroU8 {
+        let inner = match self {
+            Self::Two => 2,
+            Self::Three => 3,
+        };
+        // Safety: `inner` can be either 2 or 3.
+        unsafe { NonZeroU8::new_unchecked(inner) }
+    }
+}
+
 #[derive(Copy, Clone)]
 pub struct UnpackExArgs {
     pub before: u8,
@@ -1547,7 +1588,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).argc().get() as i32)
+            }
             ListAppend { .. } | SetAdd { .. } => -1,
             MapAdd { .. } => -2,
             PrintExpr => -1,
@@ -1734,7 +1779,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),

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();