Add duplicate_exits_without_lineno (disabled) and Block: Clone

youknowone · youknowone · commit 42b24cbe4005 · 2026-03-30T11:43:03.000+09:00
Prepare infrastructure for exit block duplication optimization.
Currently disabled pending stackdepth integration.
diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs
@@ -7575,7 +7575,8 @@ impl Compiler {
                 lower, upper, step, ..
             }) => {
                 // Try constant slice folding first
-                if self.try_fold_constant_slice(lower.as_deref(), upper.as_deref(), step.as_deref()) {
+                if self.try_fold_constant_slice(lower.as_deref(), upper.as_deref(), step.as_deref())
+                {
                     return Ok(());
                 }
                 let mut compile_bound = |bound: Option<&ast::Expr>| match bound {
@@ -8998,17 +8999,16 @@ impl Compiler {
         fn to_const(expr: Option<&ast::Expr>) -> Option<ConstantData> {
             match expr {
                 None | Some(ast::Expr::NoneLiteral(_)) => Some(ConstantData::None),
-                Some(ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. })) => {
-                    match value {
-                        ast::Number::Int(i) => Some(ConstantData::Integer {
-                            value: i.as_i64()?.into(),
-                        }),
-                        ast::Number::Float(f) => Some(ConstantData::Float { value: *f }),
-                        ast::Number::Complex { real, imag } => Some(ConstantData::Complex {
-                            value: num_complex::Complex64::new(*real, *imag),
-                        }),
-                    }
-                }
+                Some(ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. })) => match value
+                {
+                    ast::Number::Int(i) => Some(ConstantData::Integer {
+                        value: i.as_i64()?.into(),
+                    }),
+                    ast::Number::Float(f) => Some(ConstantData::Float { value: *f }),
+                    ast::Number::Complex { real, imag } => Some(ConstantData::Complex {
+                        value: num_complex::Complex64::new(*real, *imag),
+                    }),
+                },
                 Some(ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. })) => {
                     Some(ConstantData::Boolean { value: *value })
                 }
diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs
@@ -125,7 +125,7 @@ pub struct ExceptHandlerInfo {
 // spell-checker:ignore petgraph
 // TODO: look into using petgraph for handling blocks and stuff? it's heavier than this, but it
 // might enable more analysis/optimizations
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Block {
     pub instructions: Vec<InstructionInfo>,
     pub next: BlockIdx,
@@ -225,6 +225,8 @@ impl CodeInfo {
         normalize_jumps(&mut self.blocks);
         self.dce(); // re-run within-block DCE after normalize_jumps creates new instructions
         self.eliminate_unreachable_blocks();
+        // TODO: duplicate_exits_without_lineno disabled pending stackdepth fix
+        // duplicate_exits_without_lineno(&mut self.blocks);
         duplicate_end_returns(&mut self.blocks);
         self.dce(); // truncate after terminal in blocks that got return duplicated
         self.eliminate_unreachable_blocks(); // remove now-unreachable last block
@@ -2239,15 +2241,14 @@ fn jump_threading(blocks: &mut [Block]) {
             }
             // Check if target block's first instruction is an unconditional jump
             let target_block = &blocks[target.idx()];
-            if let Some(target_ins) = target_block.instructions.first() {
-                if target_ins.instr.is_unconditional_jump()
-                    && target_ins.target != BlockIdx::NULL
-                    && target_ins.target != target
-                {
-                    let final_target = target_ins.target;
-                    blocks[bi].instructions[last_idx].target = final_target;
-                    changed = true;
-                }
+            if let Some(target_ins) = target_block.instructions.first()
+                && target_ins.instr.is_unconditional_jump()
+                && target_ins.target != BlockIdx::NULL
+                && target_ins.target != target
+            {
+                let final_target = target_ins.target;
+                blocks[bi].instructions[last_idx].target = final_target;
+                changed = true;
             }
         }
     }
@@ -2301,8 +2302,7 @@ fn normalize_jumps(blocks: &mut Vec<Block>) {
 
     visited.fill(false);
 
-    for vi in 0..visit_order.len() {
-        let block_idx = visit_order[vi];
+    for &block_idx in &visit_order {
         let idx = block_idx.idx();
         visited[idx] = true;
 
@@ -2323,76 +2323,80 @@ fn normalize_jumps(blocks: &mut Vec<Block>) {
 
         // Normalize conditional jumps: forward gets NOT_TAKEN, backward gets inverted
         let last = blocks[idx].instructions.last();
-        if let Some(last_ins) = last {
-            if is_conditional_jump(&last_ins.instr) && last_ins.target != BlockIdx::NULL {
-                let target = last_ins.target;
-                let is_forward = !visited[target.idx()];
-
-                if is_forward {
-                    // Insert NOT_TAKEN after forward conditional jump
-                    let not_taken = InstructionInfo {
+        if let Some(last_ins) = last
+            && is_conditional_jump(&last_ins.instr)
+            && last_ins.target != BlockIdx::NULL
+        {
+            let target = last_ins.target;
+            let is_forward = !visited[target.idx()];
+
+            if is_forward {
+                // Insert NOT_TAKEN after forward conditional jump
+                let not_taken = InstructionInfo {
+                    instr: Instruction::NotTaken.into(),
+                    arg: OpArg::new(0),
+                    target: BlockIdx::NULL,
+                    location: last_ins.location,
+                    end_location: last_ins.end_location,
+                    except_handler: last_ins.except_handler,
+                    lineno_override: None,
+                    cache_entries: 0,
+                };
+                blocks[idx].instructions.push(not_taken);
+            } else {
+                // Backward conditional jump: invert and create new block
+                // Transform: `cond_jump T` (backward)
+                // Into: `reversed_cond_jump b_next` + new block [NOT_TAKEN, JUMP T]
+                let loc = last_ins.location;
+                let end_loc = last_ins.end_location;
+                let exc_handler = last_ins.except_handler;
+
+                if let Some(reversed) = reversed_conditional(&last_ins.instr) {
+                    let old_next = blocks[idx].next;
+                    let is_cold = blocks[idx].cold;
+
+                    // Create new block with NOT_TAKEN + JUMP to original backward target
+                    let new_block_idx = BlockIdx(blocks.len() as u32);
+                    let mut new_block = Block {
+                        cold: is_cold,
+                        ..Block::default()
+                    };
+                    new_block.instructions.push(InstructionInfo {
                         instr: Instruction::NotTaken.into(),
                         arg: OpArg::new(0),
                         target: BlockIdx::NULL,
-                        location: last_ins.location,
-                        end_location: last_ins.end_location,
-                        except_handler: last_ins.except_handler,
+                        location: loc,
+                        end_location: end_loc,
+                        except_handler: exc_handler,
                         lineno_override: None,
                         cache_entries: 0,
-                    };
-                    blocks[idx].instructions.push(not_taken);
-                } else {
-                    // Backward conditional jump: invert and create new block
-                    // Transform: `cond_jump T` (backward)
-                    // Into: `reversed_cond_jump b_next` + new block [NOT_TAKEN, JUMP T]
-                    let loc = last_ins.location;
-                    let end_loc = last_ins.end_location;
-                    let exc_handler = last_ins.except_handler;
-
-                    if let Some(reversed) = reversed_conditional(&last_ins.instr) {
-                        let old_next = blocks[idx].next;
-                        let is_cold = blocks[idx].cold;
-
-                        // Create new block with NOT_TAKEN + JUMP to original backward target
-                        let new_block_idx = BlockIdx(blocks.len() as u32);
-                        let mut new_block = Block {
-                            cold: is_cold,
-                            ..Block::default()
-                        };
-                        new_block.instructions.push(InstructionInfo {
-                            instr: Instruction::NotTaken.into(),
-                            arg: OpArg::new(0),
-                            target: BlockIdx::NULL,
-                            location: loc,
-                            end_location: end_loc,
-                            except_handler: exc_handler,
-                            lineno_override: None,
-                            cache_entries: 0,
-                        });
-                        new_block.instructions.push(InstructionInfo {
-                            instr: PseudoInstruction::Jump { delta: Arg::marker() }.into(),
-                            arg: OpArg::new(0),
-                            target,
-                            location: loc,
-                            end_location: end_loc,
-                            except_handler: exc_handler,
-                            lineno_override: None,
-                            cache_entries: 0,
-                        });
-                        new_block.next = old_next;
-
-                        // Update the conditional jump: invert opcode, target = old next block
-                        let last_mut = blocks[idx].instructions.last_mut().unwrap();
-                        last_mut.instr = reversed;
-                        last_mut.target = old_next;
-
-                        // Splice new block between current and old next
-                        blocks[idx].next = new_block_idx;
-                        blocks.push(new_block);
-
-                        // Extend visited array and update visit order
-                        visited.push(true);
-                    }
+                    });
+                    new_block.instructions.push(InstructionInfo {
+                        instr: PseudoInstruction::Jump {
+                            delta: Arg::marker(),
+                        }
+                        .into(),
+                        arg: OpArg::new(0),
+                        target,
+                        location: loc,
+                        end_location: end_loc,
+                        except_handler: exc_handler,
+                        lineno_override: None,
+                        cache_entries: 0,
+                    });
+                    new_block.next = old_next;
+
+                    // Update the conditional jump: invert opcode, target = old next block
+                    let last_mut = blocks[idx].instructions.last_mut().unwrap();
+                    last_mut.instr = reversed;
+                    last_mut.target = old_next;
+
+                    // Splice new block between current and old next
+                    blocks[idx].next = new_block_idx;
+                    blocks.push(new_block);
+
+                    // Extend visited array and update visit order
+                    visited.push(true);
                 }
             }
         }
@@ -2496,10 +2500,111 @@ fn normalize_jumps(blocks: &mut Vec<Block>) {
     }
 }
 
+#[allow(dead_code)]
+fn is_simple_return_block(block: &Block) -> bool {
+    block.instructions.len() == 1
+        && matches!(
+            block.instructions[0].instr.real(),
+            Some(Instruction::ReturnValue)
+        )
+}
+
+/// Follow chain of empty blocks to find first non-empty block.
+fn next_nonempty_block(blocks: &[Block], mut idx: BlockIdx) -> BlockIdx {
+    while idx != BlockIdx::NULL
+        && blocks[idx.idx()].instructions.is_empty()
+        && blocks[idx.idx()].next != BlockIdx::NULL
+    {
+        idx = blocks[idx.idx()].next;
+    }
+    idx
+}
+
+#[allow(dead_code)]
+fn duplicate_exits_without_lineno(blocks: &mut Vec<Block>) {
+    // Count predecessors for each block
+    let mut predecessors = vec![0u32; blocks.len()];
+    let mut current = BlockIdx(0);
+    while current != BlockIdx::NULL {
+        let block = &blocks[current.idx()];
+        // Fall-through predecessor
+        let next = next_nonempty_block(blocks, block.next);
+        if next != BlockIdx::NULL {
+            let has_fallthrough = block
+                .instructions
+                .last()
+                .is_none_or(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump());
+            if has_fallthrough {
+                predecessors[next.idx()] += 1;
+            }
+        }
+        // Jump target predecessor
+        for ins in &block.instructions {
+            if ins.target != BlockIdx::NULL {
+                let target = next_nonempty_block(blocks, ins.target);
+                if target != BlockIdx::NULL {
+                    predecessors[target.idx()] += 1;
+                }
+            }
+        }
+        current = block.next;
+    }
+
+    // For each block, if its last instruction jumps to an exit block without lineno
+    // that has >1 predecessor, copy that exit block.
+    current = BlockIdx(0);
+    while current != BlockIdx::NULL {
+        let block = &blocks[current.idx()];
+        let last = match block.instructions.last() {
+            Some(ins) if ins.target != BlockIdx::NULL => ins,
+            _ => {
+                current = blocks[current.idx()].next;
+                continue;
+            }
+        };
+        if !last.instr.is_unconditional_jump() && !is_conditional_jump(&last.instr) {
+            current = blocks[current.idx()].next;
+            continue;
+        }
+
+        let target = next_nonempty_block(blocks, last.target);
+        if target == BlockIdx::NULL || !is_simple_return_block(&blocks[target.idx()]) {
+            current = blocks[current.idx()].next;
+            continue;
+        }
+        if predecessors[target.idx()] <= 1 {
+            current = blocks[current.idx()].next;
+            continue;
+        }
+
+        // Copy the exit block
+        let new_idx = BlockIdx(blocks.len() as u32);
+        let mut new_block = blocks[target.idx()].clone();
+        // Set location from the jump instruction
+        let jump_loc = blocks[current.idx()].instructions.last().unwrap().location;
+        let jump_end_loc = blocks[current.idx()]
+            .instructions
+            .last()
+            .unwrap()
+            .end_location;
+        if let Some(first) = new_block.instructions.first_mut() {
+            first.location = jump_loc;
+            first.end_location = jump_end_loc;
+        }
+        new_block.next = blocks[target.idx()].next;
+        blocks.push(new_block);
+
+        // Update the jump target
+        let last_mut = blocks[current.idx()].instructions.last_mut().unwrap();
+        last_mut.target = new_idx;
+        predecessors[target.idx()] -= 1;
+
+        current = blocks[current.idx()].next;
+    }
+}
+
 /// Duplicate `LOAD_CONST None + RETURN_VALUE` for blocks that fall through
-/// to the final return block. Matches CPython's behavior of ensuring every
-/// code path that reaches the end of a function/module has its own explicit
-/// return instruction.
+/// to the final return block.
 fn duplicate_end_returns(blocks: &mut [Block]) {
     // Walk the block chain to find the last block
     let mut last_block = BlockIdx(0);
diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs
@@ -851,17 +851,37 @@ impl CodeUnits {
 /// ```
 #[derive(Debug, Clone)]
 pub enum ConstantData {
-    Tuple { elements: Vec<ConstantData> },
-    Integer { value: BigInt },
-    Float { value: f64 },
-    Complex { value: Complex64 },
-    Boolean { value: bool },
-    Str { value: Wtf8Buf },
-    Bytes { value: Vec<u8> },
-    Code { code: Box<CodeObject> },
+    Tuple {
+        elements: Vec<ConstantData>,
+    },
+    Integer {
+        value: BigInt,
+    },
+    Float {
+        value: f64,
+    },
+    Complex {
+        value: Complex64,
+    },
+    Boolean {
+        value: bool,
+    },
+    Str {
+        value: Wtf8Buf,
+    },
+    Bytes {
+        value: Vec<u8>,
+    },
+    Code {
+        code: Box<CodeObject>,
+    },
     /// Constant slice(start, stop, step)
-    Slice { elements: Box<[ConstantData; 3]> },
-    Frozenset { elements: Vec<ConstantData> },
+    Slice {
+        elements: Box<[ConstantData; 3]>,
+    },
+    Frozenset {
+        elements: Vec<ConstantData>,
+    },
     None,
     Ellipsis,
 }
diff --git a/crates/compiler-core/src/marshal.rs b/crates/compiler-core/src/marshal.rs
