more optimization (#6860)

Lib/test/test_exceptions.py

@@ -1798,7 +1798,6 @@ class ExceptionTests(unittest.TestCase):
                 next(i)
                 next(i)
 
-    @unittest.expectedFailure # TODO: RUSTPYTHON
     @unittest.skipUnless(__debug__, "Won't work if __debug__ is False")
     def test_assert_shadowing(self):
         # Shadowing AssertionError would cause the assert statement to

crates/codegen/src/compile.rs

@@ -36,6 +36,51 @@ use rustpython_compiler_core::{
 };
 use rustpython_wtf8::Wtf8Buf;
 
+/// Extension trait for `ast::Expr` to add constant checking methods
+trait ExprExt {
+    /// Check if an expression is a constant literal
+    fn is_constant(&self) -> bool;
+
+    /// Check if a slice expression has all constant elements
+    fn is_constant_slice(&self) -> bool;
+
+    /// Check if we should use BINARY_SLICE/STORE_SLICE optimization
+    fn should_use_slice_optimization(&self) -> bool;
+}
+
+impl ExprExt for ast::Expr {
+    fn is_constant(&self) -> bool {
+        matches!(
+            self,
+            ast::Expr::NumberLiteral(_)
+                | ast::Expr::StringLiteral(_)
+                | ast::Expr::BytesLiteral(_)
+                | ast::Expr::NoneLiteral(_)
+                | ast::Expr::BooleanLiteral(_)
+                | ast::Expr::EllipsisLiteral(_)
+        )
+    }
+
+    fn is_constant_slice(&self) -> bool {
+        match self {
+            ast::Expr::Slice(s) => {
+                let lower_const =
+                    s.lower.is_none() || s.lower.as_deref().is_some_and(|e| e.is_constant());
+                let upper_const =
+                    s.upper.is_none() || s.upper.as_deref().is_some_and(|e| e.is_constant());
+                let step_const =
+                    s.step.is_none() || s.step.as_deref().is_some_and(|e| e.is_constant());
+                lower_const && upper_const && step_const
+            }
+            _ => false,
+        }
+    }
+
+    fn should_use_slice_optimization(&self) -> bool {
+        !self.is_constant_slice() && matches!(self, ast::Expr::Slice(s) if s.step.is_none())
+    }
+}
+
 const MAXBLOCKS: usize = 20;
 
 #[derive(Debug, Clone, Copy)]
@@ -425,39 +470,25 @@ impl Compiler {
         }
     }
 
-    /// Check if the slice is a two-element slice (no step)
-    // = is_two_element_slice
-    const fn is_two_element_slice(slice: &ast::Expr) -> bool {
-        matches!(slice, ast::Expr::Slice(s) if s.step.is_none())
-    }
-
-    /// Compile a slice expression
-    // = compiler_slice
-    fn compile_slice(&mut self, s: &ast::ExprSlice) -> CompileResult<BuildSliceArgCount> {
-        // Compile lower
+    /// Compile just start and stop of a slice (for BINARY_SLICE/STORE_SLICE)
+    // = codegen_slice_two_parts
+    fn compile_slice_two_parts(&mut self, s: &ast::ExprSlice) -> CompileResult<()> {
+        // Compile lower (or None)
         if let Some(lower) = &s.lower {
             self.compile_expression(lower)?;
         } else {
             self.emit_load_const(ConstantData::None);
         }
 
-        // Compile upper
+        // Compile upper (or None)
         if let Some(upper) = &s.upper {
             self.compile_expression(upper)?;
         } else {
             self.emit_load_const(ConstantData::None);
         }
 
-        Ok(match &s.step {
-            Some(step) => {
-                // Compile step if present
-                self.compile_expression(step)?;
-                BuildSliceArgCount::Three
-            }
-            None => BuildSliceArgCount::Two,
-        })
+        Ok(())
     }
-
     /// Compile a subscript expression
     // = compiler_subscript
     fn compile_subscript(
@@ -480,27 +511,20 @@ impl Compiler {
         // VISIT(c, expr, e->v.Subscript.value)
         self.compile_expression(value)?;
 
-        // Handle two-element slice (for Load/Store, not Del)
-        if Self::is_two_element_slice(slice) && !matches!(ctx, ast::ExprContext::Del) {
-            let argc = match slice {
-                ast::Expr::Slice(s) => self.compile_slice(s)?,
+        // Handle two-element non-constant slice with BINARY_SLICE/STORE_SLICE
+        if slice.should_use_slice_optimization() && !matches!(ctx, ast::ExprContext::Del) {
+            match slice {
+                ast::Expr::Slice(s) => self.compile_slice_two_parts(s)?,
                 _ => unreachable!(
-                    "is_two_element_slice should only return true for ast::Expr::Slice"
+                    "should_use_slice_optimization should only return true for ast::Expr::Slice"
                 ),
             };
             match ctx {
                 ast::ExprContext::Load => {
-                    emit!(self, Instruction::BuildSlice { argc });
-                    emit!(
-                        self,
-                        Instruction::BinaryOp {
-                            op: BinaryOperator::Subscr
-                        }
-                    );
+                    emit!(self, Instruction::BinarySlice);
                 }
                 ast::ExprContext::Store => {
-                    emit!(self, Instruction::BuildSlice { argc });
-                    emit!(self, Instruction::StoreSubscr);
+                    emit!(self, Instruction::StoreSlice);
                 }
                 _ => unreachable!(),
             }
@@ -1269,8 +1293,12 @@ impl Compiler {
         emit!(self, Instruction::PopJumpIfFalse { target: body_block });
 
         // Raise NotImplementedError
-        let not_implemented_error = self.name("NotImplementedError");
-        emit!(self, Instruction::LoadGlobal(not_implemented_error));
+        emit!(
+            self,
+            Instruction::LoadCommonConstant {
+                idx: bytecode::CommonConstant::NotImplementedError
+            }
+        );
         emit!(
             self,
             Instruction::RaiseVarargs {
@@ -2345,8 +2373,12 @@ impl Compiler {
                     let after_block = self.new_block();
                     self.compile_jump_if(test, true, after_block)?;
 
-                    let assertion_error = self.name("AssertionError");
-                    emit!(self, Instruction::LoadGlobal(assertion_error));
+                    emit!(
+                        self,
+                        Instruction::LoadCommonConstant {
+                            idx: bytecode::CommonConstant::AssertionError
+                        }
+                    );
                     emit!(self, Instruction::PushNull);
                     match msg {
                         Some(e) => {
@@ -3326,11 +3358,9 @@ impl Compiler {
             // ADDOP_I(c, loc, COPY, 1);
             // ADDOP_JUMP(c, loc, POP_JUMP_IF_NONE, no_match);
             emit!(self, Instruction::Copy { index: 1 });
-            self.emit_load_const(ConstantData::None);
-            emit!(self, Instruction::IsOp(bytecode::Invert::No)); // is None?
             emit!(
                 self,
-                Instruction::PopJumpIfTrue {
+                Instruction::PopJumpIfNone {
                     target: no_match_block
                 }
             );
@@ -3455,11 +3485,9 @@ impl Compiler {
         // Stack: [prev_exc, result, result]
 
         // POP_JUMP_IF_NOT_NONE reraise
-        self.emit_load_const(ConstantData::None);
-        emit!(self, Instruction::IsOp(bytecode::Invert::Yes)); // is not None?
         emit!(
             self,
-            Instruction::PopJumpIfTrue {
+            Instruction::PopJumpIfNotNone {
                 target: reraise_block
             }
         );

crates/codegen/src/ir.rs

@@ -161,6 +161,7 @@ impl CodeInfo {
     ) -> crate::InternalResult<CodeObject> {
         if opts.optimize > 0 {
             self.dce();
+            self.peephole_optimize();
         }
 
         let max_stackdepth = self.max_stackdepth()?;
@@ -403,6 +404,70 @@ impl CodeInfo {
         }
     }
 
+    /// Peephole optimization: combine consecutive instructions into super-instructions
+    fn peephole_optimize(&mut self) {
+        for block in &mut self.blocks {
+            let mut i = 0;
+            while i + 1 < block.instructions.len() {
+                let combined = {
+                    let curr = &block.instructions[i];
+                    let next = &block.instructions[i + 1];
+
+                    // Only combine if both are real instructions (not pseudo)
+                    let (Some(curr_instr), Some(next_instr)) =
+                        (curr.instr.real(), next.instr.real())
+                    else {
+                        i += 1;
+                        continue;
+                    };
+
+                    match (curr_instr, next_instr) {
+                        // LoadFast + LoadFast -> LoadFastLoadFast (if both indices < 16)
+                        (Instruction::LoadFast(_), Instruction::LoadFast(_)) => {
+                            let idx1 = curr.arg.0;
+                            let idx2 = next.arg.0;
+                            if idx1 < 16 && idx2 < 16 {
+                                let packed = (idx1 << 4) | idx2;
+                                Some((
+                                    Instruction::LoadFastLoadFast { arg: Arg::marker() },
+                                    OpArg(packed),
+                                ))
+                            } else {
+                                None
+                            }
+                        }
+                        // StoreFast + StoreFast -> StoreFastStoreFast (if both indices < 16)
+                        (Instruction::StoreFast(_), Instruction::StoreFast(_)) => {
+                            let idx1 = curr.arg.0;
+                            let idx2 = next.arg.0;
+                            if idx1 < 16 && idx2 < 16 {
+                                let packed = (idx1 << 4) | idx2;
+                                Some((
+                                    Instruction::StoreFastStoreFast { arg: Arg::marker() },
+                                    OpArg(packed),
+                                ))
+                            } else {
+                                None
+                            }
+                        }
+                        _ => None,
+                    }
+                };
+
+                if let Some((new_instr, new_arg)) = combined {
+                    // Combine: keep first instruction's location, replace with combined instruction
+                    block.instructions[i].instr = new_instr.into();
+                    block.instructions[i].arg = new_arg;
+                    // Remove the second instruction
+                    block.instructions.remove(i + 1);
+                    // Don't increment i - check if we can combine again with the next instruction
+                } else {
+                    i += 1;
+                }
+            }
+        }
+    }
+
     fn max_stackdepth(&self) -> crate::InternalResult<u32> {
         let mut maxdepth = 0u32;
         let mut stack = Vec::with_capacity(self.blocks.len());

crates/compiler-core/src/bytecode.rs

@@ -21,7 +21,7 @@ pub use crate::bytecode::{
         encode_load_super_attr_arg,
     },
     oparg::{
-        BinaryOperator, BuildSliceArgCount, ComparisonOperator, ConvertValueOparg,
+        BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator, ConvertValueOparg,
         IntrinsicFunction1, IntrinsicFunction2, Invert, Label, MakeFunctionFlags, NameIdx, OpArg,
         OpArgByte, OpArgState, OpArgType, RaiseKind, ResumeType, SpecialMethod, UnpackExArgs,
     },

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

@@ -4,9 +4,10 @@ use crate::{
     bytecode::{
         BorrowedConstant, Constant, InstrDisplayContext,
         oparg::{
-            BinaryOperator, BuildSliceArgCount, ComparisonOperator, ConvertValueOparg,
-            IntrinsicFunction1, IntrinsicFunction2, Invert, Label, MakeFunctionFlags, NameIdx,
-            OpArg, OpArgByte, OpArgType, RaiseKind, SpecialMethod, UnpackExArgs,
+            BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator,
+            ConvertValueOparg, IntrinsicFunction1, IntrinsicFunction2, Invert, Label,
+            MakeFunctionFlags, NameIdx, OpArg, OpArgByte, OpArgType, RaiseKind, SpecialMethod,
+            UnpackExArgs,
         },
     },
     marshal::MarshalError,
@@ -22,8 +23,8 @@ use crate::{
 #[repr(u8)]
 pub enum Instruction {
     // No-argument instructions (opcode < HAVE_ARGUMENT=44)
-    Cache = 0,       // Placeholder
-    BinarySlice = 1, // Placeholder
+    Cache = 0, // Placeholder
+    BinarySlice = 1,
     BuildTemplate = 2,
     BinaryOpInplaceAddUnicode = 3, // Placeholder
     CallFunctionEx = 4,
@@ -59,7 +60,7 @@ pub enum Instruction {
     ReturnGenerator = 34,
     ReturnValue = 35,
     SetupAnnotations = 36,
-    StoreSlice = 37, // Placeholder
+    StoreSlice = 37,
     StoreSubscr = 38,
     ToBool = 39,
     UnaryInvert = 40,
@@ -122,7 +123,7 @@ pub enum Instruction {
     } = 59,
     CopyFreeVars {
         count: Arg<u32>,
-    } = 60, // Placeholder
+    } = 60,
     DeleteAttr {
         idx: Arg<NameIdx>,
     } = 61,
@@ -168,8 +169,8 @@ pub enum Instruction {
         idx: Arg<NameIdx>,
     } = 80,
     LoadCommonConstant {
-        idx: Arg<u32>,
-    } = 81, // Placeholder
+        idx: Arg<CommonConstant>,
+    } = 81,
     LoadConst {
         idx: Arg<u32>,
     } = 82,
@@ -180,10 +181,10 @@ pub enum Instruction {
     LoadFastBorrowLoadFastBorrow {
         arg: Arg<u32>,
     } = 87, // Placeholder
-    LoadFastCheck(Arg<NameIdx>) = 88,  // Placeholder
+    LoadFastCheck(Arg<NameIdx>) = 88,
     LoadFastLoadFast {
         arg: Arg<u32>,
-    } = 89, // Placeholder
+    } = 89,
     LoadFromDictOrDeref(Arg<NameIdx>) = 90,
     LoadFromDictOrGlobals(Arg<NameIdx>) = 91, // Placeholder
     LoadGlobal(Arg<NameIdx>) = 92,
@@ -197,7 +198,7 @@ pub enum Instruction {
     LoadSuperAttr {
         arg: Arg<u32>,
     } = 96,
-    MakeCell(Arg<NameIdx>) = 97, // Placeholder
+    MakeCell(Arg<NameIdx>) = 97,
     MapAdd {
         i: Arg<u32>,
     } = 98,
@@ -207,10 +208,10 @@ pub enum Instruction {
     } = 100,
     PopJumpIfNone {
         target: Arg<Label>,
-    } = 101, // Placeholder
+    } = 101,
     PopJumpIfNotNone {
         target: Arg<Label>,
-    } = 102, // Placeholder
+    } = 102,
     PopJumpIfTrue {
         target: Arg<Label>,
     } = 103,
@@ -243,7 +244,7 @@ pub enum Instruction {
     } = 113,
     StoreFastStoreFast {
         arg: Arg<u32>,
-    } = 114, // Placeholder
+    } = 114,
     StoreGlobal(Arg<NameIdx>) = 115,
     StoreName(Arg<NameIdx>) = 116,
     Swap {
@@ -616,14 +617,14 @@ impl InstructionMetadata for Instruction {
             Self::JumpBackwardNoInterrupt { .. } => 0,
             Self::JumpBackward { .. } => 0,
             Self::JumpForward { .. } => 0,
-            Self::LoadFastCheck(_) => 0,
+            Self::LoadFastCheck(_) => 1,
             Self::LoadFastLoadFast { .. } => 2,
             Self::LoadFastBorrowLoadFastBorrow { .. } => 2,
             Self::LoadFromDictOrGlobals(_) => 0,
             Self::MakeCell(_) => 0,
-            Self::StoreFastStoreFast { .. } => 0,
-            Self::PopJumpIfNone { .. } => -1,    // (value -- )
-            Self::PopJumpIfNotNone { .. } => -1, // (value -- )
+            Self::StoreFastStoreFast { .. } => -2, // pops 2 values
+            Self::PopJumpIfNone { .. } => -1,      // (value -- )
+            Self::PopJumpIfNotNone { .. } => -1,   // (value -- )
             Self::BinaryOpAddFloat => 0,
             Self::BinaryOpAddInt => 0,
             Self::BinaryOpAddUnicode => 0,

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

@@ -548,6 +548,38 @@ impl fmt::Display for SpecialMethod {
     }
 }
 
+op_arg_enum!(
+    /// Common constants for LOAD_COMMON_CONSTANT opcode.
+    /// pycore_opcode_utils.h CONSTANT_*
+    #[repr(u8)]
+    #[derive(Debug, Copy, Clone, PartialEq, Eq)]
+    pub enum CommonConstant {
+        /// `AssertionError` exception type
+        AssertionError = 0,
+        /// `NotImplementedError` exception type
+        NotImplementedError = 1,
+        /// Built-in `tuple` type
+        BuiltinTuple = 2,
+        /// Built-in `all` function
+        BuiltinAll = 3,
+        /// Built-in `any` function
+        BuiltinAny = 4,
+    }
+);
+
+impl fmt::Display for CommonConstant {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let name = match self {
+            Self::AssertionError => "AssertionError",
+            Self::NotImplementedError => "NotImplementedError",
+            Self::BuiltinTuple => "tuple",
+            Self::BuiltinAll => "all",
+            Self::BuiltinAny => "any",
+        };
+        write!(f, "{name}")
+    }
+}
+
 /// Specifies if a slice is built with either 2 or 3 arguments.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum BuildSliceArgCount {

crates/stdlib/src/_opcode.rs

@@ -151,7 +151,7 @@ mod _opcode {
     // prepare specialization
     #[pyattr]
     const ENABLE_SPECIALIZATION: i8 = 1;
-    #[allow(dead_code)]
+    #[pyattr]
     const ENABLE_SPECIALIZATION_FT: i8 = 1;
 
     #[derive(FromArgs)]

crates/vm/src/frame.rs

@@ -618,7 +618,22 @@ impl ExecutingFrame<'_> {
 
         match instruction {
             Instruction::BinaryOp { op } => self.execute_bin_op(vm, op.get(arg)),
-
+            Instruction::BinarySlice => {
+                // Stack: [container, start, stop] -> [result]
+                let stop = self.pop_value();
+                let start = self.pop_value();
+                let container = self.pop_value();
+                let slice: PyObjectRef = PySlice {
+                    start: Some(start),
+                    stop,
+                    step: None,
+                }
+                .into_ref(&vm.ctx)
+                .into();
+                let result = container.get_item(&*slice, vm)?;
+                self.push_value(result);
+                Ok(None)
+            }
             Instruction::BuildList { size } => {
                 let sz = size.get(arg) as usize;
                 let elements = self.pop_multiple(sz).collect();
@@ -787,6 +802,10 @@ impl ExecutingFrame<'_> {
                 self.push_value_opt(value);
                 Ok(None)
             }
+            Instruction::CopyFreeVars { .. } => {
+                // Free vars are already set up at frame creation time in RustPython
+                Ok(None)
+            }
             Instruction::DeleteAttr { idx } => self.delete_attr(vm, idx.get(arg)),
             Instruction::DeleteDeref(i) => {
                 self.cells_frees[i.get(arg) as usize].set(None);
@@ -1199,6 +1218,22 @@ impl ExecutingFrame<'_> {
                 self.push_value(self.code.constants[idx.get(arg) as usize].clone().into());
                 Ok(None)
             }
+            Instruction::LoadCommonConstant { idx } => {
+                use bytecode::CommonConstant;
+                let value = match idx.get(arg) {
+                    CommonConstant::AssertionError => {
+                        vm.ctx.exceptions.assertion_error.to_owned().into()
+                    }
+                    CommonConstant::NotImplementedError => {
+                        vm.ctx.exceptions.not_implemented_error.to_owned().into()
+                    }
+                    CommonConstant::BuiltinTuple => vm.ctx.types.tuple_type.to_owned().into(),
+                    CommonConstant::BuiltinAll => vm.builtins.get_attr("all", vm)?,
+                    CommonConstant::BuiltinAny => vm.builtins.get_attr("any", vm)?,
+                };
+                self.push_value(value);
+                Ok(None)
+            }
             Instruction::LoadSmallInt { idx } => {
                 // Push small integer (-5..=256) directly without constant table lookup
                 let value = vm.ctx.new_int(idx.get(arg) as i32);
@@ -1241,6 +1276,52 @@ impl ExecutingFrame<'_> {
                 self.push_value(x);
                 Ok(None)
             }
+            Instruction::LoadFastCheck(idx) => {
+                // Same as LoadFast but explicitly checks for unbound locals
+                // (LoadFast in RustPython already does this check)
+                let idx = idx.get(arg) as usize;
+                let x = self.fastlocals.lock()[idx].clone().ok_or_else(|| {
+                    vm.new_exception_msg(
+                        vm.ctx.exceptions.unbound_local_error.to_owned(),
+                        format!(
+                            "local variable '{}' referenced before assignment",
+                            self.code.varnames[idx]
+                        ),
+                    )
+                })?;
+                self.push_value(x);
+                Ok(None)
+            }
+            Instruction::LoadFastLoadFast { arg: packed } => {
+                // 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 fastlocals = self.fastlocals.lock();
+                let x1 = fastlocals[idx1].clone().ok_or_else(|| {
+                    vm.new_exception_msg(
+                        vm.ctx.exceptions.unbound_local_error.to_owned(),
+                        format!(
+                            "local variable '{}' referenced before assignment",
+                            self.code.varnames[idx1]
+                        ),
+                    )
+                })?;
+                let x2 = fastlocals[idx2].clone().ok_or_else(|| {
+                    vm.new_exception_msg(
+                        vm.ctx.exceptions.unbound_local_error.to_owned(),
+                        format!(
+                            "local variable '{}' referenced before assignment",
+                            self.code.varnames[idx2]
+                        ),
+                    )
+                })?;
+                drop(fastlocals);
+                self.push_value(x1);
+                self.push_value(x2);
+                Ok(None)
+            }
             Instruction::LoadGlobal(idx) => {
                 let name = &self.code.names[idx.get(arg) as usize];
                 let x = self.load_global_or_builtin(name, vm)?;
@@ -1294,6 +1375,10 @@ impl ExecutingFrame<'_> {
                 Ok(None)
             }
             Instruction::MakeFunction => self.execute_make_function(vm),
+            Instruction::MakeCell(_) => {
+                // Cell creation is handled at frame creation time in RustPython
+                Ok(None)
+            }
             Instruction::MapAdd { i } => {
                 let value = self.pop_value();
                 let key = self.pop_value();
@@ -1718,6 +1803,21 @@ impl ExecutingFrame<'_> {
                 self.push_value(load_value);
                 Ok(None)
             }
+            Instruction::StoreFastStoreFast { arg: packed } => {
+                // Store two values to two local variables at once
+                // STORE_FAST idx1 executes first: pops TOS -> locals[idx1]
+                // STORE_FAST idx2 executes second: pops new TOS -> locals[idx2]
+                // oparg encoding: (idx1 << 4) | idx2
+                let oparg = packed.get(arg);
+                let idx1 = (oparg >> 4) as usize;
+                let idx2 = (oparg & 15) as usize;
+                let value1 = self.pop_value(); // TOS -> idx1
+                let value2 = self.pop_value(); // second -> idx2
+                let mut fastlocals = self.fastlocals.lock();
+                fastlocals[idx1] = Some(value1);
+                fastlocals[idx2] = Some(value2);
+                Ok(None)
+            }
             Instruction::StoreGlobal(idx) => {
                 let value = self.pop_value();
                 self.globals
@@ -1730,6 +1830,22 @@ impl ExecutingFrame<'_> {
                 self.locals.mapping().ass_subscript(name, Some(value), vm)?;
                 Ok(None)
             }
+            Instruction::StoreSlice => {
+                // Stack: [value, container, start, stop] -> []
+                let stop = self.pop_value();
+                let start = self.pop_value();
+                let container = self.pop_value();
+                let value = self.pop_value();
+                let slice: PyObjectRef = PySlice {
+                    start: Some(start),
+                    stop,
+                    step: None,
+                }
+                .into_ref(&vm.ctx)
+                .into();
+                container.set_item(&*slice, value, vm)?;
+                Ok(None)
+            }
             Instruction::StoreSubscr => self.execute_store_subscript(vm),
             Instruction::Swap { index } => {
                 let len = self.state.stack.len();