Skip to content

Commit c6eb979

Browse files
youknowoneyouknowone
committed
more optimization
1 parent f8d4d99 commit c6eb979

File tree

7 files changed

+304
-63
lines changed

7 files changed

+304
-63
lines changed

‎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
},

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

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@ use crate::{
44
bytecode::{
55
BorrowedConstant, Constant, InstrDisplayContext,
66
oparg::{
7-
BinaryOperator, BuildSliceArgCount, ComparisonOperator, ConvertValueOparg,
8-
IntrinsicFunction1, IntrinsicFunction2, Invert, Label, MakeFunctionFlags, NameIdx,
9-
OpArg, OpArgByte, OpArgType, RaiseKind, SpecialMethod, UnpackExArgs,
7+
BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator,
8+
ConvertValueOparg, IntrinsicFunction1, IntrinsicFunction2, Invert, Label,
9+
MakeFunctionFlags, NameIdx, OpArg, OpArgByte, OpArgType, RaiseKind, SpecialMethod,
10+
UnpackExArgs,
1011
},
1112
},
1213
marshal::MarshalError,
@@ -22,8 +23,8 @@ use crate::{
2223
#[repr(u8)]
2324
pub enum Instruction {
2425
// No-argument instructions (opcode < HAVE_ARGUMENT=44)
25-
Cache = 0, // Placeholder
26-
BinarySlice = 1, // Placeholder
26+
Cache = 0, // Placeholder
27+
BinarySlice = 1,
2728
BuildTemplate = 2,
2829
BinaryOpInplaceAddUnicode = 3, // Placeholder
2930
CallFunctionEx = 4,
@@ -59,7 +60,7 @@ pub enum Instruction {
5960
ReturnGenerator = 34,
6061
ReturnValue = 35,
6162
SetupAnnotations = 36,
62-
StoreSlice = 37, // Placeholder
63+
StoreSlice = 37,
6364
StoreSubscr = 38,
6465
ToBool = 39,
6566
UnaryInvert = 40,
@@ -122,7 +123,7 @@ pub enum Instruction {
122123
} = 59,
123124
CopyFreeVars {
124125
count: Arg<u32>,
125-
} = 60, // Placeholder
126+
} = 60,
126127
DeleteAttr {
127128
idx: Arg<NameIdx>,
128129
} = 61,
@@ -168,8 +169,8 @@ pub enum Instruction {
168169
idx: Arg<NameIdx>,
169170
} = 80,
170171
LoadCommonConstant {
171-
idx: Arg<u32>,
172-
} = 81, // Placeholder
172+
idx: Arg<CommonConstant>,
173+
} = 81,
173174
LoadConst {
174175
idx: Arg<u32>,
175176
} = 82,
@@ -180,10 +181,10 @@ pub enum Instruction {
180181
LoadFastBorrowLoadFastBorrow {
181182
arg: Arg<u32>,
182183
} = 87, // Placeholder
183-
LoadFastCheck(Arg<NameIdx>) = 88, // Placeholder
184+
LoadFastCheck(Arg<NameIdx>) = 88,
184185
LoadFastLoadFast {
185186
arg: Arg<u32>,
186-
} = 89, // Placeholder
187+
} = 89,
187188
LoadFromDictOrDeref(Arg<NameIdx>) = 90,
188189
LoadFromDictOrGlobals(Arg<NameIdx>) = 91, // Placeholder
189190
LoadGlobal(Arg<NameIdx>) = 92,
@@ -197,7 +198,7 @@ pub enum Instruction {
197198
LoadSuperAttr {
198199
arg: Arg<u32>,
199200
} = 96,
200-
MakeCell(Arg<NameIdx>) = 97, // Placeholder
201+
MakeCell(Arg<NameIdx>) = 97,
201202
MapAdd {
202203
i: Arg<u32>,
203204
} = 98,
@@ -207,10 +208,10 @@ pub enum Instruction {
207208
} = 100,
208209
PopJumpIfNone {
209210
target: Arg<Label>,
210-
} = 101, // Placeholder
211+
} = 101,
211212
PopJumpIfNotNone {
212213
target: Arg<Label>,
213-
} = 102, // Placeholder
214+
} = 102,
214215
PopJumpIfTrue {
215216
target: Arg<Label>,
216217
} = 103,
@@ -243,7 +244,7 @@ pub enum Instruction {
243244
} = 113,
244245
StoreFastStoreFast {
245246
arg: Arg<u32>,
246-
} = 114, // Placeholder
247+
} = 114,
247248
StoreGlobal(Arg<NameIdx>) = 115,
248249
StoreName(Arg<NameIdx>) = 116,
249250
Swap {
@@ -616,7 +617,7 @@ impl InstructionMetadata for Instruction {
616617
Self::JumpBackwardNoInterrupt { .. } => 0,
617618
Self::JumpBackward { .. } => 0,
618619
Self::JumpForward { .. } => 0,
619-
Self::LoadFastCheck(_) => 0,
620+
Self::LoadFastCheck(_) => 1,
620621
Self::LoadFastLoadFast { .. } => 2,
621622
Self::LoadFastBorrowLoadFastBorrow { .. } => 2,
622623
Self::LoadFromDictOrGlobals(_) => 0,

0 commit comments

Comments
 (0)