Skip to content

Commit 1da0e65

Browse files
youknowoneyouknowone
committed
cold block reordering and jump normalization
Add mark_cold, push_cold_blocks_to_end, and normalize_jumps passes to the codegen CFG pipeline. Use JumpNoInterrupt for exception handler exit paths in try-except-finally compilation.
1 parent 0244657 commit 1da0e65

10 files changed

+408
-142
lines changed

‎crates/codegen/src/compile.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3088,7 +3088,7 @@ impl Compiler {
30883088
let handler_normal_exit = self.new_block();
30893089
emit!(
30903090
self,
3091-
PseudoInstruction::Jump {
3091+
PseudoInstruction::JumpNoInterrupt {
30923092
target: handler_normal_exit,
30933093
}
30943094
);
@@ -3141,10 +3141,10 @@ impl Compiler {
31413141
emit!(self, PseudoInstruction::PopBlock);
31423142
}
31433143

3144-
// Jump to finally block
3144+
// Jump to finally block (JUMP_NO_INTERRUPT: exception cleanup path)
31453145
emit!(
31463146
self,
3147-
PseudoInstruction::Jump {
3147+
PseudoInstruction::JumpNoInterrupt {
31483148
target: finally_block,
31493149
}
31503150
);

‎crates/codegen/src/ir.rs‎

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use alloc::collections::VecDeque;
12
use core::ops;
23

34
use crate::{IndexMap, IndexSet, error::InternalError};
@@ -132,6 +133,9 @@ pub struct Block {
132133
pub preserve_lasti: bool,
133134
/// Stack depth at block entry, set by stack depth analysis
134135
pub start_depth: Option<u32>,
136+
/// Whether this block is "cold" (only reachable via exception table).
137+
/// Cold blocks are pushed to the end during optimization.
138+
pub cold: bool,
135139
}
136140

137141
impl Default for Block {
@@ -142,6 +146,7 @@ impl Default for Block {
142146
except_handler: false,
143147
preserve_lasti: false,
144148
start_depth: None,
149+
cold: false,
145150
}
146151
}
147152
}
@@ -205,6 +210,8 @@ impl CodeInfo {
205210
// Post-codegen CFG analysis passes (flowgraph.c pipeline)
206211
mark_except_handlers(&mut self.blocks);
207212
label_exception_targets(&mut self.blocks);
213+
push_cold_blocks_to_end(&mut self.blocks);
214+
normalize_jumps(&mut self.blocks);
208215

209216
let max_stackdepth = self.max_stackdepth()?;
210217
let cell2arg = self.cell2arg();
@@ -1154,6 +1161,248 @@ pub(crate) fn mark_except_handlers(blocks: &mut [Block]) {
11541161
}
11551162
}
11561163

1164+
/// Mark cold blocks: blocks only reachable via exception table.
1165+
/// BFS from entry following normal control flow (next + jump targets),
1166+
/// skipping edges into except_handler blocks. Unvisited blocks are cold.
1167+
/// flowgraph.c mark_cold
1168+
fn mark_cold(blocks: &mut [Block]) {
1169+
let n = blocks.len();
1170+
let mut warm = vec![false; n];
1171+
let mut queue = VecDeque::new();
1172+
1173+
// Entry block is always warm
1174+
warm[0] = true;
1175+
queue.push_back(BlockIdx(0));
1176+
1177+
while let Some(block_idx) = queue.pop_front() {
1178+
let block = &blocks[block_idx.idx()];
1179+
1180+
// Follow fall-through (block.next)
1181+
let has_fallthrough = block
1182+
.instructions
1183+
.last()
1184+
.map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
1185+
.unwrap_or(true);
1186+
if has_fallthrough && block.next != BlockIdx::NULL {
1187+
let next_idx = block.next.idx();
1188+
if !blocks[next_idx].except_handler && !warm[next_idx] {
1189+
warm[next_idx] = true;
1190+
queue.push_back(block.next);
1191+
}
1192+
}
1193+
1194+
// Follow jump targets in instructions
1195+
for instr in &block.instructions {
1196+
if instr.target != BlockIdx::NULL {
1197+
let target_idx = instr.target.idx();
1198+
if !blocks[target_idx].except_handler && !warm[target_idx] {
1199+
warm[target_idx] = true;
1200+
queue.push_back(instr.target);
1201+
}
1202+
}
1203+
}
1204+
}
1205+
1206+
// Mark non-warm blocks as cold
1207+
for (i, block) in blocks.iter_mut().enumerate() {
1208+
block.cold = !warm[i];
1209+
}
1210+
}
1211+
1212+
/// Reorder the block linked list to push cold blocks to the end.
1213+
/// If a cold block falls through to a warm block, insert an explicit
1214+
/// JUMP_NO_INTERRUPT to maintain control flow.
1215+
/// flowgraph.c push_cold_blocks_to_end
1216+
fn push_cold_blocks_to_end(blocks: &mut Vec<Block>) {
1217+
// Single block, nothing to reorder
1218+
if blocks.len() <= 1 {
1219+
return;
1220+
}
1221+
1222+
mark_cold(blocks);
1223+
1224+
// If a cold block falls through to a warm block, add an explicit jump
1225+
let fixups: Vec<(BlockIdx, BlockIdx)> = iter_blocks(blocks)
1226+
.filter(|(_, block)| {
1227+
block.cold
1228+
&& block.next != BlockIdx::NULL
1229+
&& !blocks[block.next.idx()].cold
1230+
&& block
1231+
.instructions
1232+
.last()
1233+
.map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
1234+
.unwrap_or(true)
1235+
})
1236+
.map(|(idx, block)| (idx, block.next))
1237+
.collect();
1238+
1239+
for (cold_idx, warm_next) in fixups {
1240+
// Create a new block with an explicit jump
1241+
let jump_block_idx = BlockIdx(blocks.len() as u32);
1242+
let loc = blocks[cold_idx.idx()]
1243+
.instructions
1244+
.last()
1245+
.map(|i| i.location)
1246+
.unwrap_or_default();
1247+
let end_loc = blocks[cold_idx.idx()]
1248+
.instructions
1249+
.last()
1250+
.map(|i| i.end_location)
1251+
.unwrap_or_default();
1252+
let mut jump_block = Block {
1253+
cold: true,
1254+
..Block::default()
1255+
};
1256+
jump_block.instructions.push(InstructionInfo {
1257+
instr: PseudoInstruction::JumpNoInterrupt {
1258+
target: Arg::marker(),
1259+
}
1260+
.into(),
1261+
arg: OpArg::new(0),
1262+
target: warm_next,
1263+
location: loc,
1264+
end_location: end_loc,
1265+
except_handler: None,
1266+
lineno_override: None,
1267+
});
1268+
jump_block.next = blocks[cold_idx.idx()].next;
1269+
blocks[cold_idx.idx()].next = jump_block_idx;
1270+
blocks.push(jump_block);
1271+
}
1272+
1273+
// Now reorder: extract cold block streaks and append at the end
1274+
let mut cold_head: BlockIdx = BlockIdx::NULL;
1275+
let mut cold_tail: BlockIdx = BlockIdx::NULL;
1276+
1277+
// Walk the chain, collect and remove cold blocks
1278+
let mut current = BlockIdx(0);
1279+
// Entry block should never be cold
1280+
assert!(!blocks[0].cold);
1281+
1282+
while current != BlockIdx::NULL {
1283+
let next = blocks[current.idx()].next;
1284+
if next == BlockIdx::NULL {
1285+
break;
1286+
}
1287+
1288+
if blocks[next.idx()].cold {
1289+
// Start of a cold streak
1290+
let cold_start = next;
1291+
let mut cold_end = next;
1292+
while blocks[cold_end.idx()].next != BlockIdx::NULL
1293+
&& blocks[blocks[cold_end.idx()].next.idx()].cold
1294+
{
1295+
cold_end = blocks[cold_end.idx()].next;
1296+
}
1297+
1298+
// Unlink cold streak from main chain
1299+
let after_cold = blocks[cold_end.idx()].next;
1300+
blocks[current.idx()].next = after_cold;
1301+
blocks[cold_end.idx()].next = BlockIdx::NULL;
1302+
1303+
// Append to cold list
1304+
if cold_head == BlockIdx::NULL {
1305+
cold_head = cold_start;
1306+
} else {
1307+
blocks[cold_tail.idx()].next = cold_start;
1308+
}
1309+
cold_tail = cold_end;
1310+
1311+
// Don't advance current - check the new next
1312+
} else {
1313+
current = next;
1314+
}
1315+
}
1316+
1317+
// Append cold blocks at the end of main chain
1318+
if cold_head != BlockIdx::NULL {
1319+
// Find end of main chain
1320+
let mut last = current;
1321+
while blocks[last.idx()].next != BlockIdx::NULL {
1322+
last = blocks[last.idx()].next;
1323+
}
1324+
blocks[last.idx()].next = cold_head;
1325+
}
1326+
}
1327+
1328+
/// Returns true if the instruction is a conditional jump (POP_JUMP_IF_*).
1329+
fn is_conditional_jump(instr: &AnyInstruction) -> bool {
1330+
matches!(
1331+
instr.real(),
1332+
Some(
1333+
Instruction::PopJumpIfFalse { .. }
1334+
| Instruction::PopJumpIfTrue { .. }
1335+
| Instruction::PopJumpIfNone { .. }
1336+
| Instruction::PopJumpIfNotNone { .. }
1337+
)
1338+
)
1339+
}
1340+
1341+
/// Remove redundant unconditional jumps and add NOT_TAKEN after forward
1342+
/// conditional jumps (scanning all instructions in each block).
1343+
/// flowgraph.c normalize_jumps + remove_redundant_jumps
1344+
fn normalize_jumps(blocks: &mut [Block]) {
1345+
// Walk linked list to determine visit order (for forward/backward detection)
1346+
let mut visit_order = Vec::new();
1347+
let mut visited = vec![false; blocks.len()];
1348+
let mut current = BlockIdx(0);
1349+
while current != BlockIdx::NULL {
1350+
visit_order.push(current);
1351+
visited[current.idx()] = true;
1352+
current = blocks[current.idx()].next;
1353+
}
1354+
1355+
// Reset visited for forward/backward detection during second pass
1356+
visited.fill(false);
1357+
1358+
for block_idx in visit_order {
1359+
let idx = block_idx.idx();
1360+
visited[idx] = true;
1361+
1362+
// Remove redundant unconditional jump to next block
1363+
let next = blocks[idx].next;
1364+
if next != BlockIdx::NULL {
1365+
let last = blocks[idx].instructions.last();
1366+
let is_jump_to_next = last.is_some_and(|ins| {
1367+
ins.instr.is_unconditional_jump()
1368+
&& ins.target != BlockIdx::NULL
1369+
&& ins.target == next
1370+
});
1371+
if is_jump_to_next && let Some(last_instr) = blocks[idx].instructions.last_mut() {
1372+
last_instr.instr = Instruction::Nop.into();
1373+
last_instr.target = BlockIdx::NULL;
1374+
}
1375+
}
1376+
1377+
// Collect positions where NOT_TAKEN should be inserted
1378+
let mut insert_positions: Vec<(usize, InstructionInfo)> = Vec::new();
1379+
for (i, ins) in blocks[idx].instructions.iter().enumerate() {
1380+
if is_conditional_jump(&ins.instr)
1381+
&& ins.target != BlockIdx::NULL
1382+
&& !visited[ins.target.idx()]
1383+
{
1384+
insert_positions.push((
1385+
i + 1,
1386+
InstructionInfo {
1387+
instr: Instruction::NotTaken.into(),
1388+
arg: OpArg::new(0),
1389+
target: BlockIdx::NULL,
1390+
location: ins.location,
1391+
end_location: ins.end_location,
1392+
except_handler: ins.except_handler,
1393+
lineno_override: None,
1394+
},
1395+
));
1396+
}
1397+
}
1398+
1399+
// Insert NOT_TAKEN in reverse order to preserve indices
1400+
for (pos, info) in insert_positions.into_iter().rev() {
1401+
blocks[idx].instructions.insert(pos, info);
1402+
}
1403+
}
1404+
}
1405+
11571406
/// Label exception targets: walk CFG with except stack, set per-instruction
11581407
/// handler info and block preserve_lasti flag. Converts POP_BLOCK to NOP.
11591408
/// flowgraph.c label_exception_targets + push_except_block

‎crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap‎

Lines changed: 10 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap‎

Lines changed: 13 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap‎

Lines changed: 10 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)