Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion Lib/test/test_exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -1798,7 +1798,6 @@ def g():
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
Expand Down
116 changes: 72 additions & 44 deletions crates/codegen/src/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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(
Expand All @@ -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!(),
}
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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) => {
Expand Down Expand Up @@ -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
}
);
Expand Down Expand Up @@ -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
}
);
Expand Down
65 changes: 65 additions & 0 deletions crates/codegen/src/ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,7 @@ impl CodeInfo {
) -> crate::InternalResult<CodeObject> {
if opts.optimize > 0 {
self.dce();
self.peephole_optimize();
}

let max_stackdepth = self.max_stackdepth()?;
Expand Down Expand Up @@ -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());
Expand Down
2 changes: 1 addition & 1 deletion crates/compiler-core/src/bytecode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
Expand Down
Loading
Loading