Skip to main content

rustpython_codegen/
ir.rs

1use alloc::collections::VecDeque;
2use core::ops;
3
4use crate::{IndexMap, IndexSet, error::InternalError};
5use malachite_bigint::BigInt;
6use num_traits::{ToPrimitive, Zero};
7
8use rustpython_compiler_core::{
9    OneIndexed, SourceLocation,
10    bytecode::{
11        AnyInstruction, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, CodeFlags,
12        CodeObject, CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, InstrDisplayContext,
13        Instruction, InstructionMetadata, Label, OpArg, PseudoInstruction, PyCodeLocationInfoKind,
14        encode_exception_table, oparg,
15    },
16    varint::{write_signed_varint, write_varint},
17};
18
19/// Location info for linetable generation (allows line 0 for RESUME)
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21struct LineTableLocation {
22    line: i32,
23    end_line: i32,
24    col: i32,
25    end_col: i32,
26}
27
28const MAX_INT_SIZE_BITS: u64 = 128;
29const MIN_CONST_SEQUENCE_SIZE: usize = 3;
30
31/// Metadata for a code unit
32// = _PyCompile_CodeUnitMetadata
33#[derive(Clone, Debug)]
34pub struct CodeUnitMetadata {
35    pub name: String,                        // u_name (obj_name)
36    pub qualname: Option<String>,            // u_qualname
37    pub consts: IndexSet<ConstantData>,      // u_consts
38    pub names: IndexSet<String>,             // u_names
39    pub varnames: IndexSet<String>,          // u_varnames
40    pub cellvars: IndexSet<String>,          // u_cellvars
41    pub freevars: IndexSet<String>,          // u_freevars
42    pub fast_hidden: IndexMap<String, bool>, // u_fast_hidden
43    pub argcount: u32,                       // u_argcount
44    pub posonlyargcount: u32,                // u_posonlyargcount
45    pub kwonlyargcount: u32,                 // u_kwonlyargcount
46    pub firstlineno: OneIndexed,             // u_firstlineno
47}
48// use rustpython_parser_core::source_code::{LineNumber, SourceLocation};
49
50#[derive(Copy, Clone, PartialEq, Eq, Debug)]
51pub struct BlockIdx(u32);
52
53impl BlockIdx {
54    pub const NULL: Self = Self::new(u32::MAX);
55
56    /// Creates a new instance of [`BlockIdx`] from a [`u32`].
57    #[must_use]
58    pub const fn new(value: u32) -> Self {
59        Self(value)
60    }
61
62    /// Returns the inner value as a [`usize`].
63    #[must_use]
64    pub const fn idx(self) -> usize {
65        self.0 as usize
66    }
67}
68
69impl From<BlockIdx> for u32 {
70    fn from(block_idx: BlockIdx) -> Self {
71        block_idx.0
72    }
73}
74
75impl ops::Index<BlockIdx> for [Block] {
76    type Output = Block;
77
78    fn index(&self, idx: BlockIdx) -> &Block {
79        &self[idx.idx()]
80    }
81}
82
83impl ops::IndexMut<BlockIdx> for [Block] {
84    fn index_mut(&mut self, idx: BlockIdx) -> &mut Block {
85        &mut self[idx.idx()]
86    }
87}
88
89impl ops::Index<BlockIdx> for Vec<Block> {
90    type Output = Block;
91
92    fn index(&self, idx: BlockIdx) -> &Block {
93        &self[idx.idx()]
94    }
95}
96
97impl ops::IndexMut<BlockIdx> for Vec<Block> {
98    fn index_mut(&mut self, idx: BlockIdx) -> &mut Block {
99        &mut self[idx.idx()]
100    }
101}
102
103#[derive(Clone, Copy, Debug)]
104pub struct InstructionInfo {
105    pub instr: AnyInstruction,
106    pub arg: OpArg,
107    pub target: BlockIdx,
108    pub location: SourceLocation,
109    pub end_location: SourceLocation,
110    pub except_handler: Option<ExceptHandlerInfo>,
111    /// Override line number for linetable (e.g., line 0 for module RESUME)
112    pub lineno_override: Option<i32>,
113    /// Number of CACHE code units emitted after this instruction
114    pub cache_entries: u32,
115}
116
117/// Exception handler information for an instruction.
118#[derive(Clone, Copy, Debug, PartialEq, Eq)]
119pub struct ExceptHandlerInfo {
120    /// Block to jump to when exception occurs
121    pub handler_block: BlockIdx,
122    /// Stack depth at handler entry
123    pub stack_depth: u32,
124    /// Whether to push lasti before exception
125    pub preserve_lasti: bool,
126}
127
128// spell-checker:ignore petgraph
129// TODO: look into using petgraph for handling blocks and stuff? it's heavier than this, but it
130// might enable more analysis/optimizations
131#[derive(Debug, Clone)]
132pub struct Block {
133    pub instructions: Vec<InstructionInfo>,
134    pub next: BlockIdx,
135    // Post-codegen analysis fields (set by label_exception_targets)
136    /// Whether this block is an exception handler target (b_except_handler)
137    pub except_handler: bool,
138    /// Whether to preserve lasti for this handler block (b_preserve_lasti)
139    pub preserve_lasti: bool,
140    /// Stack depth at block entry, set by stack depth analysis
141    pub start_depth: Option<u32>,
142    /// Whether this block is only reachable via exception table (b_cold)
143    pub cold: bool,
144}
145
146impl Default for Block {
147    fn default() -> Self {
148        Self {
149            instructions: Vec::new(),
150            next: BlockIdx::NULL,
151            except_handler: false,
152            preserve_lasti: false,
153            start_depth: None,
154            cold: false,
155        }
156    }
157}
158
159pub struct CodeInfo {
160    pub flags: CodeFlags,
161    pub source_path: String,
162    pub private: Option<String>, // For private name mangling, mostly for class
163
164    pub blocks: Vec<Block>,
165    pub current_block: BlockIdx,
166
167    pub metadata: CodeUnitMetadata,
168
169    // For class scopes: attributes accessed via self.X
170    pub static_attributes: Option<IndexSet<String>>,
171
172    // True if compiling an inlined comprehension
173    pub in_inlined_comp: bool,
174
175    // Block stack for tracking nested control structures
176    pub fblock: Vec<crate::compile::FBlockInfo>,
177
178    // Reference to the symbol table for this scope
179    pub symbol_table_index: usize,
180
181    // PEP 649: Track nesting depth inside conditional blocks (if/for/while/etc.)
182    // u_in_conditional_block
183    pub in_conditional_block: u32,
184
185    // PEP 649: Next index for conditional annotation tracking
186    // u_next_conditional_annotation_index
187    pub next_conditional_annotation_index: u32,
188}
189
190impl CodeInfo {
191    pub fn finalize_code(
192        mut self,
193        opts: &crate::compile::CompileOpts,
194    ) -> crate::InternalResult<CodeObject> {
195        // Constant folding passes
196        self.fold_binop_constants();
197        self.remove_nops();
198        self.fold_unary_negative();
199        self.fold_binop_constants(); // re-run after unary folding: -1 + 2 → 1
200        self.remove_nops(); // remove NOPs so tuple/list/set see contiguous LOADs
201        self.fold_tuple_constants();
202        self.fold_list_constants();
203        self.fold_set_constants();
204        self.remove_nops(); // remove NOPs from collection folding
205        self.fold_const_iterable_for_iter();
206        self.convert_to_load_small_int();
207        self.remove_unused_consts();
208        self.remove_nops();
209
210        // DCE always runs (removes dead code after terminal instructions)
211        self.dce();
212        // Peephole optimizer creates superinstructions matching CPython
213        self.peephole_optimize();
214
215        // Phase 1: _PyCfg_OptimizeCodeUnit (flowgraph.c)
216        // Split blocks so each block has at most one branch as its last instruction
217        split_blocks_at_jumps(&mut self.blocks);
218        mark_except_handlers(&mut self.blocks);
219        label_exception_targets(&mut self.blocks);
220        // optimize_cfg: jump threading (before push_cold_blocks_to_end)
221        jump_threading(&mut self.blocks);
222        self.eliminate_unreachable_blocks();
223        self.remove_nops();
224        // TODO: insert_superinstructions disabled pending StoreFastLoadFast VM fix
225        push_cold_blocks_to_end(&mut self.blocks);
226
227        // Phase 2: _PyCfg_OptimizedCfgToInstructionSequence (flowgraph.c)
228        normalize_jumps(&mut self.blocks);
229        inline_small_or_no_lineno_blocks(&mut self.blocks);
230        self.dce(); // re-run within-block DCE after normalize_jumps creates new instructions
231        self.eliminate_unreachable_blocks();
232        resolve_line_numbers(&mut self.blocks);
233        duplicate_end_returns(&mut self.blocks);
234        self.dce(); // truncate after terminal in blocks that got return duplicated
235        self.eliminate_unreachable_blocks(); // remove now-unreachable last block
236        // optimize_load_fast: after normalize_jumps
237        self.optimize_load_fast_borrow();
238        self.optimize_load_global_push_null();
239
240        let max_stackdepth = self.max_stackdepth()?;
241
242        let Self {
243            flags,
244            source_path,
245            private: _, // private is only used during compilation
246
247            mut blocks,
248            current_block: _,
249            metadata,
250            static_attributes: _,
251            in_inlined_comp: _,
252            fblock: _,
253            symbol_table_index: _,
254            in_conditional_block: _,
255            next_conditional_annotation_index: _,
256        } = self;
257
258        let CodeUnitMetadata {
259            name: obj_name,
260            qualname,
261            consts: constants,
262            names: name_cache,
263            varnames: varname_cache,
264            cellvars: cellvar_cache,
265            freevars: freevar_cache,
266            fast_hidden,
267            argcount: arg_count,
268            posonlyargcount: posonlyarg_count,
269            kwonlyargcount: kwonlyarg_count,
270            firstlineno: first_line_number,
271        } = metadata;
272
273        let mut instructions = Vec::new();
274        let mut locations = Vec::new();
275        let mut linetable_locations: Vec<LineTableLocation> = Vec::new();
276
277        // Build cellfixedoffsets for cell-local merging
278        let cellfixedoffsets =
279            build_cellfixedoffsets(&varname_cache, &cellvar_cache, &freevar_cache);
280        // Convert pseudo ops (LoadClosure uses cellfixedoffsets) and fixup DEREF opargs
281        convert_pseudo_ops(&mut blocks, &cellfixedoffsets);
282        fixup_deref_opargs(&mut blocks, &cellfixedoffsets);
283        // Remove redundant NOPs, keeping line-marker NOPs only when
284        // they are needed to preserve tracing.
285        let mut block_order = Vec::new();
286        let mut current = BlockIdx(0);
287        while current != BlockIdx::NULL {
288            block_order.push(current);
289            current = blocks[current.idx()].next;
290        }
291        for block_idx in block_order {
292            let bi = block_idx.idx();
293            let mut src_instructions = core::mem::take(&mut blocks[bi].instructions);
294            let mut kept = Vec::with_capacity(src_instructions.len());
295            let mut prev_lineno = -1i32;
296
297            for src in 0..src_instructions.len() {
298                let instr = src_instructions[src];
299                let lineno = instr
300                    .lineno_override
301                    .unwrap_or_else(|| instr.location.line.get() as i32);
302                let mut remove = false;
303
304                if matches!(instr.instr.real(), Some(Instruction::Nop)) {
305                    // Remove location-less NOPs.
306                    if lineno < 0 || prev_lineno == lineno {
307                        remove = true;
308                    }
309                    // Remove if the next instruction has same line or no line.
310                    else if src < src_instructions.len() - 1 {
311                        let next_lineno =
312                            src_instructions[src + 1]
313                                .lineno_override
314                                .unwrap_or_else(|| {
315                                    src_instructions[src + 1].location.line.get() as i32
316                                });
317                        if next_lineno == lineno {
318                            remove = true;
319                        } else if next_lineno < 0 {
320                            src_instructions[src + 1].lineno_override = Some(lineno);
321                            remove = true;
322                        }
323                    }
324                    // Last instruction in block: compare with first real location
325                    // in the next non-empty block.
326                    else {
327                        let mut next = blocks[bi].next;
328                        while next != BlockIdx::NULL && blocks[next.idx()].instructions.is_empty() {
329                            next = blocks[next.idx()].next;
330                        }
331                        if next != BlockIdx::NULL {
332                            let mut next_lineno = None;
333                            for next_instr in &blocks[next.idx()].instructions {
334                                let line = next_instr
335                                    .lineno_override
336                                    .unwrap_or_else(|| next_instr.location.line.get() as i32);
337                                if matches!(next_instr.instr.real(), Some(Instruction::Nop))
338                                    && line < 0
339                                {
340                                    continue;
341                                }
342                                next_lineno = Some(line);
343                                break;
344                            }
345                            if next_lineno.is_some_and(|line| line == lineno) {
346                                remove = true;
347                            }
348                        }
349                    }
350                }
351
352                if !remove {
353                    kept.push(instr);
354                    prev_lineno = lineno;
355                }
356            }
357
358            blocks[bi].instructions = kept;
359        }
360
361        // Final DCE: truncate instructions after terminal ops in linearized blocks.
362        // This catches dead code created by normalize_jumps after the initial DCE.
363        for block in blocks.iter_mut() {
364            if let Some(pos) = block
365                .instructions
366                .iter()
367                .position(|ins| ins.instr.is_scope_exit() || ins.instr.is_unconditional_jump())
368            {
369                block.instructions.truncate(pos + 1);
370            }
371        }
372
373        // Pre-compute cache_entries for real (non-pseudo) instructions
374        for block in blocks.iter_mut() {
375            for instr in &mut block.instructions {
376                if let AnyInstruction::Real(op) = instr.instr {
377                    instr.cache_entries = op.cache_entries() as u32;
378                }
379            }
380        }
381
382        let mut block_to_offset = vec![Label::from_u32(0); blocks.len()];
383        // block_to_index: maps block idx to instruction index (for exception table)
384        // This is the index into the final instructions array, including EXTENDED_ARG and CACHE
385        let mut block_to_index = vec![0u32; blocks.len()];
386        // The offset (in code units) of END_SEND from SEND in the yield-from sequence.
387        const END_SEND_OFFSET: u32 = 5;
388        loop {
389            let mut num_instructions = 0;
390            for (idx, block) in iter_blocks(&blocks) {
391                block_to_offset[idx.idx()] = Label::from_u32(num_instructions as u32);
392                // block_to_index uses the same value as block_to_offset but as u32
393                // because lasti in frame.rs is the index into instructions array
394                // and instructions array index == byte offset (each instruction is 1 CodeUnit)
395                block_to_index[idx.idx()] = num_instructions as u32;
396                for instr in &block.instructions {
397                    num_instructions += instr.arg.instr_size() + instr.cache_entries as usize;
398                }
399            }
400
401            instructions.reserve_exact(num_instructions);
402            locations.reserve_exact(num_instructions);
403
404            let mut recompile = false;
405            let mut next_block = BlockIdx(0);
406            while next_block != BlockIdx::NULL {
407                let block = &mut blocks[next_block];
408                // Track current instruction offset for jump direction resolution
409                let mut current_offset = block_to_offset[next_block.idx()].as_u32();
410                for info in &mut block.instructions {
411                    let target = info.target;
412                    let mut op = info.instr.expect_real();
413                    let old_arg_size = info.arg.instr_size();
414                    let old_cache_entries = info.cache_entries;
415                    // Keep offsets fixed within this pass: changes in jump
416                    // arg/cache sizes only take effect in the next iteration.
417                    let offset_after = current_offset + old_arg_size as u32 + old_cache_entries;
418
419                    if target != BlockIdx::NULL {
420                        let target_offset = block_to_offset[target.idx()].as_u32();
421                        // Direction must be based on concrete instruction offsets.
422                        // Empty blocks can share offsets, so block-order-based resolution
423                        // may classify some jumps incorrectly.
424                        op = match op {
425                            Instruction::JumpForward { .. } if target_offset <= current_offset => {
426                                Instruction::JumpBackward {
427                                    delta: Arg::marker(),
428                                }
429                            }
430                            Instruction::JumpBackward { .. } if target_offset > current_offset => {
431                                Instruction::JumpForward {
432                                    delta: Arg::marker(),
433                                }
434                            }
435                            Instruction::JumpBackwardNoInterrupt { .. }
436                                if target_offset > current_offset =>
437                            {
438                                Instruction::JumpForward {
439                                    delta: Arg::marker(),
440                                }
441                            }
442                            _ => op,
443                        };
444                        info.instr = op.into();
445                        let updated_cache = op.cache_entries() as u32;
446                        recompile |= updated_cache != old_cache_entries;
447                        info.cache_entries = updated_cache;
448                        let new_arg = if matches!(op, Instruction::EndAsyncFor) {
449                            let arg = offset_after
450                                .checked_sub(target_offset + END_SEND_OFFSET)
451                                .expect("END_ASYNC_FOR target must be before instruction");
452                            OpArg::new(arg)
453                        } else if matches!(
454                            op,
455                            Instruction::JumpBackward { .. }
456                                | Instruction::JumpBackwardNoInterrupt { .. }
457                        ) {
458                            let arg = offset_after
459                                .checked_sub(target_offset)
460                                .expect("backward jump target must be before instruction");
461                            OpArg::new(arg)
462                        } else {
463                            let arg = target_offset
464                                .checked_sub(offset_after)
465                                .expect("forward jump target must be after instruction");
466                            OpArg::new(arg)
467                        };
468                        recompile |= new_arg.instr_size() != old_arg_size;
469                        info.arg = new_arg;
470                    }
471
472                    let cache_count = info.cache_entries as usize;
473                    let (extras, lo_arg) = info.arg.split();
474                    let loc_pair = (info.location, info.end_location);
475                    locations.extend(core::iter::repeat_n(
476                        loc_pair,
477                        info.arg.instr_size() + cache_count,
478                    ));
479                    // Collect linetable locations with lineno_override support
480                    let lt_loc = LineTableLocation {
481                        line: info
482                            .lineno_override
483                            .unwrap_or_else(|| info.location.line.get() as i32),
484                        end_line: info.end_location.line.get() as i32,
485                        col: info.location.character_offset.to_zero_indexed() as i32,
486                        end_col: info.end_location.character_offset.to_zero_indexed() as i32,
487                    };
488                    linetable_locations.extend(core::iter::repeat_n(lt_loc, info.arg.instr_size()));
489                    // CACHE entries inherit parent instruction's location
490                    if cache_count > 0 {
491                        linetable_locations.extend(core::iter::repeat_n(lt_loc, cache_count));
492                    }
493                    instructions.extend(
494                        extras
495                            .map(|byte| CodeUnit::new(Instruction::ExtendedArg, byte))
496                            .chain([CodeUnit { op, arg: lo_arg }]),
497                    );
498                    // Emit CACHE code units after the instruction (all zeroed)
499                    if cache_count > 0 {
500                        instructions.extend(core::iter::repeat_n(
501                            CodeUnit::new(Instruction::Cache, 0.into()),
502                            cache_count,
503                        ));
504                    }
505                    current_offset = offset_after;
506                }
507                next_block = block.next;
508            }
509
510            if !recompile {
511                break;
512            }
513
514            instructions.clear();
515            locations.clear();
516            linetable_locations.clear();
517        }
518
519        // Generate linetable from linetable_locations (supports line 0 for RESUME)
520        let linetable = generate_linetable(
521            &linetable_locations,
522            first_line_number.get() as i32,
523            opts.debug_ranges,
524        );
525
526        // Generate exception table before moving source_path
527        let exceptiontable = generate_exception_table(&blocks, &block_to_index);
528
529        // Build localspluskinds with cell-local merging
530        let nlocals = varname_cache.len();
531        let ncells = cellvar_cache.len();
532        let nfrees = freevar_cache.len();
533        let numdropped = cellvar_cache
534            .iter()
535            .filter(|cv| varname_cache.contains(cv.as_str()))
536            .count();
537        let nlocalsplus = nlocals + ncells - numdropped + nfrees;
538        let mut localspluskinds = vec![0u8; nlocalsplus];
539        // Mark locals
540        for kind in localspluskinds.iter_mut().take(nlocals) {
541            *kind = CO_FAST_LOCAL;
542        }
543        // Mark cells (merged and non-merged)
544        for (i, cellvar) in cellvar_cache.iter().enumerate() {
545            let idx = cellfixedoffsets[i] as usize;
546            if varname_cache.contains(cellvar.as_str()) {
547                localspluskinds[idx] |= CO_FAST_CELL; // merged: LOCAL | CELL
548            } else {
549                localspluskinds[idx] = CO_FAST_CELL;
550            }
551        }
552        // Mark frees
553        for i in 0..nfrees {
554            let idx = cellfixedoffsets[ncells + i] as usize;
555            localspluskinds[idx] = CO_FAST_FREE;
556        }
557        // Apply CO_FAST_HIDDEN for inlined comprehension variables
558        for (name, &hidden) in &fast_hidden {
559            if hidden && let Some(idx) = varname_cache.get_index_of(name.as_str()) {
560                localspluskinds[idx] |= CO_FAST_HIDDEN;
561            }
562        }
563
564        Ok(CodeObject {
565            flags,
566            posonlyarg_count,
567            arg_count,
568            kwonlyarg_count,
569            source_path,
570            first_line_number: Some(first_line_number),
571            obj_name: obj_name.clone(),
572            qualname: qualname.unwrap_or(obj_name),
573
574            max_stackdepth,
575            instructions: CodeUnits::from(instructions),
576            locations: locations.into_boxed_slice(),
577            constants: constants.into_iter().collect(),
578            names: name_cache.into_iter().collect(),
579            varnames: varname_cache.into_iter().collect(),
580            cellvars: cellvar_cache.into_iter().collect(),
581            freevars: freevar_cache.into_iter().collect(),
582            localspluskinds: localspluskinds.into_boxed_slice(),
583            linetable,
584            exceptiontable,
585        })
586    }
587
588    fn dce(&mut self) {
589        // Truncate instructions after terminal instructions within each block
590        for block in &mut self.blocks {
591            let mut last_instr = None;
592            for (i, ins) in block.instructions.iter().enumerate() {
593                if ins.instr.is_scope_exit() || ins.instr.is_unconditional_jump() {
594                    last_instr = Some(i);
595                    break;
596                }
597            }
598            if let Some(i) = last_instr {
599                block.instructions.truncate(i + 1);
600            }
601        }
602    }
603
604    /// Clear blocks that are unreachable (not entry, not a jump target,
605    /// and only reachable via fall-through from a terminal block).
606    fn eliminate_unreachable_blocks(&mut self) {
607        let mut reachable = vec![false; self.blocks.len()];
608        reachable[0] = true;
609
610        // Fixpoint: only mark targets of already-reachable blocks
611        let mut changed = true;
612        while changed {
613            changed = false;
614            for i in 0..self.blocks.len() {
615                if !reachable[i] {
616                    continue;
617                }
618                // Mark jump targets and exception handlers
619                for ins in &self.blocks[i].instructions {
620                    if ins.target != BlockIdx::NULL && !reachable[ins.target.idx()] {
621                        reachable[ins.target.idx()] = true;
622                        changed = true;
623                    }
624                    if let Some(eh) = &ins.except_handler
625                        && !reachable[eh.handler_block.idx()]
626                    {
627                        reachable[eh.handler_block.idx()] = true;
628                        changed = true;
629                    }
630                }
631                // Mark fall-through
632                let next = self.blocks[i].next;
633                if next != BlockIdx::NULL
634                    && !reachable[next.idx()]
635                    && !self.blocks[i].instructions.last().is_some_and(|ins| {
636                        ins.instr.is_scope_exit() || ins.instr.is_unconditional_jump()
637                    })
638                {
639                    reachable[next.idx()] = true;
640                    changed = true;
641                }
642            }
643        }
644
645        for (i, block) in self.blocks.iter_mut().enumerate() {
646            if !reachable[i] {
647                block.instructions.clear();
648            }
649        }
650    }
651
652    /// Fold LOAD_CONST/LOAD_SMALL_INT + UNARY_NEGATIVE → LOAD_CONST (negative value)
653    fn fold_unary_negative(&mut self) {
654        for block in &mut self.blocks {
655            let mut i = 0;
656            while i + 1 < block.instructions.len() {
657                let next = &block.instructions[i + 1];
658                let Some(Instruction::UnaryNegative) = next.instr.real() else {
659                    i += 1;
660                    continue;
661                };
662                let curr = &block.instructions[i];
663                let value = match curr.instr.real() {
664                    Some(Instruction::LoadConst { .. }) => {
665                        let idx = u32::from(curr.arg) as usize;
666                        match self.metadata.consts.get_index(idx) {
667                            Some(ConstantData::Integer { value }) => {
668                                Some(ConstantData::Integer { value: -value })
669                            }
670                            Some(ConstantData::Float { value }) => {
671                                Some(ConstantData::Float { value: -value })
672                            }
673                            _ => None,
674                        }
675                    }
676                    Some(Instruction::LoadSmallInt { .. }) => {
677                        let v = u32::from(curr.arg) as i32;
678                        Some(ConstantData::Integer {
679                            value: BigInt::from(-v),
680                        })
681                    }
682                    _ => None,
683                };
684                if let Some(neg_const) = value {
685                    let (const_idx, _) = self.metadata.consts.insert_full(neg_const);
686                    // Replace LOAD_CONST/LOAD_SMALL_INT with new LOAD_CONST
687                    let load_location = block.instructions[i].location;
688                    block.instructions[i].instr = Instruction::LoadConst {
689                        consti: Arg::marker(),
690                    }
691                    .into();
692                    block.instructions[i].arg = OpArg::new(const_idx as u32);
693                    // Replace UNARY_NEGATIVE with NOP, inheriting the LOAD_CONST
694                    // location so that remove_nops can clean it up
695                    block.instructions[i + 1].instr = Instruction::Nop.into();
696                    block.instructions[i + 1].location = load_location;
697                    block.instructions[i + 1].end_location = block.instructions[i].end_location;
698                    // Skip the NOP, don't re-check
699                    i += 2;
700                } else {
701                    i += 1;
702                }
703            }
704        }
705    }
706
707    /// Constant folding: fold LOAD_CONST/LOAD_SMALL_INT + LOAD_CONST/LOAD_SMALL_INT + BINARY_OP
708    /// into a single LOAD_CONST when the result is computable at compile time.
709    /// = fold_binops_on_constants in CPython flowgraph.c
710    fn fold_binop_constants(&mut self) {
711        use oparg::BinaryOperator as BinOp;
712
713        for block in &mut self.blocks {
714            let mut i = 0;
715            while i + 2 < block.instructions.len() {
716                // Check pattern: LOAD_CONST/LOAD_SMALL_INT, LOAD_CONST/LOAD_SMALL_INT, BINARY_OP
717                let Some(Instruction::BinaryOp { .. }) = block.instructions[i + 2].instr.real()
718                else {
719                    i += 1;
720                    continue;
721                };
722
723                let op_raw = u32::from(block.instructions[i + 2].arg);
724                let Ok(op) = BinOp::try_from(op_raw) else {
725                    i += 1;
726                    continue;
727                };
728
729                let left = Self::get_const_value_from(&self.metadata, &block.instructions[i]);
730                let right = Self::get_const_value_from(&self.metadata, &block.instructions[i + 1]);
731
732                let (Some(left_val), Some(right_val)) = (left, right) else {
733                    i += 1;
734                    continue;
735                };
736
737                let result = Self::eval_binop(&left_val, &right_val, op);
738
739                if let Some(result_const) = result {
740                    // Check result size limit (CPython limits to 4096 bytes)
741                    if Self::const_too_big(&result_const) {
742                        i += 1;
743                        continue;
744                    }
745                    let (const_idx, _) = self.metadata.consts.insert_full(result_const);
746                    // Replace first instruction with LOAD_CONST result
747                    block.instructions[i].instr = Instruction::LoadConst {
748                        consti: Arg::marker(),
749                    }
750                    .into();
751                    block.instructions[i].arg = OpArg::new(const_idx as u32);
752                    // NOP out the second and third instructions
753                    let loc = block.instructions[i].location;
754                    let end_loc = block.instructions[i].end_location;
755                    block.instructions[i + 1].instr = Instruction::Nop.into();
756                    block.instructions[i + 1].location = loc;
757                    block.instructions[i + 1].end_location = end_loc;
758                    block.instructions[i + 2].instr = Instruction::Nop.into();
759                    block.instructions[i + 2].location = loc;
760                    block.instructions[i + 2].end_location = end_loc;
761                    // Don't advance - check if the result can be folded again
762                    // (e.g., 2 ** 31 - 1)
763                    i = i.saturating_sub(1); // re-check with previous instruction
764                } else {
765                    i += 1;
766                }
767            }
768        }
769    }
770
771    fn get_const_value_from(
772        metadata: &CodeUnitMetadata,
773        info: &InstructionInfo,
774    ) -> Option<ConstantData> {
775        match info.instr.real() {
776            Some(Instruction::LoadConst { .. }) => {
777                let idx = u32::from(info.arg) as usize;
778                metadata.consts.get_index(idx).cloned()
779            }
780            Some(Instruction::LoadSmallInt { .. }) => {
781                let v = u32::from(info.arg) as i32;
782                Some(ConstantData::Integer {
783                    value: BigInt::from(v),
784                })
785            }
786            _ => None,
787        }
788    }
789
790    fn eval_binop(
791        left: &ConstantData,
792        right: &ConstantData,
793        op: oparg::BinaryOperator,
794    ) -> Option<ConstantData> {
795        use oparg::BinaryOperator as BinOp;
796        match (left, right) {
797            (ConstantData::Integer { value: l }, ConstantData::Integer { value: r }) => {
798                let result = match op {
799                    BinOp::Add => l + r,
800                    BinOp::Subtract => l - r,
801                    BinOp::Multiply => {
802                        if !l.is_zero() && !r.is_zero() && l.bits() + r.bits() > MAX_INT_SIZE_BITS {
803                            return None;
804                        }
805                        l * r
806                    }
807                    BinOp::FloorDivide => {
808                        if r.is_zero() {
809                            return None;
810                        }
811                        // Python floor division: round towards negative infinity
812                        let (q, rem) = (l.clone() / r.clone(), l.clone() % r.clone());
813                        if !rem.is_zero() && (rem < BigInt::from(0)) != (*r < BigInt::from(0)) {
814                            q - 1
815                        } else {
816                            q
817                        }
818                    }
819                    BinOp::Remainder => {
820                        if r.is_zero() {
821                            return None;
822                        }
823                        // Python modulo: result has same sign as divisor
824                        let rem = l.clone() % r.clone();
825                        if !rem.is_zero() && (rem < BigInt::from(0)) != (*r < BigInt::from(0)) {
826                            rem + r
827                        } else {
828                            rem
829                        }
830                    }
831                    BinOp::Power => {
832                        let exp: u64 = r.try_into().ok()?;
833                        let exp_usize = usize::try_from(exp).ok()?;
834                        if !l.is_zero() && exp > 0 && l.bits() > MAX_INT_SIZE_BITS / exp {
835                            return None;
836                        }
837                        num_traits::pow::pow(l.clone(), exp_usize)
838                    }
839                    BinOp::Lshift => {
840                        let shift: u64 = r.try_into().ok()?;
841                        let shift_usize = usize::try_from(shift).ok()?;
842                        if shift > MAX_INT_SIZE_BITS
843                            || (!l.is_zero() && l.bits() > MAX_INT_SIZE_BITS - shift)
844                        {
845                            return None;
846                        }
847                        l << shift_usize
848                    }
849                    BinOp::Rshift => {
850                        let shift: u32 = r.try_into().ok()?;
851                        l >> (shift as usize)
852                    }
853                    BinOp::And => l & r,
854                    BinOp::Or => l | r,
855                    BinOp::Xor => l ^ r,
856                    _ => return None,
857                };
858                Some(ConstantData::Integer { value: result })
859            }
860            (ConstantData::Float { value: l }, ConstantData::Float { value: r }) => {
861                let result = match op {
862                    BinOp::Add => l + r,
863                    BinOp::Subtract => l - r,
864                    BinOp::Multiply => l * r,
865                    BinOp::TrueDivide => {
866                        if *r == 0.0 {
867                            return None;
868                        }
869                        l / r
870                    }
871                    BinOp::FloorDivide => {
872                        // Float floor division uses runtime semantics; skip folding
873                        return None;
874                    }
875                    BinOp::Remainder => {
876                        // Float modulo uses fmod() at runtime; Rust arithmetic differs
877                        return None;
878                    }
879                    BinOp::Power => l.powf(*r),
880                    _ => return None,
881                };
882                if !result.is_finite() {
883                    return None;
884                }
885                Some(ConstantData::Float { value: result })
886            }
887            // Int op Float or Float op Int → Float
888            (ConstantData::Integer { value: l }, ConstantData::Float { value: r }) => {
889                let l_f = l.to_f64()?;
890                Self::eval_binop(
891                    &ConstantData::Float { value: l_f },
892                    &ConstantData::Float { value: *r },
893                    op,
894                )
895            }
896            (ConstantData::Float { value: l }, ConstantData::Integer { value: r }) => {
897                let r_f = r.to_f64()?;
898                Self::eval_binop(
899                    &ConstantData::Float { value: *l },
900                    &ConstantData::Float { value: r_f },
901                    op,
902                )
903            }
904            // String concatenation and repetition
905            (ConstantData::Str { value: l }, ConstantData::Str { value: r })
906                if matches!(op, BinOp::Add) =>
907            {
908                let mut result = l.to_string();
909                result.push_str(&r.to_string());
910                Some(ConstantData::Str {
911                    value: result.into(),
912                })
913            }
914            (ConstantData::Str { value: s }, ConstantData::Integer { value: n })
915                if matches!(op, BinOp::Multiply) =>
916            {
917                let n: usize = n.try_into().ok()?;
918                if n > 4096 {
919                    return None;
920                }
921                let result = s.to_string().repeat(n);
922                Some(ConstantData::Str {
923                    value: result.into(),
924                })
925            }
926            _ => None,
927        }
928    }
929
930    fn const_too_big(c: &ConstantData) -> bool {
931        match c {
932            ConstantData::Integer { value } => value.bits() > 4096 * 8,
933            ConstantData::Str { value } => value.len() > 4096,
934            _ => false,
935        }
936    }
937
938    /// Constant folding: fold LOAD_CONST/LOAD_SMALL_INT + BUILD_TUPLE into LOAD_CONST tuple
939    /// fold_tuple_of_constants
940    fn fold_tuple_constants(&mut self) {
941        for block in &mut self.blocks {
942            let mut i = 0;
943            while i < block.instructions.len() {
944                let instr = &block.instructions[i];
945                // Look for BUILD_TUPLE
946                let Some(Instruction::BuildTuple { .. }) = instr.instr.real() else {
947                    i += 1;
948                    continue;
949                };
950
951                let tuple_size = u32::from(instr.arg) as usize;
952                if tuple_size == 0 {
953                    // BUILD_TUPLE 0 → LOAD_CONST ()
954                    let (const_idx, _) = self.metadata.consts.insert_full(ConstantData::Tuple {
955                        elements: Vec::new(),
956                    });
957                    block.instructions[i].instr = Instruction::LoadConst {
958                        consti: Arg::marker(),
959                    }
960                    .into();
961                    block.instructions[i].arg = OpArg::new(const_idx as u32);
962                    i += 1;
963                    continue;
964                }
965                if i < tuple_size {
966                    i += 1;
967                    continue;
968                }
969
970                // Check if all preceding instructions are constant-loading
971                let start_idx = i - tuple_size;
972                let mut elements = Vec::with_capacity(tuple_size);
973                let mut all_const = true;
974
975                for j in start_idx..i {
976                    let load_instr = &block.instructions[j];
977                    match load_instr.instr.real() {
978                        Some(Instruction::LoadConst { .. }) => {
979                            let const_idx = u32::from(load_instr.arg) as usize;
980                            if let Some(constant) =
981                                self.metadata.consts.get_index(const_idx).cloned()
982                            {
983                                elements.push(constant);
984                            } else {
985                                all_const = false;
986                                break;
987                            }
988                        }
989                        Some(Instruction::LoadSmallInt { .. }) => {
990                            // arg is the i32 value stored as u32 (two's complement)
991                            let value = u32::from(load_instr.arg) as i32;
992                            elements.push(ConstantData::Integer {
993                                value: BigInt::from(value),
994                            });
995                        }
996                        _ => {
997                            all_const = false;
998                            break;
999                        }
1000                    }
1001                }
1002
1003                if !all_const {
1004                    i += 1;
1005                    continue;
1006                }
1007
1008                // Note: The first small int is added to co_consts during compilation
1009                // (in compile_default_arguments).
1010                // We don't need to add it here again.
1011
1012                // Create tuple constant and add to consts
1013                let tuple_const = ConstantData::Tuple { elements };
1014                let (const_idx, _) = self.metadata.consts.insert_full(tuple_const);
1015
1016                // Replace preceding LOAD instructions with NOP at the
1017                // BUILD_TUPLE location so remove_nops() can eliminate them.
1018                let folded_loc = block.instructions[i].location;
1019                for j in start_idx..i {
1020                    block.instructions[j].instr = Instruction::Nop.into();
1021                    block.instructions[j].location = folded_loc;
1022                }
1023
1024                // Replace BUILD_TUPLE with LOAD_CONST
1025                block.instructions[i].instr = Instruction::LoadConst {
1026                    consti: Arg::marker(),
1027                }
1028                .into();
1029                block.instructions[i].arg = OpArg::new(const_idx as u32);
1030
1031                i += 1;
1032            }
1033        }
1034    }
1035
1036    /// Fold constant list literals: LOAD_CONST* + BUILD_LIST N →
1037    /// BUILD_LIST 0 + LOAD_CONST (tuple) + LIST_EXTEND 1
1038    fn fold_list_constants(&mut self) {
1039        for block in &mut self.blocks {
1040            let mut i = 0;
1041            while i < block.instructions.len() {
1042                let instr = &block.instructions[i];
1043                let Some(Instruction::BuildList { .. }) = instr.instr.real() else {
1044                    i += 1;
1045                    continue;
1046                };
1047
1048                let list_size = u32::from(instr.arg) as usize;
1049                if list_size == 0 || i < list_size {
1050                    i += 1;
1051                    continue;
1052                }
1053
1054                let start_idx = i - list_size;
1055                let mut elements = Vec::with_capacity(list_size);
1056                let mut all_const = true;
1057
1058                for j in start_idx..i {
1059                    let load_instr = &block.instructions[j];
1060                    match load_instr.instr.real() {
1061                        Some(Instruction::LoadConst { .. }) => {
1062                            let const_idx = u32::from(load_instr.arg) as usize;
1063                            if let Some(constant) =
1064                                self.metadata.consts.get_index(const_idx).cloned()
1065                            {
1066                                elements.push(constant);
1067                            } else {
1068                                all_const = false;
1069                                break;
1070                            }
1071                        }
1072                        Some(Instruction::LoadSmallInt { .. }) => {
1073                            let value = u32::from(load_instr.arg) as i32;
1074                            elements.push(ConstantData::Integer {
1075                                value: BigInt::from(value),
1076                            });
1077                        }
1078                        _ => {
1079                            all_const = false;
1080                            break;
1081                        }
1082                    }
1083                }
1084
1085                if !all_const || list_size < MIN_CONST_SEQUENCE_SIZE {
1086                    i += 1;
1087                    continue;
1088                }
1089
1090                let tuple_const = ConstantData::Tuple { elements };
1091                let (const_idx, _) = self.metadata.consts.insert_full(tuple_const);
1092
1093                let folded_loc = block.instructions[i].location;
1094                let end_loc = block.instructions[i].end_location;
1095                let eh = block.instructions[i].except_handler;
1096
1097                // slot[start_idx] → BUILD_LIST 0
1098                block.instructions[start_idx].instr = Instruction::BuildList {
1099                    count: Arg::marker(),
1100                }
1101                .into();
1102                block.instructions[start_idx].arg = OpArg::new(0);
1103                block.instructions[start_idx].location = folded_loc;
1104                block.instructions[start_idx].end_location = end_loc;
1105                block.instructions[start_idx].except_handler = eh;
1106
1107                // slot[start_idx+1] → LOAD_CONST (tuple)
1108                block.instructions[start_idx + 1].instr = Instruction::LoadConst {
1109                    consti: Arg::marker(),
1110                }
1111                .into();
1112                block.instructions[start_idx + 1].arg = OpArg::new(const_idx as u32);
1113                block.instructions[start_idx + 1].location = folded_loc;
1114                block.instructions[start_idx + 1].end_location = end_loc;
1115                block.instructions[start_idx + 1].except_handler = eh;
1116
1117                // NOP the rest
1118                for j in (start_idx + 2)..i {
1119                    block.instructions[j].instr = Instruction::Nop.into();
1120                    block.instructions[j].location = folded_loc;
1121                }
1122
1123                // slot[i] (was BUILD_LIST) → LIST_EXTEND 1
1124                block.instructions[i].instr = Instruction::ListExtend { i: Arg::marker() }.into();
1125                block.instructions[i].arg = OpArg::new(1);
1126
1127                i += 1;
1128            }
1129        }
1130    }
1131
1132    /// Convert constant list construction before GET_ITER to just LOAD_CONST tuple.
1133    /// BUILD_LIST 0 + LOAD_CONST (tuple) + LIST_EXTEND 1 + GET_ITER
1134    /// → LOAD_CONST (tuple) + GET_ITER
1135    fn fold_const_iterable_for_iter(&mut self) {
1136        for block in &mut self.blocks {
1137            let mut i = 0;
1138            while i + 1 < block.instructions.len() {
1139                let is_build = matches!(
1140                    block.instructions[i].instr.real(),
1141                    Some(Instruction::BuildList { .. })
1142                ) && u32::from(block.instructions[i].arg) == 0;
1143
1144                let is_const = matches!(
1145                    block
1146                        .instructions
1147                        .get(i + 1)
1148                        .and_then(|instr| instr.instr.real()),
1149                    Some(Instruction::LoadConst { .. })
1150                );
1151
1152                let is_extend = matches!(
1153                    block
1154                        .instructions
1155                        .get(i + 2)
1156                        .and_then(|instr| instr.instr.real()),
1157                    Some(Instruction::ListExtend { .. })
1158                ) && block
1159                    .instructions
1160                    .get(i + 2)
1161                    .is_some_and(|instr| u32::from(instr.arg) == 1);
1162
1163                let is_iter = matches!(
1164                    block
1165                        .instructions
1166                        .get(i + 3)
1167                        .and_then(|instr| instr.instr.real()),
1168                    Some(Instruction::GetIter)
1169                );
1170
1171                if is_build && is_const && is_extend && is_iter {
1172                    // Replace: BUILD_X 0 → NOP, keep LOAD_CONST, LIST_EXTEND → NOP
1173                    let loc = block.instructions[i].location;
1174                    block.instructions[i].instr = Instruction::Nop.into();
1175                    block.instructions[i].location = loc;
1176                    block.instructions[i + 2].instr = Instruction::Nop.into();
1177                    block.instructions[i + 2].location = loc;
1178                    i += 4;
1179                } else if matches!(
1180                    block.instructions[i].instr.real(),
1181                    Some(Instruction::BuildList { .. })
1182                ) && matches!(
1183                    block.instructions[i + 1].instr.real(),
1184                    Some(Instruction::GetIter)
1185                ) {
1186                    let seq_size = u32::from(block.instructions[i].arg) as usize;
1187
1188                    if seq_size != 0 && i >= seq_size {
1189                        let start_idx = i - seq_size;
1190                        let mut elements = Vec::with_capacity(seq_size);
1191                        let mut all_const = true;
1192
1193                        for j in start_idx..i {
1194                            match Self::get_const_value_from(&self.metadata, &block.instructions[j])
1195                            {
1196                                Some(constant) => elements.push(constant),
1197                                None => {
1198                                    all_const = false;
1199                                    break;
1200                                }
1201                            }
1202                        }
1203
1204                        if all_const {
1205                            let const_data = ConstantData::Tuple { elements };
1206                            let (const_idx, _) = self.metadata.consts.insert_full(const_data);
1207                            let folded_loc = block.instructions[i].location;
1208
1209                            for j in start_idx..i {
1210                                block.instructions[j].instr = Instruction::Nop.into();
1211                                block.instructions[j].location = folded_loc;
1212                            }
1213
1214                            block.instructions[i].instr = Instruction::LoadConst {
1215                                consti: Arg::marker(),
1216                            }
1217                            .into();
1218                            block.instructions[i].arg = OpArg::new(const_idx as u32);
1219                            i += 2;
1220                            continue;
1221                        }
1222                    }
1223
1224                    block.instructions[i].instr = Instruction::BuildTuple {
1225                        count: Arg::marker(),
1226                    }
1227                    .into();
1228                    i += 2;
1229                } else {
1230                    i += 1;
1231                }
1232            }
1233        }
1234    }
1235
1236    /// Fold constant set literals: LOAD_CONST* + BUILD_SET N →
1237    /// BUILD_SET 0 + LOAD_CONST (frozenset-as-tuple) + SET_UPDATE 1
1238    fn fold_set_constants(&mut self) {
1239        for block in &mut self.blocks {
1240            let mut i = 0;
1241            while i < block.instructions.len() {
1242                let instr = &block.instructions[i];
1243                let Some(Instruction::BuildSet { .. }) = instr.instr.real() else {
1244                    i += 1;
1245                    continue;
1246                };
1247
1248                let set_size = u32::from(instr.arg) as usize;
1249                if set_size < 3 || i < set_size {
1250                    i += 1;
1251                    continue;
1252                }
1253
1254                let start_idx = i - set_size;
1255                let mut elements = Vec::with_capacity(set_size);
1256                let mut all_const = true;
1257
1258                for j in start_idx..i {
1259                    let load_instr = &block.instructions[j];
1260                    match load_instr.instr.real() {
1261                        Some(Instruction::LoadConst { .. }) => {
1262                            let const_idx = u32::from(load_instr.arg) as usize;
1263                            if let Some(constant) =
1264                                self.metadata.consts.get_index(const_idx).cloned()
1265                            {
1266                                elements.push(constant);
1267                            } else {
1268                                all_const = false;
1269                                break;
1270                            }
1271                        }
1272                        Some(Instruction::LoadSmallInt { .. }) => {
1273                            let value = u32::from(load_instr.arg) as i32;
1274                            elements.push(ConstantData::Integer {
1275                                value: BigInt::from(value),
1276                            });
1277                        }
1278                        _ => {
1279                            all_const = false;
1280                            break;
1281                        }
1282                    }
1283                }
1284
1285                if !all_const {
1286                    i += 1;
1287                    continue;
1288                }
1289
1290                // Use FrozenSet constant (stored as Tuple for now)
1291                let const_data = ConstantData::Tuple { elements };
1292                let (const_idx, _) = self.metadata.consts.insert_full(const_data);
1293
1294                let folded_loc = block.instructions[i].location;
1295                let end_loc = block.instructions[i].end_location;
1296                let eh = block.instructions[i].except_handler;
1297
1298                block.instructions[start_idx].instr = Instruction::BuildSet {
1299                    count: Arg::marker(),
1300                }
1301                .into();
1302                block.instructions[start_idx].arg = OpArg::new(0);
1303                block.instructions[start_idx].location = folded_loc;
1304                block.instructions[start_idx].end_location = end_loc;
1305                block.instructions[start_idx].except_handler = eh;
1306
1307                block.instructions[start_idx + 1].instr = Instruction::LoadConst {
1308                    consti: Arg::marker(),
1309                }
1310                .into();
1311                block.instructions[start_idx + 1].arg = OpArg::new(const_idx as u32);
1312                block.instructions[start_idx + 1].location = folded_loc;
1313                block.instructions[start_idx + 1].end_location = end_loc;
1314                block.instructions[start_idx + 1].except_handler = eh;
1315
1316                for j in (start_idx + 2)..i {
1317                    block.instructions[j].instr = Instruction::Nop.into();
1318                    block.instructions[j].location = folded_loc;
1319                }
1320
1321                block.instructions[i].instr = Instruction::SetUpdate { i: Arg::marker() }.into();
1322                block.instructions[i].arg = OpArg::new(1);
1323
1324                i += 1;
1325            }
1326        }
1327    }
1328
1329    /// Peephole optimization: combine consecutive instructions into super-instructions
1330    fn peephole_optimize(&mut self) {
1331        for block in &mut self.blocks {
1332            let mut i = 0;
1333            while i + 1 < block.instructions.len() {
1334                let combined = {
1335                    let curr = &block.instructions[i];
1336                    let next = &block.instructions[i + 1];
1337
1338                    // Only combine if both are real instructions (not pseudo)
1339                    let (Some(curr_instr), Some(next_instr)) =
1340                        (curr.instr.real(), next.instr.real())
1341                    else {
1342                        i += 1;
1343                        continue;
1344                    };
1345
1346                    match (curr_instr, next_instr) {
1347                        // LoadFast + LoadFast -> LoadFastLoadFast (if both indices < 16)
1348                        (Instruction::LoadFast { .. }, Instruction::LoadFast { .. }) => {
1349                            let idx1 = u32::from(curr.arg);
1350                            let idx2 = u32::from(next.arg);
1351                            if idx1 < 16 && idx2 < 16 {
1352                                let packed = (idx1 << 4) | idx2;
1353                                Some((
1354                                    Instruction::LoadFastLoadFast {
1355                                        var_nums: Arg::marker(),
1356                                    },
1357                                    OpArg::new(packed),
1358                                ))
1359                            } else {
1360                                None
1361                            }
1362                        }
1363                        // StoreFast + StoreFast -> StoreFastStoreFast (if both indices < 16)
1364                        (Instruction::StoreFast { .. }, Instruction::StoreFast { .. }) => {
1365                            let idx1 = u32::from(curr.arg);
1366                            let idx2 = u32::from(next.arg);
1367                            if idx1 < 16 && idx2 < 16 {
1368                                let packed = (idx1 << 4) | idx2;
1369                                Some((
1370                                    Instruction::StoreFastStoreFast {
1371                                        var_nums: Arg::marker(),
1372                                    },
1373                                    OpArg::new(packed),
1374                                ))
1375                            } else {
1376                                None
1377                            }
1378                        }
1379                        // Note: StoreFast + LoadFast → StoreFastLoadFast is done in a
1380                        // separate pass AFTER optimize_load_fast_borrow, because CPython
1381                        // only combines STORE_FAST + LOAD_FAST (not LOAD_FAST_BORROW).
1382                        (Instruction::LoadConst { consti }, Instruction::ToBool) => {
1383                            let consti = consti.get(curr.arg);
1384                            let constant = &self.metadata.consts[consti.as_usize()];
1385                            if let ConstantData::Boolean { .. } = constant {
1386                                Some((curr_instr, OpArg::from(consti.as_u32())))
1387                            } else {
1388                                None
1389                            }
1390                        }
1391                        (Instruction::LoadConst { consti }, Instruction::UnaryNot) => {
1392                            let constant = &self.metadata.consts[consti.get(curr.arg).as_usize()];
1393                            match constant {
1394                                ConstantData::Boolean { value } => {
1395                                    let (const_idx, _) = self
1396                                        .metadata
1397                                        .consts
1398                                        .insert_full(ConstantData::Boolean { value: !value });
1399                                    Some((
1400                                        (Instruction::LoadConst {
1401                                            consti: Arg::marker(),
1402                                        }),
1403                                        OpArg::new(const_idx as u32),
1404                                    ))
1405                                }
1406                                _ => None,
1407                            }
1408                        }
1409                        _ => None,
1410                    }
1411                };
1412
1413                if let Some((new_instr, new_arg)) = combined {
1414                    // Combine: keep first instruction's location, replace with combined instruction
1415                    block.instructions[i].instr = new_instr.into();
1416                    block.instructions[i].arg = new_arg;
1417                    // Remove the second instruction
1418                    block.instructions.remove(i + 1);
1419                    // Don't increment i - check if we can combine again with the next instruction
1420                } else {
1421                    i += 1;
1422                }
1423            }
1424        }
1425    }
1426
1427    /// LOAD_GLOBAL <even> + PUSH_NULL -> LOAD_GLOBAL <odd>, NOP
1428    fn optimize_load_global_push_null(&mut self) {
1429        for block in &mut self.blocks {
1430            let mut i = 0;
1431            while i + 1 < block.instructions.len() {
1432                let curr = &block.instructions[i];
1433                let next = &block.instructions[i + 1];
1434
1435                let (Some(Instruction::LoadGlobal { .. }), Some(Instruction::PushNull)) =
1436                    (curr.instr.real(), next.instr.real())
1437                else {
1438                    i += 1;
1439                    continue;
1440                };
1441
1442                let oparg = u32::from(block.instructions[i].arg);
1443                if (oparg & 1) != 0 {
1444                    i += 1;
1445                    continue;
1446                }
1447
1448                block.instructions[i].arg = OpArg::new(oparg | 1);
1449                block.instructions.remove(i + 1);
1450            }
1451        }
1452    }
1453
1454    /// Convert LOAD_CONST for small integers to LOAD_SMALL_INT
1455    /// maybe_instr_make_load_smallint
1456    fn convert_to_load_small_int(&mut self) {
1457        for block in &mut self.blocks {
1458            for instr in &mut block.instructions {
1459                // Check if it's a LOAD_CONST instruction
1460                let Some(Instruction::LoadConst { .. }) = instr.instr.real() else {
1461                    continue;
1462                };
1463
1464                // Get the constant value
1465                let const_idx = u32::from(instr.arg) as usize;
1466                let Some(constant) = self.metadata.consts.get_index(const_idx) else {
1467                    continue;
1468                };
1469
1470                // Check if it's a small integer
1471                let ConstantData::Integer { value } = constant else {
1472                    continue;
1473                };
1474
1475                // LOAD_SMALL_INT oparg is unsigned, so only 0..=255 can be encoded
1476                if let Some(small) = value.to_i32().filter(|v| (0..=255).contains(v)) {
1477                    // Convert LOAD_CONST to LOAD_SMALL_INT
1478                    instr.instr = Instruction::LoadSmallInt { i: Arg::marker() }.into();
1479                    // The arg is the i32 value stored as u32 (two's complement)
1480                    instr.arg = OpArg::new(small as u32);
1481                }
1482            }
1483        }
1484    }
1485
1486    /// Remove constants that are no longer referenced by LOAD_CONST instructions.
1487    /// remove_unused_consts
1488    fn remove_unused_consts(&mut self) {
1489        let nconsts = self.metadata.consts.len();
1490        if nconsts == 0 {
1491            return;
1492        }
1493
1494        // Mark used constants
1495        // The first constant (index 0) is always kept (may be docstring)
1496        let mut used = vec![false; nconsts];
1497        used[0] = true;
1498
1499        for block in &self.blocks {
1500            for instr in &block.instructions {
1501                if let Some(Instruction::LoadConst { .. }) = instr.instr.real() {
1502                    let idx = u32::from(instr.arg) as usize;
1503                    if idx < nconsts {
1504                        used[idx] = true;
1505                    }
1506                }
1507            }
1508        }
1509
1510        // Check if any constants can be removed
1511        let n_used: usize = used.iter().filter(|&&u| u).count();
1512        if n_used == nconsts {
1513            return; // Nothing to remove
1514        }
1515
1516        // Build old_to_new index mapping
1517        let mut old_to_new = vec![0usize; nconsts];
1518        let mut new_idx = 0usize;
1519        for (old_idx, &is_used) in used.iter().enumerate() {
1520            if is_used {
1521                old_to_new[old_idx] = new_idx;
1522                new_idx += 1;
1523            }
1524        }
1525
1526        // Build new consts list
1527        let old_consts: Vec<_> = self.metadata.consts.iter().cloned().collect();
1528        self.metadata.consts.clear();
1529        for (old_idx, constant) in old_consts.into_iter().enumerate() {
1530            if used[old_idx] {
1531                self.metadata.consts.insert(constant);
1532            }
1533        }
1534
1535        // Update LOAD_CONST instruction arguments
1536        for block in &mut self.blocks {
1537            for instr in &mut block.instructions {
1538                if let Some(Instruction::LoadConst { .. }) = instr.instr.real() {
1539                    let old_idx = u32::from(instr.arg) as usize;
1540                    if old_idx < nconsts {
1541                        instr.arg = OpArg::new(old_to_new[old_idx] as u32);
1542                    }
1543                }
1544            }
1545        }
1546    }
1547
1548    /// Remove NOP instructions from all blocks, but keep NOPs that introduce
1549    /// a new source line (they serve as line markers for monitoring LINE events).
1550    fn remove_nops(&mut self) {
1551        for block in &mut self.blocks {
1552            let mut prev_line = None;
1553            block.instructions.retain(|ins| {
1554                if matches!(ins.instr.real(), Some(Instruction::Nop)) {
1555                    let line = ins.location.line;
1556                    if prev_line == Some(line) {
1557                        return false;
1558                    }
1559                }
1560                prev_line = Some(ins.location.line);
1561                true
1562            });
1563        }
1564    }
1565
1566    /// Optimize LOAD_FAST to LOAD_FAST_BORROW where safe.
1567    ///
1568    /// insert_superinstructions (flowgraph.c): Combine STORE_FAST + LOAD_FAST →
1569    /// STORE_FAST_LOAD_FAST. Currently disabled pending VM stack null investigation.
1570    #[allow(dead_code)]
1571    fn combine_store_fast_load_fast(&mut self) {
1572        for block in &mut self.blocks {
1573            let mut i = 0;
1574            while i + 1 < block.instructions.len() {
1575                let curr = &block.instructions[i];
1576                let next = &block.instructions[i + 1];
1577                let (Some(Instruction::StoreFast { .. }), Some(Instruction::LoadFast { .. })) =
1578                    (curr.instr.real(), next.instr.real())
1579                else {
1580                    i += 1;
1581                    continue;
1582                };
1583                // Skip if instructions are on different lines (matching make_super_instruction)
1584                let line1 = curr.location.line;
1585                let line2 = next.location.line;
1586                if line1 != line2 {
1587                    i += 1;
1588                    continue;
1589                }
1590                let idx1 = u32::from(curr.arg);
1591                let idx2 = u32::from(next.arg);
1592                if idx1 < 16 && idx2 < 16 {
1593                    let packed = (idx1 << 4) | idx2;
1594                    block.instructions[i].instr = Instruction::StoreFastLoadFast {
1595                        var_nums: Arg::marker(),
1596                    }
1597                    .into();
1598                    block.instructions[i].arg = OpArg::new(packed);
1599                    // Replace second instruction with NOP (CPython: INSTR_SET_OP0(inst2, NOP))
1600                    block.instructions[i + 1].instr = Instruction::Nop.into();
1601                    block.instructions[i + 1].arg = OpArg::new(0);
1602                    i += 2; // skip the NOP
1603                } else {
1604                    i += 1;
1605                }
1606            }
1607        }
1608    }
1609
1610    fn optimize_load_fast_borrow(&mut self) {
1611        // NOT_LOCAL marker: instruction didn't come from a LOAD_FAST
1612        const NOT_LOCAL: usize = usize::MAX;
1613
1614        for block in &mut self.blocks {
1615            if block.instructions.is_empty() {
1616                continue;
1617            }
1618
1619            // Track which instructions' outputs are still on stack at block end
1620            // For each instruction, we track if its pushed value(s) are unconsumed
1621            let mut unconsumed = vec![false; block.instructions.len()];
1622
1623            // Simulate stack: each entry is the instruction index that pushed it
1624            // (or NOT_LOCAL if not from LOAD_FAST/LOAD_FAST_LOAD_FAST).
1625            //
1626            // CPython (flowgraph.c optimize_load_fast) pre-fills the stack with
1627            // dummy refs for values inherited from predecessor blocks. We take
1628            // the simpler approach of aborting the optimisation for the whole
1629            // block on stack underflow.
1630            let mut stack: Vec<usize> = Vec::new();
1631            let mut underflow = false;
1632
1633            for (i, info) in block.instructions.iter().enumerate() {
1634                let Some(instr) = info.instr.real() else {
1635                    continue;
1636                };
1637
1638                let stack_effect_info = instr.stack_effect_info(info.arg.into());
1639                let (pushes, pops) = (stack_effect_info.pushed(), stack_effect_info.popped());
1640
1641                // Pop values from stack
1642                for _ in 0..pops {
1643                    if stack.pop().is_none() {
1644                        // Stack underflow — block receives values from a predecessor.
1645                        // Abort optimisation for the entire block.
1646                        underflow = true;
1647                        break;
1648                    }
1649                }
1650                if underflow {
1651                    break;
1652                }
1653
1654                // Push values to stack with source instruction index
1655                let source = match instr {
1656                    Instruction::LoadFast { .. } | Instruction::LoadFastLoadFast { .. } => i,
1657                    _ => NOT_LOCAL,
1658                };
1659                for _ in 0..pushes {
1660                    stack.push(source);
1661                }
1662            }
1663
1664            if underflow {
1665                continue;
1666            }
1667
1668            // Mark instructions whose values remain on stack at block end
1669            for &src in &stack {
1670                if src != NOT_LOCAL {
1671                    unconsumed[src] = true;
1672                }
1673            }
1674
1675            // Convert LOAD_FAST to LOAD_FAST_BORROW where value is fully consumed
1676            for (i, info) in block.instructions.iter_mut().enumerate() {
1677                if unconsumed[i] {
1678                    continue;
1679                }
1680                let Some(instr) = info.instr.real() else {
1681                    continue;
1682                };
1683                match instr {
1684                    Instruction::LoadFast { .. } => {
1685                        info.instr = Instruction::LoadFastBorrow {
1686                            var_num: Arg::marker(),
1687                        }
1688                        .into();
1689                    }
1690                    Instruction::LoadFastLoadFast { .. } => {
1691                        info.instr = Instruction::LoadFastBorrowLoadFastBorrow {
1692                            var_nums: Arg::marker(),
1693                        }
1694                        .into();
1695                    }
1696                    _ => {}
1697                }
1698            }
1699        }
1700    }
1701
1702    fn max_stackdepth(&mut self) -> crate::InternalResult<u32> {
1703        let mut maxdepth = 0u32;
1704        let mut stack = Vec::with_capacity(self.blocks.len());
1705        let mut start_depths = vec![u32::MAX; self.blocks.len()];
1706        stackdepth_push(&mut stack, &mut start_depths, BlockIdx(0), 0);
1707        const DEBUG: bool = false;
1708        'process_blocks: while let Some(block_idx) = stack.pop() {
1709            let idx = block_idx.idx();
1710            let mut depth = start_depths[idx];
1711            if DEBUG {
1712                eprintln!("===BLOCK {}===", block_idx.0);
1713            }
1714            let block = &self.blocks[block_idx];
1715            for ins in &block.instructions {
1716                let instr = &ins.instr;
1717                let effect = instr.stack_effect(ins.arg.into());
1718                if DEBUG {
1719                    let display_arg = if ins.target == BlockIdx::NULL {
1720                        ins.arg
1721                    } else {
1722                        OpArg::new(ins.target.0)
1723                    };
1724                    let instr_display = instr.display(display_arg, self);
1725                    eprint!("{instr_display}: {depth} {effect:+} => ");
1726                }
1727                let new_depth = depth.checked_add_signed(effect).ok_or({
1728                    if effect < 0 {
1729                        InternalError::StackUnderflow
1730                    } else {
1731                        InternalError::StackOverflow
1732                    }
1733                })?;
1734                if DEBUG {
1735                    eprintln!("{new_depth}");
1736                }
1737                if new_depth > maxdepth {
1738                    maxdepth = new_depth
1739                }
1740                // Process target blocks for branching instructions
1741                if ins.target != BlockIdx::NULL {
1742                    if instr.is_block_push() {
1743                        // SETUP_* pseudo ops: target is a handler block.
1744                        // Handler entry depth uses the jump-path stack effect:
1745                        //   SETUP_FINALLY:  +1  (pushes exc)
1746                        //   SETUP_CLEANUP:  +2  (pushes lasti + exc)
1747                        //   SETUP_WITH:     +1  (pops __enter__ result, pushes lasti + exc)
1748                        let handler_effect: u32 = match instr.pseudo() {
1749                            Some(PseudoInstruction::SetupCleanup { .. }) => 2,
1750                            _ => 1, // SetupFinally and SetupWith
1751                        };
1752                        let handler_depth = depth + handler_effect;
1753                        if handler_depth > maxdepth {
1754                            maxdepth = handler_depth;
1755                        }
1756                        stackdepth_push(&mut stack, &mut start_depths, ins.target, handler_depth);
1757                    } else {
1758                        // SEND jumps to END_SEND with receiver still on stack.
1759                        // END_SEND performs the receiver pop.
1760                        let jump_effect = match instr.real() {
1761                            Some(Instruction::Send { .. }) => 0i32,
1762                            _ => effect,
1763                        };
1764                        let target_depth = depth.checked_add_signed(jump_effect).ok_or({
1765                            if jump_effect < 0 {
1766                                InternalError::StackUnderflow
1767                            } else {
1768                                InternalError::StackOverflow
1769                            }
1770                        })?;
1771                        if target_depth > maxdepth {
1772                            maxdepth = target_depth
1773                        }
1774                        stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth);
1775                    }
1776                }
1777                depth = new_depth;
1778                if instr.is_scope_exit() || instr.is_unconditional_jump() {
1779                    continue 'process_blocks;
1780                }
1781            }
1782            // Only push next block if it's not NULL
1783            if block.next != BlockIdx::NULL {
1784                stackdepth_push(&mut stack, &mut start_depths, block.next, depth);
1785            }
1786        }
1787        if DEBUG {
1788            eprintln!("DONE: {maxdepth}");
1789        }
1790
1791        for (block, &start_depth) in self.blocks.iter_mut().zip(&start_depths) {
1792            block.start_depth = (start_depth != u32::MAX).then_some(start_depth);
1793        }
1794
1795        // Fix up handler stack_depth in ExceptHandlerInfo using start_depths
1796        // computed above: depth = start_depth - 1 - preserve_lasti
1797        for block in self.blocks.iter_mut() {
1798            for ins in &mut block.instructions {
1799                if let Some(ref mut handler) = ins.except_handler {
1800                    let h_start = start_depths[handler.handler_block.idx()];
1801                    if h_start != u32::MAX {
1802                        let adjustment = 1 + handler.preserve_lasti as u32;
1803                        debug_assert!(
1804                            h_start >= adjustment,
1805                            "handler start depth {h_start} too shallow for adjustment {adjustment}"
1806                        );
1807                        handler.stack_depth = h_start.saturating_sub(adjustment);
1808                    }
1809                }
1810            }
1811        }
1812
1813        Ok(maxdepth)
1814    }
1815}
1816
1817impl InstrDisplayContext for CodeInfo {
1818    type Constant = ConstantData;
1819
1820    fn get_constant(&self, consti: oparg::ConstIdx) -> &ConstantData {
1821        &self.metadata.consts[consti.as_usize()]
1822    }
1823
1824    fn get_name(&self, i: usize) -> &str {
1825        self.metadata.names[i].as_ref()
1826    }
1827
1828    fn get_varname(&self, var_num: oparg::VarNum) -> &str {
1829        self.metadata.varnames[var_num.as_usize()].as_ref()
1830    }
1831
1832    fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str {
1833        let idx = var_num.as_usize();
1834        let nlocals = self.metadata.varnames.len();
1835        if idx < nlocals {
1836            self.metadata.varnames[idx].as_ref()
1837        } else {
1838            let cell_idx = idx - nlocals;
1839            self.metadata
1840                .cellvars
1841                .get_index(cell_idx)
1842                .unwrap_or_else(|| &self.metadata.freevars[cell_idx - self.metadata.cellvars.len()])
1843                .as_ref()
1844        }
1845    }
1846}
1847
1848fn stackdepth_push(
1849    stack: &mut Vec<BlockIdx>,
1850    start_depths: &mut [u32],
1851    target: BlockIdx,
1852    depth: u32,
1853) {
1854    let idx = target.idx();
1855    let block_depth = &mut start_depths[idx];
1856    if depth > *block_depth || *block_depth == u32::MAX {
1857        *block_depth = depth;
1858        stack.push(target);
1859    }
1860}
1861
1862fn iter_blocks(blocks: &[Block]) -> impl Iterator<Item = (BlockIdx, &Block)> + '_ {
1863    let mut next = BlockIdx(0);
1864    core::iter::from_fn(move || {
1865        if next == BlockIdx::NULL {
1866            return None;
1867        }
1868        let (idx, b) = (next, &blocks[next]);
1869        next = b.next;
1870        Some((idx, b))
1871    })
1872}
1873
1874/// Generate Python 3.11+ format linetable from source locations
1875fn generate_linetable(
1876    locations: &[LineTableLocation],
1877    first_line: i32,
1878    debug_ranges: bool,
1879) -> Box<[u8]> {
1880    if locations.is_empty() {
1881        return Box::new([]);
1882    }
1883
1884    let mut linetable = Vec::new();
1885    // Initialize prev_line to first_line
1886    // The first entry's delta is relative to co_firstlineno
1887    let mut prev_line = first_line;
1888    let mut i = 0;
1889
1890    while i < locations.len() {
1891        let loc = &locations[i];
1892
1893        // Count consecutive instructions with the same location
1894        let mut length = 1;
1895        while i + length < locations.len() && locations[i + length] == locations[i] {
1896            length += 1;
1897        }
1898
1899        // Process in chunks of up to 8 instructions
1900        while length > 0 {
1901            let entry_length = length.min(8);
1902
1903            // Get line information
1904            let line = loc.line;
1905
1906            // NO_LOCATION: emit PyCodeLocationInfoKind::None entries (CACHE, etc.)
1907            if line == -1 {
1908                linetable.push(
1909                    0x80 | ((PyCodeLocationInfoKind::None as u8) << 3) | ((entry_length - 1) as u8),
1910                );
1911                // Do NOT update prev_line
1912                length -= entry_length;
1913                i += entry_length;
1914                continue;
1915            }
1916
1917            let end_line = loc.end_line;
1918            let line_delta = line - prev_line;
1919            let end_line_delta = end_line - line;
1920
1921            // When debug_ranges is disabled, only emit line info (NoColumns format)
1922            if !debug_ranges {
1923                // NoColumns format (code 13): line info only, no column data
1924                linetable.push(
1925                    0x80 | ((PyCodeLocationInfoKind::NoColumns as u8) << 3)
1926                        | ((entry_length - 1) as u8),
1927                );
1928                write_signed_varint(&mut linetable, line_delta);
1929
1930                prev_line = line;
1931                length -= entry_length;
1932                i += entry_length;
1933                continue;
1934            }
1935
1936            // Get column information (only when debug_ranges is enabled)
1937            let col = loc.col;
1938            let end_col = loc.end_col;
1939
1940            // Choose the appropriate encoding based on line delta and column info
1941            if line_delta == 0 && end_line_delta == 0 {
1942                if col < 80 && end_col - col < 16 && end_col >= col {
1943                    // Short form (codes 0-9) for common cases
1944                    let code = (col / 8).min(9) as u8; // Short0 to Short9
1945                    linetable.push(0x80 | (code << 3) | ((entry_length - 1) as u8));
1946                    let col_byte = (((col % 8) as u8) << 4) | ((end_col - col) as u8 & 0xf);
1947                    linetable.push(col_byte);
1948                } else if col < 128 && end_col < 128 {
1949                    // One-line form (code 10) for same line
1950                    linetable.push(
1951                        0x80 | ((PyCodeLocationInfoKind::OneLine0 as u8) << 3)
1952                            | ((entry_length - 1) as u8),
1953                    );
1954                    linetable.push(col as u8);
1955                    linetable.push(end_col as u8);
1956                } else {
1957                    // Long form for columns >= 128
1958                    linetable.push(
1959                        0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3)
1960                            | ((entry_length - 1) as u8),
1961                    );
1962                    write_signed_varint(&mut linetable, 0); // line_delta = 0
1963                    write_varint(&mut linetable, 0); // end_line delta = 0
1964                    write_varint(&mut linetable, (col as u32) + 1);
1965                    write_varint(&mut linetable, (end_col as u32) + 1);
1966                }
1967            } else if line_delta > 0 && line_delta < 3 && end_line_delta == 0 {
1968                // One-line form (codes 11-12) for line deltas 1-2
1969                if col < 128 && end_col < 128 {
1970                    let code = (PyCodeLocationInfoKind::OneLine0 as u8) + (line_delta as u8);
1971                    linetable.push(0x80 | (code << 3) | ((entry_length - 1) as u8));
1972                    linetable.push(col as u8);
1973                    linetable.push(end_col as u8);
1974                } else {
1975                    // Long form for columns >= 128
1976                    linetable.push(
1977                        0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3)
1978                            | ((entry_length - 1) as u8),
1979                    );
1980                    write_signed_varint(&mut linetable, line_delta);
1981                    write_varint(&mut linetable, 0); // end_line delta = 0
1982                    write_varint(&mut linetable, (col as u32) + 1);
1983                    write_varint(&mut linetable, (end_col as u32) + 1);
1984                }
1985            } else {
1986                // Long form (code 14) for all other cases
1987                // Handles: line_delta < 0, line_delta >= 3, multi-line spans, or columns >= 128
1988                linetable.push(
1989                    0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3) | ((entry_length - 1) as u8),
1990                );
1991                write_signed_varint(&mut linetable, line_delta);
1992                write_varint(&mut linetable, end_line_delta as u32);
1993                write_varint(&mut linetable, (col as u32) + 1);
1994                write_varint(&mut linetable, (end_col as u32) + 1);
1995            }
1996
1997            prev_line = line;
1998            length -= entry_length;
1999            i += entry_length;
2000        }
2001    }
2002
2003    linetable.into_boxed_slice()
2004}
2005
2006/// Generate Python 3.11+ exception table from instruction handler info
2007fn generate_exception_table(blocks: &[Block], block_to_index: &[u32]) -> Box<[u8]> {
2008    let mut entries: Vec<ExceptionTableEntry> = Vec::new();
2009    let mut current_entry: Option<(ExceptHandlerInfo, u32)> = None; // (handler_info, start_index)
2010    let mut instr_index = 0u32;
2011
2012    // Iterate through all instructions in block order
2013    // instr_index is the index into the final instructions array (including EXTENDED_ARG)
2014    // This matches how frame.rs uses lasti
2015    for (_, block) in iter_blocks(blocks) {
2016        for instr in &block.instructions {
2017            // instr_size includes EXTENDED_ARG and CACHE entries
2018            let instr_size = instr.arg.instr_size() as u32 + instr.cache_entries;
2019
2020            match (&current_entry, instr.except_handler) {
2021                // No current entry, no handler - nothing to do
2022                (None, None) => {}
2023
2024                // No current entry, handler starts - begin new entry
2025                (None, Some(handler)) => {
2026                    current_entry = Some((handler, instr_index));
2027                }
2028
2029                // Current entry exists, same handler - continue
2030                (Some((curr_handler, _)), Some(handler))
2031                    if curr_handler.handler_block == handler.handler_block
2032                        && curr_handler.stack_depth == handler.stack_depth
2033                        && curr_handler.preserve_lasti == handler.preserve_lasti => {}
2034
2035                // Current entry exists, different handler - finish current, start new
2036                (Some((curr_handler, start)), Some(handler)) => {
2037                    let target_index = block_to_index[curr_handler.handler_block.idx()];
2038                    entries.push(ExceptionTableEntry::new(
2039                        *start,
2040                        instr_index,
2041                        target_index,
2042                        curr_handler.stack_depth as u16,
2043                        curr_handler.preserve_lasti,
2044                    ));
2045                    current_entry = Some((handler, instr_index));
2046                }
2047
2048                // Current entry exists, no handler - finish current entry
2049                (Some((curr_handler, start)), None) => {
2050                    let target_index = block_to_index[curr_handler.handler_block.idx()];
2051                    entries.push(ExceptionTableEntry::new(
2052                        *start,
2053                        instr_index,
2054                        target_index,
2055                        curr_handler.stack_depth as u16,
2056                        curr_handler.preserve_lasti,
2057                    ));
2058                    current_entry = None;
2059                }
2060            }
2061
2062            instr_index += instr_size; // Account for EXTENDED_ARG instructions
2063        }
2064    }
2065
2066    // Finish any remaining entry
2067    if let Some((curr_handler, start)) = current_entry {
2068        let target_index = block_to_index[curr_handler.handler_block.idx()];
2069        entries.push(ExceptionTableEntry::new(
2070            start,
2071            instr_index,
2072            target_index,
2073            curr_handler.stack_depth as u16,
2074            curr_handler.preserve_lasti,
2075        ));
2076    }
2077
2078    encode_exception_table(&entries)
2079}
2080
2081/// Mark exception handler target blocks.
2082/// flowgraph.c mark_except_handlers
2083pub(crate) fn mark_except_handlers(blocks: &mut [Block]) {
2084    // Reset handler flags
2085    for block in blocks.iter_mut() {
2086        block.except_handler = false;
2087        block.preserve_lasti = false;
2088    }
2089    // Mark target blocks of SETUP_* as except handlers
2090    let targets: Vec<usize> = blocks
2091        .iter()
2092        .flat_map(|b| b.instructions.iter())
2093        .filter(|i| i.instr.is_block_push() && i.target != BlockIdx::NULL)
2094        .map(|i| i.target.idx())
2095        .collect();
2096    for idx in targets {
2097        blocks[idx].except_handler = true;
2098    }
2099}
2100
2101/// flowgraph.c mark_cold
2102fn mark_cold(blocks: &mut [Block]) {
2103    let n = blocks.len();
2104    let mut warm = vec![false; n];
2105    let mut queue = VecDeque::new();
2106
2107    warm[0] = true;
2108    queue.push_back(BlockIdx(0));
2109
2110    while let Some(block_idx) = queue.pop_front() {
2111        let block = &blocks[block_idx.idx()];
2112
2113        let has_fallthrough = block
2114            .instructions
2115            .last()
2116            .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
2117            .unwrap_or(true);
2118        if has_fallthrough && block.next != BlockIdx::NULL {
2119            let next_idx = block.next.idx();
2120            if !blocks[next_idx].except_handler && !warm[next_idx] {
2121                warm[next_idx] = true;
2122                queue.push_back(block.next);
2123            }
2124        }
2125
2126        for instr in &block.instructions {
2127            if instr.target != BlockIdx::NULL {
2128                let target_idx = instr.target.idx();
2129                if !blocks[target_idx].except_handler && !warm[target_idx] {
2130                    warm[target_idx] = true;
2131                    queue.push_back(instr.target);
2132                }
2133            }
2134        }
2135    }
2136
2137    for (i, block) in blocks.iter_mut().enumerate() {
2138        block.cold = !warm[i];
2139    }
2140}
2141
2142/// flowgraph.c push_cold_blocks_to_end
2143fn push_cold_blocks_to_end(blocks: &mut Vec<Block>) {
2144    if blocks.len() <= 1 {
2145        return;
2146    }
2147
2148    mark_cold(blocks);
2149
2150    // If a cold block falls through to a warm block, add an explicit jump
2151    let fixups: Vec<(BlockIdx, BlockIdx)> = iter_blocks(blocks)
2152        .filter(|(_, block)| {
2153            block.cold
2154                && block.next != BlockIdx::NULL
2155                && !blocks[block.next.idx()].cold
2156                && block
2157                    .instructions
2158                    .last()
2159                    .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
2160                    .unwrap_or(true)
2161        })
2162        .map(|(idx, block)| (idx, block.next))
2163        .collect();
2164
2165    for (cold_idx, warm_next) in fixups {
2166        let jump_block_idx = BlockIdx(blocks.len() as u32);
2167        let loc = blocks[cold_idx.idx()]
2168            .instructions
2169            .last()
2170            .map(|i| i.location)
2171            .unwrap_or_default();
2172        let end_loc = blocks[cold_idx.idx()]
2173            .instructions
2174            .last()
2175            .map(|i| i.end_location)
2176            .unwrap_or_default();
2177        let mut jump_block = Block {
2178            cold: true,
2179            ..Block::default()
2180        };
2181        jump_block.instructions.push(InstructionInfo {
2182            instr: PseudoInstruction::JumpNoInterrupt {
2183                delta: Arg::marker(),
2184            }
2185            .into(),
2186            arg: OpArg::new(0),
2187            target: warm_next,
2188            location: loc,
2189            end_location: end_loc,
2190            except_handler: None,
2191            lineno_override: None,
2192            cache_entries: 0,
2193        });
2194        jump_block.next = blocks[cold_idx.idx()].next;
2195        blocks[cold_idx.idx()].next = jump_block_idx;
2196        blocks.push(jump_block);
2197    }
2198
2199    // Extract cold block streaks and append at the end
2200    let mut cold_head: BlockIdx = BlockIdx::NULL;
2201    let mut cold_tail: BlockIdx = BlockIdx::NULL;
2202    let mut current = BlockIdx(0);
2203    assert!(!blocks[0].cold);
2204
2205    while current != BlockIdx::NULL {
2206        let next = blocks[current.idx()].next;
2207        if next == BlockIdx::NULL {
2208            break;
2209        }
2210
2211        if blocks[next.idx()].cold {
2212            let cold_start = next;
2213            let mut cold_end = next;
2214            while blocks[cold_end.idx()].next != BlockIdx::NULL
2215                && blocks[blocks[cold_end.idx()].next.idx()].cold
2216            {
2217                cold_end = blocks[cold_end.idx()].next;
2218            }
2219
2220            let after_cold = blocks[cold_end.idx()].next;
2221            blocks[current.idx()].next = after_cold;
2222            blocks[cold_end.idx()].next = BlockIdx::NULL;
2223
2224            if cold_head == BlockIdx::NULL {
2225                cold_head = cold_start;
2226            } else {
2227                blocks[cold_tail.idx()].next = cold_start;
2228            }
2229            cold_tail = cold_end;
2230        } else {
2231            current = next;
2232        }
2233    }
2234
2235    if cold_head != BlockIdx::NULL {
2236        let mut last = current;
2237        while blocks[last.idx()].next != BlockIdx::NULL {
2238            last = blocks[last.idx()].next;
2239        }
2240        blocks[last.idx()].next = cold_head;
2241    }
2242}
2243
2244/// Split blocks at branch points so each block has at most one branch
2245/// (conditional/unconditional jump) as its last instruction.
2246/// This matches CPython's CFG structure where each basic block has one exit.
2247fn split_blocks_at_jumps(blocks: &mut Vec<Block>) {
2248    let mut bi = 0;
2249    while bi < blocks.len() {
2250        // Find the first jump/branch instruction in the block
2251        let split_at = {
2252            let block = &blocks[bi];
2253            let mut found = None;
2254            for (i, ins) in block.instructions.iter().enumerate() {
2255                if is_conditional_jump(&ins.instr)
2256                    || ins.instr.is_unconditional_jump()
2257                    || ins.instr.is_scope_exit()
2258                {
2259                    if i + 1 < block.instructions.len() {
2260                        found = Some(i + 1);
2261                    }
2262                    break;
2263                }
2264            }
2265            found
2266        };
2267        if let Some(pos) = split_at {
2268            let new_block_idx = BlockIdx(blocks.len() as u32);
2269            let tail: Vec<InstructionInfo> = blocks[bi].instructions.drain(pos..).collect();
2270            let old_next = blocks[bi].next;
2271            let cold = blocks[bi].cold;
2272            blocks[bi].next = new_block_idx;
2273            blocks.push(Block {
2274                instructions: tail,
2275                next: old_next,
2276                cold,
2277                ..Block::default()
2278            });
2279            // Don't increment bi - re-check current block (it might still have issues)
2280        } else {
2281            bi += 1;
2282        }
2283    }
2284}
2285
2286/// Jump threading: when a block's last jump targets a block whose first
2287/// instruction is an unconditional jump, redirect to the final target.
2288/// flowgraph.c optimize_basic_block + jump_thread
2289fn jump_threading(blocks: &mut [Block]) {
2290    let mut changed = true;
2291    while changed {
2292        changed = false;
2293        for bi in 0..blocks.len() {
2294            let last_idx = match blocks[bi].instructions.len().checked_sub(1) {
2295                Some(i) => i,
2296                None => continue,
2297            };
2298            let ins = &blocks[bi].instructions[last_idx];
2299            let target = ins.target;
2300            if target == BlockIdx::NULL {
2301                continue;
2302            }
2303            if !ins.instr.is_unconditional_jump() && !is_conditional_jump(&ins.instr) {
2304                continue;
2305            }
2306            // Check if target block's first instruction is an unconditional jump
2307            let target_block = &blocks[target.idx()];
2308            if let Some(target_ins) = target_block.instructions.first()
2309                && target_ins.instr.is_unconditional_jump()
2310                && target_ins.target != BlockIdx::NULL
2311                && target_ins.target != target
2312            {
2313                let final_target = target_ins.target;
2314                blocks[bi].instructions[last_idx].target = final_target;
2315                changed = true;
2316            }
2317        }
2318    }
2319}
2320
2321fn is_conditional_jump(instr: &AnyInstruction) -> bool {
2322    matches!(
2323        instr.real(),
2324        Some(
2325            Instruction::PopJumpIfFalse { .. }
2326                | Instruction::PopJumpIfTrue { .. }
2327                | Instruction::PopJumpIfNone { .. }
2328                | Instruction::PopJumpIfNotNone { .. }
2329        )
2330    )
2331}
2332
2333/// Invert a conditional jump opcode.
2334fn reversed_conditional(instr: &AnyInstruction) -> Option<AnyInstruction> {
2335    Some(match instr.real()? {
2336        Instruction::PopJumpIfFalse { .. } => Instruction::PopJumpIfTrue {
2337            delta: Arg::marker(),
2338        }
2339        .into(),
2340        Instruction::PopJumpIfTrue { .. } => Instruction::PopJumpIfFalse {
2341            delta: Arg::marker(),
2342        }
2343        .into(),
2344        Instruction::PopJumpIfNone { .. } => Instruction::PopJumpIfNotNone {
2345            delta: Arg::marker(),
2346        }
2347        .into(),
2348        Instruction::PopJumpIfNotNone { .. } => Instruction::PopJumpIfNone {
2349            delta: Arg::marker(),
2350        }
2351        .into(),
2352        _ => return None,
2353    })
2354}
2355
2356/// flowgraph.c normalize_jumps + remove_redundant_jumps
2357fn normalize_jumps(blocks: &mut Vec<Block>) {
2358    let mut visit_order = Vec::new();
2359    let mut visited = vec![false; blocks.len()];
2360    let mut current = BlockIdx(0);
2361    while current != BlockIdx::NULL {
2362        visit_order.push(current);
2363        visited[current.idx()] = true;
2364        current = blocks[current.idx()].next;
2365    }
2366
2367    visited.fill(false);
2368
2369    for &block_idx in &visit_order {
2370        let idx = block_idx.idx();
2371        visited[idx] = true;
2372
2373        // Remove redundant unconditional jump to next block
2374        let next = blocks[idx].next;
2375        if next != BlockIdx::NULL {
2376            let last = blocks[idx].instructions.last();
2377            let is_jump_to_next = last.is_some_and(|ins| {
2378                ins.instr.is_unconditional_jump()
2379                    && ins.target != BlockIdx::NULL
2380                    && ins.target == next
2381            });
2382            if is_jump_to_next && let Some(last_instr) = blocks[idx].instructions.last_mut() {
2383                last_instr.instr = Instruction::Nop.into();
2384                last_instr.target = BlockIdx::NULL;
2385            }
2386        }
2387
2388        // Normalize conditional jumps: forward gets NOT_TAKEN, backward gets inverted
2389        let last = blocks[idx].instructions.last();
2390        if let Some(last_ins) = last
2391            && is_conditional_jump(&last_ins.instr)
2392            && last_ins.target != BlockIdx::NULL
2393        {
2394            let target = last_ins.target;
2395            let is_forward = !visited[target.idx()];
2396
2397            if is_forward {
2398                // Insert NOT_TAKEN after forward conditional jump
2399                let not_taken = InstructionInfo {
2400                    instr: Instruction::NotTaken.into(),
2401                    arg: OpArg::new(0),
2402                    target: BlockIdx::NULL,
2403                    location: last_ins.location,
2404                    end_location: last_ins.end_location,
2405                    except_handler: last_ins.except_handler,
2406                    lineno_override: None,
2407                    cache_entries: 0,
2408                };
2409                blocks[idx].instructions.push(not_taken);
2410            } else {
2411                // Backward conditional jump: invert and create new block
2412                // Transform: `cond_jump T` (backward)
2413                // Into: `reversed_cond_jump b_next` + new block [NOT_TAKEN, JUMP T]
2414                let loc = last_ins.location;
2415                let end_loc = last_ins.end_location;
2416                let exc_handler = last_ins.except_handler;
2417
2418                if let Some(reversed) = reversed_conditional(&last_ins.instr) {
2419                    let old_next = blocks[idx].next;
2420                    let is_cold = blocks[idx].cold;
2421
2422                    // Create new block with NOT_TAKEN + JUMP to original backward target
2423                    let new_block_idx = BlockIdx(blocks.len() as u32);
2424                    let mut new_block = Block {
2425                        cold: is_cold,
2426                        ..Block::default()
2427                    };
2428                    new_block.instructions.push(InstructionInfo {
2429                        instr: Instruction::NotTaken.into(),
2430                        arg: OpArg::new(0),
2431                        target: BlockIdx::NULL,
2432                        location: loc,
2433                        end_location: end_loc,
2434                        except_handler: exc_handler,
2435                        lineno_override: None,
2436                        cache_entries: 0,
2437                    });
2438                    new_block.instructions.push(InstructionInfo {
2439                        instr: PseudoInstruction::Jump {
2440                            delta: Arg::marker(),
2441                        }
2442                        .into(),
2443                        arg: OpArg::new(0),
2444                        target,
2445                        location: loc,
2446                        end_location: end_loc,
2447                        except_handler: exc_handler,
2448                        lineno_override: None,
2449                        cache_entries: 0,
2450                    });
2451                    new_block.next = old_next;
2452
2453                    // Update the conditional jump: invert opcode, target = old next block
2454                    let last_mut = blocks[idx].instructions.last_mut().unwrap();
2455                    last_mut.instr = reversed;
2456                    last_mut.target = old_next;
2457
2458                    // Splice new block between current and old next
2459                    blocks[idx].next = new_block_idx;
2460                    blocks.push(new_block);
2461
2462                    // Extend visited array and update visit order
2463                    visited.push(true);
2464                }
2465            }
2466        }
2467    }
2468
2469    // Rebuild visit_order since backward normalization may have added new blocks
2470    let mut visit_order = Vec::new();
2471    let mut current = BlockIdx(0);
2472    while current != BlockIdx::NULL {
2473        visit_order.push(current);
2474        current = blocks[current.idx()].next;
2475    }
2476
2477    // Replace JUMP → value-producing-instr + RETURN_VALUE with inline return.
2478    // This matches CPython's optimize_basic_block: "Replace JUMP to a RETURN".
2479    for &block_idx in &visit_order {
2480        let idx = block_idx.idx();
2481        let mut replacements: Vec<(usize, Vec<InstructionInfo>)> = Vec::new();
2482        for (i, ins) in blocks[idx].instructions.iter().enumerate() {
2483            if !ins.instr.is_unconditional_jump() || ins.target == BlockIdx::NULL {
2484                continue;
2485            }
2486            // Follow through empty blocks (next_nonempty_block)
2487            let mut target_idx = ins.target.idx();
2488            while blocks[target_idx].instructions.is_empty()
2489                && blocks[target_idx].next != BlockIdx::NULL
2490            {
2491                target_idx = blocks[target_idx].next.idx();
2492            }
2493            let target_block = &blocks[target_idx];
2494            // Target must be exactly `value; RETURN_VALUE`.
2495            if target_block.instructions.len() == 2 {
2496                let t0 = &target_block.instructions[0];
2497                let t1 = &target_block.instructions[1];
2498                if matches!(t0.instr, AnyInstruction::Real(_))
2499                    && !t0.instr.is_scope_exit()
2500                    && !t0.instr.is_unconditional_jump()
2501                    && matches!(t1.instr.real(), Some(Instruction::ReturnValue))
2502                {
2503                    let mut load = *t0;
2504                    let mut ret = *t1;
2505                    // Use the jump's location for the inlined return
2506                    load.location = ins.location;
2507                    load.end_location = ins.end_location;
2508                    load.except_handler = ins.except_handler;
2509                    ret.location = ins.location;
2510                    ret.end_location = ins.end_location;
2511                    ret.except_handler = ins.except_handler;
2512                    replacements.push((i, vec![load, ret]));
2513                }
2514            }
2515        }
2516        // Apply replacements in reverse order
2517        for (i, new_insts) in replacements.into_iter().rev() {
2518            blocks[idx].instructions.splice(i..i + 1, new_insts);
2519        }
2520    }
2521
2522    // Resolve JUMP/JUMP_NO_INTERRUPT pseudo instructions before offset fixpoint.
2523    let mut block_order = vec![0u32; blocks.len()];
2524    for (pos, &block_idx) in visit_order.iter().enumerate() {
2525        block_order[block_idx.idx()] = pos as u32;
2526    }
2527
2528    for &block_idx in &visit_order {
2529        let source_pos = block_order[block_idx.idx()];
2530        for info in &mut blocks[block_idx.idx()].instructions {
2531            let target = info.target;
2532            if target == BlockIdx::NULL {
2533                continue;
2534            }
2535            let target_pos = block_order[target.idx()];
2536            info.instr = match info.instr {
2537                AnyInstruction::Pseudo(PseudoInstruction::Jump { .. }) => {
2538                    if target_pos > source_pos {
2539                        Instruction::JumpForward {
2540                            delta: Arg::marker(),
2541                        }
2542                        .into()
2543                    } else {
2544                        Instruction::JumpBackward {
2545                            delta: Arg::marker(),
2546                        }
2547                        .into()
2548                    }
2549                }
2550                AnyInstruction::Pseudo(PseudoInstruction::JumpNoInterrupt { .. }) => {
2551                    if target_pos > source_pos {
2552                        Instruction::JumpForward {
2553                            delta: Arg::marker(),
2554                        }
2555                        .into()
2556                    } else {
2557                        Instruction::JumpBackwardNoInterrupt {
2558                            delta: Arg::marker(),
2559                        }
2560                        .into()
2561                    }
2562                }
2563                other => other,
2564            };
2565        }
2566    }
2567}
2568
2569/// flowgraph.c inline_small_or_no_lineno_blocks
2570fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) {
2571    const MAX_COPY_SIZE: usize = 4;
2572
2573    let block_exits_scope = |block: &Block| {
2574        block
2575            .instructions
2576            .last()
2577            .is_some_and(|ins| ins.instr.is_scope_exit())
2578    };
2579    let block_has_no_lineno = |block: &Block| {
2580        block
2581            .instructions
2582            .iter()
2583            .all(|ins| !instruction_has_lineno(ins))
2584    };
2585
2586    loop {
2587        let mut changes = false;
2588        let mut current = BlockIdx(0);
2589        while current != BlockIdx::NULL {
2590            let next = blocks[current.idx()].next;
2591            let Some(last) = blocks[current.idx()].instructions.last().copied() else {
2592                current = next;
2593                continue;
2594            };
2595            if !last.instr.is_unconditional_jump() || last.target == BlockIdx::NULL {
2596                current = next;
2597                continue;
2598            }
2599
2600            let target = last.target;
2601            let small_exit_block = block_exits_scope(&blocks[target.idx()])
2602                && blocks[target.idx()].instructions.len() <= MAX_COPY_SIZE;
2603            let no_lineno_no_fallthrough = block_has_no_lineno(&blocks[target.idx()])
2604                && !block_has_fallthrough(&blocks[target.idx()]);
2605
2606            if small_exit_block || no_lineno_no_fallthrough {
2607                if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() {
2608                    last_instr.instr = Instruction::Nop.into();
2609                    last_instr.arg = OpArg::new(0);
2610                    last_instr.target = BlockIdx::NULL;
2611                }
2612                let appended = blocks[target.idx()].instructions.clone();
2613                blocks[current.idx()].instructions.extend(appended);
2614                changes = true;
2615            }
2616
2617            current = next;
2618        }
2619
2620        if !changes {
2621            break;
2622        }
2623    }
2624}
2625
2626/// Follow chain of empty blocks to find first non-empty block.
2627fn next_nonempty_block(blocks: &[Block], mut idx: BlockIdx) -> BlockIdx {
2628    while idx != BlockIdx::NULL
2629        && blocks[idx.idx()].instructions.is_empty()
2630        && blocks[idx.idx()].next != BlockIdx::NULL
2631    {
2632        idx = blocks[idx.idx()].next;
2633    }
2634    idx
2635}
2636
2637fn instruction_lineno(instr: &InstructionInfo) -> i32 {
2638    instr
2639        .lineno_override
2640        .unwrap_or_else(|| instr.location.line.get() as i32)
2641}
2642
2643fn instruction_has_lineno(instr: &InstructionInfo) -> bool {
2644    instruction_lineno(instr) > 0
2645}
2646
2647fn block_has_fallthrough(block: &Block) -> bool {
2648    block
2649        .instructions
2650        .last()
2651        .is_none_or(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
2652}
2653
2654fn is_jump_instruction(instr: &InstructionInfo) -> bool {
2655    instr.instr.is_unconditional_jump() || is_conditional_jump(&instr.instr)
2656}
2657
2658fn is_exit_without_lineno(block: &Block) -> bool {
2659    let Some(first) = block.instructions.first() else {
2660        return false;
2661    };
2662    let Some(last) = block.instructions.last() else {
2663        return false;
2664    };
2665    !instruction_has_lineno(first) && last.instr.is_scope_exit()
2666}
2667
2668fn maybe_propagate_location(
2669    instr: &mut InstructionInfo,
2670    location: SourceLocation,
2671    end_location: SourceLocation,
2672) {
2673    if !instruction_has_lineno(instr) {
2674        instr.location = location;
2675        instr.end_location = end_location;
2676        instr.lineno_override = None;
2677    }
2678}
2679
2680fn propagate_locations_in_block(
2681    block: &mut Block,
2682    location: SourceLocation,
2683    end_location: SourceLocation,
2684) {
2685    let mut prev_location = location;
2686    let mut prev_end_location = end_location;
2687    for instr in &mut block.instructions {
2688        maybe_propagate_location(instr, prev_location, prev_end_location);
2689        prev_location = instr.location;
2690        prev_end_location = instr.end_location;
2691    }
2692}
2693
2694fn compute_predecessors(blocks: &[Block]) -> Vec<u32> {
2695    let mut predecessors = vec![0u32; blocks.len()];
2696    if blocks.is_empty() {
2697        return predecessors;
2698    }
2699
2700    predecessors[0] = 1;
2701    let mut current = BlockIdx(0);
2702    while current != BlockIdx::NULL {
2703        let block = &blocks[current.idx()];
2704        if block_has_fallthrough(block) {
2705            let next = next_nonempty_block(blocks, block.next);
2706            if next != BlockIdx::NULL {
2707                predecessors[next.idx()] += 1;
2708            }
2709        }
2710        for ins in &block.instructions {
2711            if ins.target != BlockIdx::NULL {
2712                let target = next_nonempty_block(blocks, ins.target);
2713                if target != BlockIdx::NULL {
2714                    predecessors[target.idx()] += 1;
2715                }
2716            }
2717        }
2718        current = block.next;
2719    }
2720    predecessors
2721}
2722
2723fn duplicate_exits_without_lineno(blocks: &mut Vec<Block>, predecessors: &mut Vec<u32>) {
2724    let mut current = BlockIdx(0);
2725    while current != BlockIdx::NULL {
2726        let block = &blocks[current.idx()];
2727        let last = match block.instructions.last() {
2728            Some(ins) if ins.target != BlockIdx::NULL && is_jump_instruction(ins) => ins,
2729            _ => {
2730                current = blocks[current.idx()].next;
2731                continue;
2732            }
2733        };
2734
2735        let target = next_nonempty_block(blocks, last.target);
2736        if target == BlockIdx::NULL || !is_exit_without_lineno(&blocks[target.idx()]) {
2737            current = blocks[current.idx()].next;
2738            continue;
2739        }
2740        if predecessors[target.idx()] <= 1 {
2741            current = blocks[current.idx()].next;
2742            continue;
2743        }
2744
2745        // Copy the exit block and splice it into the linked list after current
2746        let new_idx = BlockIdx(blocks.len() as u32);
2747        let mut new_block = blocks[target.idx()].clone();
2748        let jump_loc = last.location;
2749        let jump_end_loc = last.end_location;
2750        propagate_locations_in_block(&mut new_block, jump_loc, jump_end_loc);
2751        let old_next = blocks[current.idx()].next;
2752        new_block.next = old_next;
2753        blocks.push(new_block);
2754        blocks[current.idx()].next = new_idx;
2755
2756        // Update the jump target
2757        let last_mut = blocks[current.idx()].instructions.last_mut().unwrap();
2758        last_mut.target = new_idx;
2759        predecessors[target.idx()] -= 1;
2760        predecessors.push(1);
2761
2762        // Skip past the newly inserted block
2763        current = old_next;
2764    }
2765
2766    current = BlockIdx(0);
2767    while current != BlockIdx::NULL {
2768        let block = &blocks[current.idx()];
2769        if let Some(last) = block.instructions.last()
2770            && block_has_fallthrough(block)
2771        {
2772            let target = next_nonempty_block(blocks, block.next);
2773            if target != BlockIdx::NULL
2774                && predecessors[target.idx()] == 1
2775                && is_exit_without_lineno(&blocks[target.idx()])
2776            {
2777                let last_location = last.location;
2778                let last_end_location = last.end_location;
2779                propagate_locations_in_block(
2780                    &mut blocks[target.idx()],
2781                    last_location,
2782                    last_end_location,
2783                );
2784            }
2785        }
2786        current = blocks[current.idx()].next;
2787    }
2788}
2789
2790fn propagate_line_numbers(blocks: &mut [Block], predecessors: &[u32]) {
2791    let mut current = BlockIdx(0);
2792    while current != BlockIdx::NULL {
2793        let last = blocks[current.idx()].instructions.last().copied();
2794        if let Some(last) = last {
2795            let (next_block, has_fallthrough) = {
2796                let block = &blocks[current.idx()];
2797                (block.next, block_has_fallthrough(block))
2798            };
2799
2800            {
2801                let block = &mut blocks[current.idx()];
2802                let mut prev_location = None;
2803                for instr in &mut block.instructions {
2804                    if let Some((location, end_location)) = prev_location {
2805                        maybe_propagate_location(instr, location, end_location);
2806                    }
2807                    prev_location = Some((instr.location, instr.end_location));
2808                }
2809            }
2810
2811            if has_fallthrough {
2812                let target = next_nonempty_block(blocks, next_block);
2813                if target != BlockIdx::NULL && predecessors[target.idx()] == 1 {
2814                    propagate_locations_in_block(
2815                        &mut blocks[target.idx()],
2816                        last.location,
2817                        last.end_location,
2818                    );
2819                }
2820            }
2821
2822            if is_jump_instruction(&last) {
2823                let target = next_nonempty_block(blocks, last.target);
2824                if target != BlockIdx::NULL && predecessors[target.idx()] == 1 {
2825                    propagate_locations_in_block(
2826                        &mut blocks[target.idx()],
2827                        last.location,
2828                        last.end_location,
2829                    );
2830                }
2831            }
2832        }
2833        current = blocks[current.idx()].next;
2834    }
2835}
2836
2837fn resolve_line_numbers(blocks: &mut Vec<Block>) {
2838    let mut predecessors = compute_predecessors(blocks);
2839    duplicate_exits_without_lineno(blocks, &mut predecessors);
2840    propagate_line_numbers(blocks, &predecessors);
2841}
2842
2843/// Duplicate `LOAD_CONST None + RETURN_VALUE` for blocks that fall through
2844/// to the final return block.
2845fn duplicate_end_returns(blocks: &mut [Block]) {
2846    // Walk the block chain and keep the last non-empty block.
2847    let mut last_block = BlockIdx::NULL;
2848    let mut current = BlockIdx(0);
2849    while current != BlockIdx::NULL {
2850        if !blocks[current.idx()].instructions.is_empty() {
2851            last_block = current;
2852        }
2853        current = blocks[current.idx()].next;
2854    }
2855    if last_block == BlockIdx::NULL {
2856        return;
2857    }
2858
2859    let last_insts = &blocks[last_block.idx()].instructions;
2860    // Only apply when the last block is EXACTLY a return-None epilogue
2861    let is_return_block = last_insts.len() == 2
2862        && matches!(
2863            last_insts[0].instr,
2864            AnyInstruction::Real(Instruction::LoadConst { .. })
2865        )
2866        && matches!(
2867            last_insts[1].instr,
2868            AnyInstruction::Real(Instruction::ReturnValue)
2869        );
2870    if !is_return_block {
2871        return;
2872    }
2873
2874    // Get the return instructions to clone
2875    let return_insts: Vec<InstructionInfo> = last_insts[last_insts.len() - 2..].to_vec();
2876
2877    // Find non-cold blocks that fall through to the last block
2878    let mut blocks_to_fix = Vec::new();
2879    current = BlockIdx(0);
2880    while current != BlockIdx::NULL {
2881        let block = &blocks[current.idx()];
2882        let next = next_nonempty_block(blocks, block.next);
2883        if current != last_block && next == last_block && !block.cold && !block.except_handler {
2884            let last_ins = block.instructions.last();
2885            let has_fallthrough = last_ins
2886                .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
2887                .unwrap_or(true);
2888            // Don't duplicate if block already ends with the same return pattern
2889            let already_has_return = block.instructions.len() >= 2 && {
2890                let n = block.instructions.len();
2891                matches!(
2892                    block.instructions[n - 2].instr,
2893                    AnyInstruction::Real(Instruction::LoadConst { .. })
2894                ) && matches!(
2895                    block.instructions[n - 1].instr,
2896                    AnyInstruction::Real(Instruction::ReturnValue)
2897                )
2898            };
2899            if has_fallthrough && !already_has_return {
2900                blocks_to_fix.push(current);
2901            }
2902        }
2903        current = blocks[current.idx()].next;
2904    }
2905
2906    // Duplicate the return instructions at the end of fall-through blocks
2907    for block_idx in blocks_to_fix {
2908        blocks[block_idx.idx()]
2909            .instructions
2910            .extend_from_slice(&return_insts);
2911    }
2912}
2913
2914/// Label exception targets: walk CFG with except stack, set per-instruction
2915/// handler info and block preserve_lasti flag. Converts POP_BLOCK to NOP.
2916/// flowgraph.c label_exception_targets + push_except_block
2917pub(crate) fn label_exception_targets(blocks: &mut [Block]) {
2918    #[derive(Clone)]
2919    struct ExceptEntry {
2920        handler_block: BlockIdx,
2921        preserve_lasti: bool,
2922    }
2923
2924    let num_blocks = blocks.len();
2925    if num_blocks == 0 {
2926        return;
2927    }
2928
2929    let mut visited = vec![false; num_blocks];
2930    let mut block_stacks: Vec<Option<Vec<ExceptEntry>>> = vec![None; num_blocks];
2931
2932    // Entry block
2933    visited[0] = true;
2934    block_stacks[0] = Some(Vec::new());
2935
2936    let mut todo = vec![BlockIdx(0)];
2937
2938    while let Some(block_idx) = todo.pop() {
2939        let bi = block_idx.idx();
2940        let mut stack = block_stacks[bi].take().unwrap_or_default();
2941        let mut last_yield_except_depth: i32 = -1;
2942
2943        let instr_count = blocks[bi].instructions.len();
2944        for i in 0..instr_count {
2945            // Read all needed fields (each temporary borrow ends immediately)
2946            let target = blocks[bi].instructions[i].target;
2947            let arg = blocks[bi].instructions[i].arg;
2948            let is_push = blocks[bi].instructions[i].instr.is_block_push();
2949            let is_pop = blocks[bi].instructions[i].instr.is_pop_block();
2950
2951            if is_push {
2952                // Determine preserve_lasti from instruction type (push_except_block)
2953                let preserve_lasti = matches!(
2954                    blocks[bi].instructions[i].instr.pseudo(),
2955                    Some(
2956                        PseudoInstruction::SetupWith { .. }
2957                            | PseudoInstruction::SetupCleanup { .. }
2958                    )
2959                );
2960
2961                // Set preserve_lasti on handler block
2962                if preserve_lasti && target != BlockIdx::NULL {
2963                    blocks[target.idx()].preserve_lasti = true;
2964                }
2965
2966                // Propagate except stack to handler block if not visited
2967                if target != BlockIdx::NULL && !visited[target.idx()] {
2968                    visited[target.idx()] = true;
2969                    block_stacks[target.idx()] = Some(stack.clone());
2970                    todo.push(target);
2971                }
2972
2973                // Push handler onto except stack
2974                stack.push(ExceptEntry {
2975                    handler_block: target,
2976                    preserve_lasti,
2977                });
2978            } else if is_pop {
2979                debug_assert!(
2980                    !stack.is_empty(),
2981                    "POP_BLOCK with empty except stack at block {bi} instruction {i}"
2982                );
2983                stack.pop();
2984                // POP_BLOCK → NOP
2985                blocks[bi].instructions[i].instr = Instruction::Nop.into();
2986            } else {
2987                // Set except_handler for this instruction from except stack top
2988                // stack_depth placeholder: filled by fixup_handler_depths
2989                let handler_info = stack.last().map(|e| ExceptHandlerInfo {
2990                    handler_block: e.handler_block,
2991                    stack_depth: 0,
2992                    preserve_lasti: e.preserve_lasti,
2993                });
2994                blocks[bi].instructions[i].except_handler = handler_info;
2995
2996                // Track YIELD_VALUE except stack depth
2997                // Record the except stack depth at the point of yield.
2998                // With the StopIteration wrapper, depth is naturally correct:
2999                // - plain yield outside try: depth=1 → DEPTH1 set
3000                // - yield inside try: depth=2+ → no DEPTH1
3001                // - yield-from/await: has internal SETUP_FINALLY → depth=2+ → no DEPTH1
3002                if let Some(Instruction::YieldValue { .. }) =
3003                    blocks[bi].instructions[i].instr.real()
3004                {
3005                    last_yield_except_depth = stack.len() as i32;
3006                }
3007
3008                // Set RESUME DEPTH1 flag based on last yield's except depth
3009                if let Some(Instruction::Resume { context }) =
3010                    blocks[bi].instructions[i].instr.real()
3011                {
3012                    let location = context.get(arg).location();
3013                    match location {
3014                        oparg::ResumeLocation::AtFuncStart => {}
3015                        _ => {
3016                            if last_yield_except_depth == 1 {
3017                                blocks[bi].instructions[i].arg =
3018                                    OpArg::new(oparg::ResumeContext::new(location, true).as_u32());
3019                            }
3020                            last_yield_except_depth = -1;
3021                        }
3022                    }
3023                }
3024
3025                // For jump instructions, propagate except stack to target
3026                if target != BlockIdx::NULL && !visited[target.idx()] {
3027                    visited[target.idx()] = true;
3028                    block_stacks[target.idx()] = Some(stack.clone());
3029                    todo.push(target);
3030                }
3031            }
3032        }
3033
3034        // Propagate to fallthrough block (block.next)
3035        let next = blocks[bi].next;
3036        if next != BlockIdx::NULL && !visited[next.idx()] {
3037            let has_fallthrough = blocks[bi]
3038                .instructions
3039                .last()
3040                .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump())
3041                .unwrap_or(true); // Empty block falls through
3042            if has_fallthrough {
3043                visited[next.idx()] = true;
3044                block_stacks[next.idx()] = Some(stack);
3045                todo.push(next);
3046            }
3047        }
3048    }
3049}
3050
3051/// Convert remaining pseudo ops to real instructions or NOP.
3052/// flowgraph.c convert_pseudo_ops
3053pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32]) {
3054    for block in blocks.iter_mut() {
3055        for info in &mut block.instructions {
3056            let Some(pseudo) = info.instr.pseudo() else {
3057                continue;
3058            };
3059            match pseudo {
3060                // Block push pseudo ops → NOP
3061                PseudoInstruction::SetupCleanup { .. }
3062                | PseudoInstruction::SetupFinally { .. }
3063                | PseudoInstruction::SetupWith { .. } => {
3064                    info.instr = Instruction::Nop.into();
3065                }
3066                // PopBlock in reachable blocks is converted to NOP by
3067                // label_exception_targets. Dead blocks may still have them.
3068                PseudoInstruction::PopBlock => {
3069                    info.instr = Instruction::Nop.into();
3070                }
3071                // LOAD_CLOSURE → LOAD_FAST (using cellfixedoffsets for merged layout)
3072                PseudoInstruction::LoadClosure { i } => {
3073                    let cell_relative = i.get(info.arg) as usize;
3074                    let new_idx = cellfixedoffsets[cell_relative];
3075                    info.arg = OpArg::new(new_idx);
3076                    info.instr = Instruction::LoadFast {
3077                        var_num: Arg::marker(),
3078                    }
3079                    .into();
3080                }
3081                // Jump pseudo ops are resolved during block linearization
3082                PseudoInstruction::Jump { .. } | PseudoInstruction::JumpNoInterrupt { .. } => {}
3083                // These should have been resolved earlier
3084                PseudoInstruction::AnnotationsPlaceholder
3085                | PseudoInstruction::JumpIfFalse { .. }
3086                | PseudoInstruction::JumpIfTrue { .. }
3087                | PseudoInstruction::StoreFastMaybeNull { .. } => {
3088                    unreachable!("Unexpected pseudo instruction in convert_pseudo_ops: {pseudo:?}")
3089                }
3090            }
3091        }
3092    }
3093}
3094
3095/// Build cellfixedoffsets mapping: cell/free index -> localsplus index.
3096/// Merged cells (cellvar also in varnames) get the local slot index.
3097/// Non-merged cells get slots after nlocals. Free vars follow.
3098pub(crate) fn build_cellfixedoffsets(
3099    varnames: &IndexSet<String>,
3100    cellvars: &IndexSet<String>,
3101    freevars: &IndexSet<String>,
3102) -> Vec<u32> {
3103    let nlocals = varnames.len();
3104    let ncells = cellvars.len();
3105    let nfrees = freevars.len();
3106    let mut fixed = Vec::with_capacity(ncells + nfrees);
3107    let mut numdropped = 0usize;
3108    for (i, cellvar) in cellvars.iter().enumerate() {
3109        if let Some(local_idx) = varnames.get_index_of(cellvar) {
3110            fixed.push(local_idx as u32);
3111            numdropped += 1;
3112        } else {
3113            fixed.push((nlocals + i - numdropped) as u32);
3114        }
3115    }
3116    for i in 0..nfrees {
3117        fixed.push((nlocals + ncells - numdropped + i) as u32);
3118    }
3119    fixed
3120}
3121
3122/// Convert DEREF instruction opargs from cell-relative indices to localsplus indices
3123/// using the cellfixedoffsets mapping.
3124pub(crate) fn fixup_deref_opargs(blocks: &mut [Block], cellfixedoffsets: &[u32]) {
3125    for block in blocks.iter_mut() {
3126        for info in &mut block.instructions {
3127            let Some(instr) = info.instr.real() else {
3128                continue;
3129            };
3130            let needs_fixup = matches!(
3131                instr,
3132                Instruction::LoadDeref { .. }
3133                    | Instruction::StoreDeref { .. }
3134                    | Instruction::DeleteDeref { .. }
3135                    | Instruction::LoadFromDictOrDeref { .. }
3136                    | Instruction::MakeCell { .. }
3137            );
3138            if needs_fixup {
3139                let cell_relative = u32::from(info.arg) as usize;
3140                info.arg = OpArg::new(cellfixedoffsets[cell_relative]);
3141            }
3142        }
3143    }
3144}