Skip to content

Commit 7eceb14

Browse files
youknowoneyouknowone
authored
more optimization (#6860)
1 parent 6595d50 commit 7eceb14

File tree

8 files changed

+308
-67
lines changed

8 files changed

+308
-67
lines changed

‎Lib/test/test_exceptions.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1798,7 +1798,6 @@ def g():
17981798
next(i)
17991799
next(i)
18001800

1801-
@unittest.expectedFailure # TODO: RUSTPYTHON
18021801
@unittest.skipUnless(__debug__, "Won't work if __debug__ is False")
18031802
def test_assert_shadowing(self):
18041803
# Shadowing AssertionError would cause the assert statement to

‎crates/codegen/src/compile.rs‎

Lines changed: 72 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,51 @@ use rustpython_compiler_core::{
3636
};
3737
use rustpython_wtf8::Wtf8Buf;
3838

39+
/// Extension trait for `ast::Expr` to add constant checking methods
40+
trait ExprExt {
41+
/// Check if an expression is a constant literal
42+
fn is_constant(&self) -> bool;
43+
44+
/// Check if a slice expression has all constant elements
45+
fn is_constant_slice(&self) -> bool;
46+
47+
/// Check if we should use BINARY_SLICE/STORE_SLICE optimization
48+
fn should_use_slice_optimization(&self) -> bool;
49+
}
50+
51+
impl ExprExt for ast::Expr {
52+
fn is_constant(&self) -> bool {
53+
matches!(
54+
self,
55+
ast::Expr::NumberLiteral(_)
56+
| ast::Expr::StringLiteral(_)
57+
| ast::Expr::BytesLiteral(_)
58+
| ast::Expr::NoneLiteral(_)
59+
| ast::Expr::BooleanLiteral(_)
60+
| ast::Expr::EllipsisLiteral(_)
61+
)
62+
}
63+
64+
fn is_constant_slice(&self) -> bool {
65+
match self {
66+
ast::Expr::Slice(s) => {
67+
let lower_const =
68+
s.lower.is_none() || s.lower.as_deref().is_some_and(|e| e.is_constant());
69+
let upper_const =
70+
s.upper.is_none() || s.upper.as_deref().is_some_and(|e| e.is_constant());
71+
let step_const =
72+
s.step.is_none() || s.step.as_deref().is_some_and(|e| e.is_constant());
73+
lower_const && upper_const && step_const
74+
}
75+
_ => false,
76+
}
77+
}
78+
79+
fn should_use_slice_optimization(&self) -> bool {
80+
!self.is_constant_slice() && matches!(self, ast::Expr::Slice(s) if s.step.is_none())
81+
}
82+
}
83+
3984
const MAXBLOCKS: usize = 20;
4085

4186
#[derive(Debug, Clone, Copy)]
@@ -425,39 +470,25 @@ impl Compiler {
425470
}
426471
}
427472

428-
/// Check if the slice is a two-element slice (no step)
429-
// = is_two_element_slice
430-
const fn is_two_element_slice(slice: &ast::Expr) -> bool {
431-
matches!(slice, ast::Expr::Slice(s) if s.step.is_none())
432-
}
433-
434-
/// Compile a slice expression
435-
// = compiler_slice
436-
fn compile_slice(&mut self, s: &ast::ExprSlice) -> CompileResult<BuildSliceArgCount> {
437-
// Compile lower
473+
/// Compile just start and stop of a slice (for BINARY_SLICE/STORE_SLICE)
474+
// = codegen_slice_two_parts
475+
fn compile_slice_two_parts(&mut self, s: &ast::ExprSlice) -> CompileResult<()> {
476+
// Compile lower (or None)
438477
if let Some(lower) = &s.lower {
439478
self.compile_expression(lower)?;
440479
} else {
441480
self.emit_load_const(ConstantData::None);
442481
}
443482

444-
// Compile upper
483+
// Compile upper (or None)
445484
if let Some(upper) = &s.upper {
446485
self.compile_expression(upper)?;
447486
} else {
448487
self.emit_load_const(ConstantData::None);
449488
}
450489

451-
Ok(match &s.step {
452-
Some(step) => {
453-
// Compile step if present
454-
self.compile_expression(step)?;
455-
BuildSliceArgCount::Three
456-
}
457-
None => BuildSliceArgCount::Two,
458-
})
490+
Ok(())
459491
}
460-
461492
/// Compile a subscript expression
462493
// = compiler_subscript
463494
fn compile_subscript(
@@ -480,27 +511,20 @@ impl Compiler {
480511
// VISIT(c, expr, e->v.Subscript.value)
481512
self.compile_expression(value)?;
482513

483-
// Handle two-element slice (for Load/Store, not Del)
484-
if Self::is_two_element_slice(slice) && !matches!(ctx, ast::ExprContext::Del) {
485-
let argc = match slice {
486-
ast::Expr::Slice(s) => self.compile_slice(s)?,
514+
// Handle two-element non-constant slice with BINARY_SLICE/STORE_SLICE
515+
if slice.should_use_slice_optimization() && !matches!(ctx, ast::ExprContext::Del) {
516+
match slice {
517+
ast::Expr::Slice(s) => self.compile_slice_two_parts(s)?,
487518
_ => unreachable!(
488-
"is_two_element_slice should only return true for ast::Expr::Slice"
519+
"should_use_slice_optimization should only return true for ast::Expr::Slice"
489520
),
490521
};
491522
match ctx {
492523
ast::ExprContext::Load => {
493-
emit!(self, Instruction::BuildSlice { argc });
494-
emit!(
495-
self,
496-
Instruction::BinaryOp {
497-
op: BinaryOperator::Subscr
498-
}
499-
);
524+
emit!(self, Instruction::BinarySlice);
500525
}
501526
ast::ExprContext::Store => {
502-
emit!(self, Instruction::BuildSlice { argc });
503-
emit!(self, Instruction::StoreSubscr);
527+
emit!(self, Instruction::StoreSlice);
504528
}
505529
_ => unreachable!(),
506530
}
@@ -1269,8 +1293,12 @@ impl Compiler {
12691293
emit!(self, Instruction::PopJumpIfFalse { target: body_block });
12701294

12711295
// Raise NotImplementedError
1272-
let not_implemented_error = self.name("NotImplementedError");
1273-
emit!(self, Instruction::LoadGlobal(not_implemented_error));
1296+
emit!(
1297+
self,
1298+
Instruction::LoadCommonConstant {
1299+
idx: bytecode::CommonConstant::NotImplementedError
1300+
}
1301+
);
12741302
emit!(
12751303
self,
12761304
Instruction::RaiseVarargs {
@@ -2345,8 +2373,12 @@ impl Compiler {
23452373
let after_block = self.new_block();
23462374
self.compile_jump_if(test, true, after_block)?;
23472375

2348-
let assertion_error = self.name("AssertionError");
2349-
emit!(self, Instruction::LoadGlobal(assertion_error));
2376+
emit!(
2377+
self,
2378+
Instruction::LoadCommonConstant {
2379+
idx: bytecode::CommonConstant::AssertionError
2380+
}
2381+
);
23502382
emit!(self, Instruction::PushNull);
23512383
match msg {
23522384
Some(e) => {
@@ -3326,11 +3358,9 @@ impl Compiler {
33263358
// ADDOP_I(c, loc, COPY, 1);
33273359
// ADDOP_JUMP(c, loc, POP_JUMP_IF_NONE, no_match);
33283360
emit!(self, Instruction::Copy { index: 1 });
3329-
self.emit_load_const(ConstantData::None);
3330-
emit!(self, Instruction::IsOp(bytecode::Invert::No)); // is None?
33313361
emit!(
33323362
self,
3333-
Instruction::PopJumpIfTrue {
3363+
Instruction::PopJumpIfNone {
33343364
target: no_match_block
33353365
}
33363366
);
@@ -3455,11 +3485,9 @@ impl Compiler {
34553485
// Stack: [prev_exc, result, result]
34563486

34573487
// POP_JUMP_IF_NOT_NONE reraise
3458-
self.emit_load_const(ConstantData::None);
3459-
emit!(self, Instruction::IsOp(bytecode::Invert::Yes)); // is not None?
34603488
emit!(
34613489
self,
3462-
Instruction::PopJumpIfTrue {
3490+
Instruction::PopJumpIfNotNone {
34633491
target: reraise_block
34643492
}
34653493
);

‎crates/codegen/src/ir.rs‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,7 @@ impl CodeInfo {
161161
) -> crate::InternalResult<CodeObject> {
162162
if opts.optimize > 0 {
163163
self.dce();
164+
self.peephole_optimize();
164165
}
165166

166167
let max_stackdepth = self.max_stackdepth()?;
@@ -403,6 +404,70 @@ impl CodeInfo {
403404
}
404405
}
405406

407+
/// Peephole optimization: combine consecutive instructions into super-instructions
408+
fn peephole_optimize(&mut self) {
409+
for block in &mut self.blocks {
410+
let mut i = 0;
411+
while i + 1 < block.instructions.len() {
412+
let combined = {
413+
let curr = &block.instructions[i];
414+
let next = &block.instructions[i + 1];
415+
416+
// Only combine if both are real instructions (not pseudo)
417+
let (Some(curr_instr), Some(next_instr)) =
418+
(curr.instr.real(), next.instr.real())
419+
else {
420+
i += 1;
421+
continue;
422+
};
423+
424+
match (curr_instr, next_instr) {
425+
// LoadFast + LoadFast -> LoadFastLoadFast (if both indices < 16)
426+
(Instruction::LoadFast(_), Instruction::LoadFast(_)) => {
427+
let idx1 = curr.arg.0;
428+
let idx2 = next.arg.0;
429+
if idx1 < 16 && idx2 < 16 {
430+
let packed = (idx1 << 4) | idx2;
431+
Some((
432+
Instruction::LoadFastLoadFast { arg: Arg::marker() },
433+
OpArg(packed),
434+
))
435+
} else {
436+
None
437+
}
438+
}
439+
// StoreFast + StoreFast -> StoreFastStoreFast (if both indices < 16)
440+
(Instruction::StoreFast(_), Instruction::StoreFast(_)) => {
441+
let idx1 = curr.arg.0;
442+
let idx2 = next.arg.0;
443+
if idx1 < 16 && idx2 < 16 {
444+
let packed = (idx1 << 4) | idx2;
445+
Some((
446+
Instruction::StoreFastStoreFast { arg: Arg::marker() },
447+
OpArg(packed),
448+
))
449+
} else {
450+
None
451+
}
452+
}
453+
_ => None,
454+
}
455+
};
456+
457+
if let Some((new_instr, new_arg)) = combined {
458+
// Combine: keep first instruction's location, replace with combined instruction
459+
block.instructions[i].instr = new_instr.into();
460+
block.instructions[i].arg = new_arg;
461+
// Remove the second instruction
462+
block.instructions.remove(i + 1);
463+
// Don't increment i - check if we can combine again with the next instruction
464+
} else {
465+
i += 1;
466+
}
467+
}
468+
}
469+
}
470+
406471
fn max_stackdepth(&self) -> crate::InternalResult<u32> {
407472
let mut maxdepth = 0u32;
408473
let mut stack = Vec::with_capacity(self.blocks.len());

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ pub use crate::bytecode::{
2121
encode_load_super_attr_arg,
2222
},
2323
oparg::{
24-
BinaryOperator, BuildSliceArgCount, ComparisonOperator, ConvertValueOparg,
24+
BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator, ConvertValueOparg,
2525
IntrinsicFunction1, IntrinsicFunction2, Invert, Label, MakeFunctionFlags, NameIdx, OpArg,
2626
OpArgByte, OpArgState, OpArgType, RaiseKind, ResumeType, SpecialMethod, UnpackExArgs,
2727
},

0 commit comments

Comments
 (0)