Skip to content

Commit f1beb4e

Browse files
youknowoneyouknowone
committed
Align bytecode codegen with CPython structure
1 parent 6daeae2 commit f1beb4e

2 files changed

Lines changed: 157 additions & 19 deletions

File tree

‎crates/codegen/src/compile.rs‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -460,6 +460,18 @@ impl Compiler {
460460
}
461461
}
462462

463+
fn constant_expr_truthiness(&mut self, expr: &ast::Expr) -> CompileResult<Option<bool>> {
464+
Ok(self
465+
.try_fold_constant_expr(expr)?
466+
.map(|constant| Self::constant_truthiness(&constant)))
467+
}
468+
469+
fn disable_load_fast_borrow_for_block(&mut self, block: BlockIdx) {
470+
if block != BlockIdx::NULL {
471+
self.current_code_info().blocks[block.idx()].disable_load_fast_borrow = true;
472+
}
473+
}
474+
463475
fn new(opts: CompileOpts, source_file: SourceFile, code_name: String) -> Self {
464476
let module_code = ir::CodeInfo {
465477
flags: bytecode::CodeFlags::NEWLOCALS,
@@ -5519,10 +5531,14 @@ impl Compiler {
55195531
body: &[ast::Stmt],
55205532
elif_else_clauses: &[ast::ElifElseClause],
55215533
) -> CompileResult<()> {
5534+
let test_truthiness = self.constant_expr_truthiness(test)?;
55225535
match elif_else_clauses {
55235536
// Only if
55245537
[] => {
55255538
let after_block = self.new_block();
5539+
if matches!(test_truthiness, Some(false)) {
5540+
self.disable_load_fast_borrow_for_block(after_block);
5541+
}
55265542
self.compile_jump_if(test, false, after_block)?;
55275543
self.compile_statements(body)?;
55285544
self.switch_to_block(after_block);
@@ -5532,6 +5548,9 @@ impl Compiler {
55325548
let after_block = self.new_block();
55335549
let mut next_block = self.new_block();
55345550

5551+
if matches!(test_truthiness, Some(false)) {
5552+
self.disable_load_fast_borrow_for_block(next_block);
5553+
}
55355554
self.compile_jump_if(test, false, next_block)?;
55365555
self.compile_statements(body)?;
55375556
emit!(self, PseudoInstruction::Jump { delta: after_block });
@@ -5573,6 +5592,9 @@ impl Compiler {
55735592

55745593
self.switch_to_block(while_block);
55755594
self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?;
5595+
if matches!(self.constant_expr_truthiness(test)?, Some(false)) {
5596+
self.disable_load_fast_borrow_for_block(else_block);
5597+
}
55765598
self.compile_jump_if(test, false, else_block)?;
55775599

55785600
let was_in_loop = self.ctx.loop_data.replace((while_block, after_block));
@@ -12397,11 +12419,17 @@ else:
1239712419
",
1239812420
);
1239912421

12400-
assert!(find_code(&code, "fallback").is_some(), "missing fallback code");
12422+
assert!(
12423+
find_code(&code, "fallback").is_some(),
12424+
"missing fallback code"
12425+
);
1240112426
let impl_code = find_code(&code, "impl").expect("missing impl code");
1240212427
assert!(
1240312428
impl_code.instructions.iter().any(|unit| {
12404-
matches!(unit.op, Instruction::LoadGlobal { .. } | Instruction::LoadName { .. })
12429+
matches!(
12430+
unit.op,
12431+
Instruction::LoadGlobal { .. } | Instruction::LoadName { .. }
12432+
)
1240512433
}),
1240612434
"expected impl to compile global name access, got ops={:?}",
1240712435
impl_code

‎crates/codegen/src/ir.rs‎

Lines changed: 127 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,8 @@ pub struct Block {
148148
pub start_depth: Option<u32>,
149149
/// Whether this block is only reachable via exception table (b_cold)
150150
pub cold: bool,
151+
/// Whether LOAD_FAST borrow optimization should be suppressed for this block.
152+
pub disable_load_fast_borrow: bool,
151153
}
152154

153155
impl Default for Block {
@@ -159,6 +161,7 @@ impl Default for Block {
159161
preserve_lasti: false,
160162
start_depth: None,
161163
cold: false,
164+
disable_load_fast_borrow: false,
162165
}
163166
}
164167
}
@@ -283,6 +286,8 @@ impl CodeInfo {
283286
// Match CPython order: pseudo ops are lowered after stackdepth
284287
// calculation but before optimize_load_fast.
285288
convert_pseudo_ops(&mut self.blocks, &cellfixedoffsets);
289+
self.compute_load_fast_start_depths();
290+
self.propagate_disable_load_fast_borrow();
286291
// optimize_load_fast: after normalize_jumps
287292
self.optimize_load_fast_borrow();
288293
self.deoptimize_borrow_after_push_exc_info();
@@ -291,6 +296,7 @@ impl CodeInfo {
291296
self.optimize_load_global_push_null();
292297
self.remove_redundant_const_pop_top_pairs();
293298
self.reorder_entry_prefix_cell_setup();
299+
self.remove_unused_consts();
294300

295301
let Self {
296302
flags,
@@ -1720,7 +1726,6 @@ impl CodeInfo {
17201726
Instruction::LoadSmallInt { i } => Some(i.get(arg) != 0),
17211727
_ => None,
17221728
};
1723-
17241729
for block in &mut self.blocks {
17251730
let mut i = 0;
17261731
while i + 1 < block.instructions.len() {
@@ -2155,7 +2160,10 @@ impl CodeInfo {
21552160
expected == Some(start_depth),
21562161
"optimize_load_fast_borrow start_depth mismatch: source={source:?} target={target:?} expected={expected:?} actual={:?} source_last={:?} target_instrs={:?}",
21572162
Some(start_depth),
2158-
blocks[source.idx()].instructions.last().and_then(|info| info.instr.real()),
2163+
blocks[source.idx()]
2164+
.instructions
2165+
.last()
2166+
.and_then(|info| info.instr.real()),
21592167
blocks[target.idx()]
21602168
.instructions
21612169
.iter()
@@ -2418,6 +2426,9 @@ impl CodeInfo {
24182426
}
24192427

24202428
let block = &mut self.blocks[block_idx];
2429+
if block.disable_load_fast_borrow {
2430+
continue;
2431+
}
24212432
for (i, info) in block.instructions.iter_mut().enumerate() {
24222433
if instr_flags[i] != 0 {
24232434
continue;
@@ -2441,6 +2452,96 @@ impl CodeInfo {
24412452
}
24422453
}
24432454

2455+
fn compute_load_fast_start_depths(&mut self) {
2456+
fn stackdepth_push(
2457+
stack: &mut Vec<BlockIdx>,
2458+
start_depths: &mut [u32],
2459+
target: BlockIdx,
2460+
depth: u32,
2461+
) {
2462+
let idx = target.idx();
2463+
let block_depth = &mut start_depths[idx];
2464+
debug_assert!(
2465+
*block_depth == u32::MAX || *block_depth == depth,
2466+
"Invalid CFG, inconsistent optimize_load_fast stackdepth for block {:?}: existing={}, new={}",
2467+
target,
2468+
*block_depth,
2469+
depth,
2470+
);
2471+
if *block_depth == u32::MAX {
2472+
*block_depth = depth;
2473+
stack.push(target);
2474+
}
2475+
}
2476+
2477+
let mut stack = Vec::with_capacity(self.blocks.len());
2478+
let mut start_depths = vec![u32::MAX; self.blocks.len()];
2479+
stackdepth_push(&mut stack, &mut start_depths, BlockIdx(0), 0);
2480+
2481+
'process_blocks: while let Some(block_idx) = stack.pop() {
2482+
let mut depth = start_depths[block_idx.idx()];
2483+
let block = &self.blocks[block_idx];
2484+
for ins in &block.instructions {
2485+
let instr = &ins.instr;
2486+
let effect = instr.stack_effect(ins.arg.into());
2487+
let new_depth = depth.saturating_add_signed(effect);
2488+
if ins.target != BlockIdx::NULL {
2489+
let jump_effect = instr.stack_effect_jump(ins.arg.into());
2490+
let target_depth = depth.saturating_add_signed(jump_effect);
2491+
stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth);
2492+
}
2493+
depth = new_depth;
2494+
if instr.is_scope_exit() || instr.is_unconditional_jump() {
2495+
continue 'process_blocks;
2496+
}
2497+
}
2498+
if block.next != BlockIdx::NULL {
2499+
stackdepth_push(&mut stack, &mut start_depths, block.next, depth);
2500+
}
2501+
}
2502+
2503+
for (block, &start_depth) in self.blocks.iter_mut().zip(&start_depths) {
2504+
block.start_depth = (start_depth != u32::MAX).then_some(start_depth);
2505+
}
2506+
}
2507+
2508+
fn propagate_disable_load_fast_borrow(&mut self) {
2509+
let mut predecessors = vec![Vec::new(); self.blocks.len()];
2510+
for (pred_idx, block) in iter_blocks(&self.blocks) {
2511+
if block.next != BlockIdx::NULL {
2512+
predecessors[block.next.idx()].push(pred_idx);
2513+
}
2514+
for info in &block.instructions {
2515+
if info.target != BlockIdx::NULL {
2516+
predecessors[info.target.idx()].push(pred_idx);
2517+
}
2518+
}
2519+
}
2520+
2521+
let mut changed = true;
2522+
while changed {
2523+
changed = false;
2524+
let block_indices: Vec<_> = iter_blocks(&self.blocks).map(|(idx, _)| idx).collect();
2525+
for idx in block_indices {
2526+
if idx == BlockIdx(0) || self.blocks[idx.idx()].disable_load_fast_borrow {
2527+
continue;
2528+
}
2529+
let preds = &predecessors[idx.idx()];
2530+
if preds.is_empty() {
2531+
continue;
2532+
}
2533+
if preds
2534+
.iter()
2535+
.copied()
2536+
.all(|pred_idx| self.blocks[pred_idx.idx()].disable_load_fast_borrow)
2537+
{
2538+
self.blocks[idx.idx()].disable_load_fast_borrow = true;
2539+
changed = true;
2540+
}
2541+
}
2542+
}
2543+
}
2544+
24442545
fn deoptimize_borrow_for_handler_return_paths(&mut self) {
24452546
for block in &mut self.blocks {
24462547
let len = block.instructions.len();
@@ -2968,9 +3069,11 @@ impl CodeInfo {
29683069
old_next
29693070
} else {
29703071
let suffix_idx = BlockIdx::new(base + annotations_blocks.len() as u32);
3072+
let disable_load_fast_borrow = self.blocks[block_idx].disable_load_fast_borrow;
29713073
let block = Block {
29723074
instructions: suffix,
29733075
next: old_next,
3076+
disable_load_fast_borrow,
29743077
..Default::default()
29753078
};
29763079
annotations_blocks.push(block);
@@ -3429,11 +3532,13 @@ fn split_blocks_at_jumps(blocks: &mut Vec<Block>) {
34293532
let tail: Vec<InstructionInfo> = blocks[bi].instructions.drain(pos..).collect();
34303533
let old_next = blocks[bi].next;
34313534
let cold = blocks[bi].cold;
3535+
let disable_load_fast_borrow = blocks[bi].disable_load_fast_borrow;
34323536
blocks[bi].next = new_block_idx;
34333537
blocks.push(Block {
34343538
instructions: tail,
34353539
next: old_next,
34363540
cold,
3541+
disable_load_fast_borrow,
34373542
..Block::default()
34383543
});
34393544
// Don't increment bi - re-check current block (it might still have issues)
@@ -3673,11 +3778,13 @@ fn normalize_jumps(blocks: &mut Vec<Block>) {
36733778
if let Some(reversed) = reversed_conditional(&last_ins.instr) {
36743779
let old_next = blocks[idx].next;
36753780
let is_cold = blocks[idx].cold;
3781+
let disable_load_fast_borrow = blocks[idx].disable_load_fast_borrow;
36763782

36773783
// Create new block with NOT_TAKEN + JUMP to original backward target
36783784
let new_block_idx = BlockIdx(blocks.len() as u32);
36793785
let mut new_block = Block {
36803786
cold: is_cold,
3787+
disable_load_fast_borrow,
36813788
..Block::default()
36823789
};
36833790
new_block.instructions.push(InstructionInfo {
@@ -3795,27 +3902,30 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) {
37953902
let is_return_epilogue_block = |block: &Block| {
37963903
matches!(
37973904
block.instructions.as_slice(),
3798-
[InstructionInfo {
3799-
instr: AnyInstruction::Real(Instruction::LoadConst { .. }),
3800-
..
3801-
}, InstructionInfo {
3802-
instr: AnyInstruction::Real(Instruction::ReturnValue),
3803-
..
3804-
}]
3805-
| [InstructionInfo {
3806-
instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }),
3905+
[
3906+
InstructionInfo {
3907+
instr: AnyInstruction::Real(Instruction::LoadConst { .. }),
38073908
..
3808-
}, InstructionInfo {
3909+
},
3910+
InstructionInfo {
38093911
instr: AnyInstruction::Real(Instruction::ReturnValue),
38103912
..
3811-
}]
3812-
| [InstructionInfo {
3913+
}
3914+
] | [
3915+
InstructionInfo {
3916+
instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }),
3917+
..
3918+
},
3919+
InstructionInfo {
38133920
instr: AnyInstruction::Real(Instruction::ReturnValue),
38143921
..
3815-
}]
3922+
}
3923+
] | [InstructionInfo {
3924+
instr: AnyInstruction::Real(Instruction::ReturnValue),
3925+
..
3926+
}]
38163927
)
38173928
};
3818-
38193929
loop {
38203930
let mut changes = false;
38213931
let mut current = BlockIdx(0);
@@ -3848,7 +3958,6 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) {
38483958
current = next;
38493959
continue;
38503960
}
3851-
38523961
if small_exit_block || no_lineno_no_fallthrough {
38533962
if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() {
38543963
set_to_nop(last_instr);
@@ -4712,6 +4821,7 @@ fn duplicate_end_returns(blocks: &mut Vec<Block>) {
47124821
let new_block = Block {
47134822
cold: blocks[last_block.idx()].cold,
47144823
except_handler: blocks[last_block.idx()].except_handler,
4824+
disable_load_fast_borrow: blocks[last_block.idx()].disable_load_fast_borrow,
47154825
instructions: cloned_return,
47164826
next: if is_conditional {
47174827
last_block

0 commit comments

Comments
 (0)