diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 51b513c758..62c7178591 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -1513,12 +1513,12 @@ impl Compiler { FBlockType::ForLoop => { // When returning from a for-loop, CPython swaps the preserved - // value with the iterator and uses POP_TOP for the iterator slot. + // value with the iterator and uses POP_TOP for loop cleanup. if preserve_tos { emit!(self, Instruction::Swap { i: 2 }); emit!(self, Instruction::PopTop); } else { - emit!(self, Instruction::PopIter); + emit!(self, Instruction::PopTop); } } @@ -5809,7 +5809,10 @@ impl Compiler { emit!(self, Instruction::PopTop); // pop lasti emit!(self, Instruction::PopTop); // pop self_exit emit!(self, Instruction::PopTop); // pop exit_func - emit!(self, PseudoInstruction::Jump { delta: after_block }); + emit!( + self, + PseudoInstruction::JumpNoInterrupt { delta: after_block } + ); // ===== Cleanup block (for nested exception during __exit__) ===== // Stack: [..., __exit__, lasti, prev_exc, lasti2, exc2] @@ -9920,9 +9923,9 @@ impl Compiler { } } - // For break in a for loop, pop the iterator + // CPython unwinds a for-loop break with POP_TOP rather than POP_ITER. if is_break && is_for_loop { - emit!(self, Instruction::PopIter); + emit!(self, Instruction::PopTop); } // Jump to target @@ -11155,12 +11158,12 @@ def f(base, cls, state): fn test_loop_store_subscr_threads_direct_backedge() { let code = compile_exec( "\ -def f(kwonlyargs, kwonlydefaults, arg2value): +def f(kwonlyargs, kw_only_defaults, arg2value): missing = 0 for kwarg in kwonlyargs: if kwarg not in arg2value: - if kwonlydefaults and kwarg in kwonlydefaults: - arg2value[kwarg] = kwonlydefaults[kwarg] + if kw_only_defaults and kwarg in kw_only_defaults: + arg2value[kwarg] = kw_only_defaults[kwarg] else: missing += 1 return missing @@ -11244,6 +11247,91 @@ def f(obj): ); } + #[test] + fn test_shared_final_return_is_cloned_for_jump_target() { + let code = compile_exec( + "\ +def f(node): + if not isinstance( + node, (AsyncFunctionDef, FunctionDef, ClassDef, Module) + ) or len(node.body) < 1: + return None + node = node.body[0] + if not isinstance(node, Expr): + return None + node = node.value + if isinstance(node, Constant) and isinstance(node.value, str): + return node +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let return_count = ops + .iter() + .filter(|op| matches!(op, Instruction::ReturnValue)) + .count(); + assert!( + return_count >= 3, + "expected multiple explicit return sites for shared final return case, got ops={ops:?}" + ); + } + + #[test] + fn test_for_break_uses_poptop_cleanup() { + let code = compile_exec( + "\ +def f(parts): + for value in parts: + if value: + break +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let pop_iter_count = ops + .iter() + .filter(|op| matches!(op, Instruction::PopIter)) + .count(); + assert_eq!( + pop_iter_count, 1, + "expected only the loop-exhaustion POP_ITER, got ops={ops:?}" + ); + + let break_cleanup_idx = ops + .windows(3) + .position(|window| { + matches!( + window, + [ + Instruction::PopTop, + Instruction::LoadConst { .. }, + Instruction::ReturnValue + ] + ) + }) + .expect("missing POP_TOP/LOAD_CONST/RETURN_VALUE break cleanup"); + let end_for_idx = ops + .iter() + .position(|op| matches!(op, Instruction::EndFor)) + .expect("missing END_FOR"); + assert!( + break_cleanup_idx < end_for_idx, + "expected break cleanup before END_FOR, got ops={ops:?}" + ); + } + #[test] fn test_assert_without_message_raises_class_directly() { let code = compile_exec( @@ -11335,7 +11423,7 @@ def f(names, cls): .filter(|unit| matches!(unit.op, Instruction::ReturnValue)) .count(); - assert_eq!(return_count, 2); + assert_eq!(return_count, 1); } #[test] diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 9b48cab08f..a91c43e898 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -2724,6 +2724,56 @@ fn jump_threading_unconditional(blocks: &mut [Block]) { jump_threading_impl(blocks, false); } +#[derive(Clone, Copy, PartialEq, Eq)] +enum JumpThreadKind { + Plain, + NoInterrupt, +} + +fn jump_thread_kind(instr: AnyInstruction) -> Option { + match instr { + AnyInstruction::Pseudo(PseudoInstruction::Jump { .. }) + | AnyInstruction::Real(Instruction::JumpForward { .. }) + | AnyInstruction::Real(Instruction::JumpBackward { .. }) => Some(JumpThreadKind::Plain), + AnyInstruction::Pseudo(PseudoInstruction::JumpNoInterrupt { .. }) + | AnyInstruction::Real(Instruction::JumpBackwardNoInterrupt { .. }) => { + Some(JumpThreadKind::NoInterrupt) + } + _ => None, + } +} + +fn threaded_jump_instr( + source: AnyInstruction, + target: AnyInstruction, + conditional: bool, +) -> Option { + let target_kind = jump_thread_kind(target)?; + if conditional { + return (target_kind == JumpThreadKind::Plain).then_some(source); + } + + let source_kind = jump_thread_kind(source)?; + if source_kind == JumpThreadKind::NoInterrupt { + return Some(source); + } + Some(match source { + AnyInstruction::Pseudo(_) => PseudoInstruction::Jump { + delta: Arg::marker(), + } + .into(), + AnyInstruction::Real(Instruction::JumpBackwardNoInterrupt { .. }) => { + Instruction::JumpBackward { + delta: Arg::marker(), + } + .into() + } + AnyInstruction::Real(Instruction::JumpForward { .. }) + | AnyInstruction::Real(Instruction::JumpBackward { .. }) => source, + _ => return None, + }) +} + fn jump_threading_impl(blocks: &mut [Block], include_conditional: bool) { let mut changed = true; while changed { @@ -2734,7 +2784,7 @@ fn jump_threading_impl(blocks: &mut [Block], include_conditional: bool) { None => continue, }; let ins = blocks[bi].instructions[last_idx]; - let target = ins.target; + let mut target = ins.target; if target == BlockIdx::NULL { continue; } @@ -2754,6 +2804,10 @@ fn jump_threading_impl(blocks: &mut [Block], include_conditional: bool) { continue; } } + target = next_nonempty_block(blocks, target); + if target == BlockIdx::NULL { + continue; + } // Check if target block's first instruction is an unconditional jump let target_jump = blocks[target.idx()] .instructions @@ -2765,12 +2819,19 @@ fn jump_threading_impl(blocks: &mut [Block], include_conditional: bool) { && target_ins.target != BlockIdx::NULL && target_ins.target != target { + let conditional = is_conditional_jump(&ins.instr); + let Some(threaded_instr) = + threaded_jump_instr(ins.instr, target_ins.instr, conditional) + else { + continue; + }; let final_target = target_ins.target; if ins.target == final_target { continue; } set_to_nop(&mut blocks[bi].instructions[last_idx]); let mut threaded = ins; + threaded.instr = threaded_instr; threaded.arg = OpArg::new(0); threaded.target = final_target; threaded.location = target_ins.location; @@ -3017,6 +3078,12 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { } let target = last.target; + if block_is_exceptional(&blocks[current.idx()]) + || block_is_exceptional(&blocks[target.idx()]) + { + current = next; + continue; + } let small_exit_block = block_exits_scope(&blocks[target.idx()]) && blocks[target.idx()].instructions.len() <= MAX_COPY_SIZE; let no_lineno_no_fallthrough = block_has_no_lineno(&blocks[target.idx()]) @@ -3214,6 +3281,10 @@ fn is_scope_exit_block(block: &Block) -> bool { .is_some_and(|instr| instr.instr.is_scope_exit()) } +fn block_is_exceptional(block: &Block) -> bool { + block.except_handler || block.preserve_lasti +} + fn trailing_conditional_jump_index(block: &Block) -> Option { let last_idx = block.instructions.len().checked_sub(1)?; if is_conditional_jump(&block.instructions[last_idx].instr) @@ -3264,6 +3335,10 @@ fn reorder_conditional_exit_and_jump_blocks(blocks: &mut [Block]) { let mut cursor = exit_start; let mut exit_segment_valid = true; while cursor != BlockIdx::NULL && cursor != jump_start { + if block_is_exceptional(&blocks[cursor.idx()]) { + exit_segment_valid = false; + break; + } if !blocks[cursor.idx()].instructions.is_empty() { if exit_block != BlockIdx::NULL { exit_segment_valid = false; @@ -3288,6 +3363,10 @@ fn reorder_conditional_exit_and_jump_blocks(blocks: &mut [Block]) { let mut jump_block = BlockIdx::NULL; cursor = jump_start; while cursor != BlockIdx::NULL { + if block_is_exceptional(&blocks[cursor.idx()]) { + jump_block = BlockIdx::NULL; + break; + } jump_end = cursor; if blocks[cursor.idx()].instructions.is_empty() { cursor = blocks[cursor.idx()].next; @@ -3345,6 +3424,10 @@ fn reorder_conditional_jump_and_exit_blocks(blocks: &mut [Block]) { let mut cursor = jump_start; let mut jump_segment_valid = true; while cursor != BlockIdx::NULL && cursor != exit_start { + if block_is_exceptional(&blocks[cursor.idx()]) { + jump_segment_valid = false; + break; + } if !blocks[cursor.idx()].instructions.is_empty() { if jump_block != BlockIdx::NULL || !is_jump_only_block(&blocks[cursor.idx()]) { jump_segment_valid = false; @@ -3374,6 +3457,10 @@ fn reorder_conditional_jump_and_exit_blocks(blocks: &mut [Block]) { if cursor == BlockIdx::NULL { break BlockIdx::NULL; } + if block_is_exceptional(&blocks[cursor.idx()]) { + exit_block = BlockIdx::NULL; + break BlockIdx::NULL; + } if !blocks[cursor.idx()].instructions.is_empty() { if exit_block != BlockIdx::NULL { break cursor; @@ -3560,6 +3647,9 @@ fn duplicate_jump_targets_without_lineno(blocks: &mut Vec, predecessors: last_mut.target = new_idx; predecessors[target.idx()] -= 1; predecessors.push(1); + if old_next != BlockIdx::NULL { + predecessors[old_next.idx()] += 1; + } current = old_next; } @@ -3621,7 +3711,7 @@ fn resolve_line_numbers(blocks: &mut Vec) { /// Duplicate `LOAD_CONST None + RETURN_VALUE` for blocks that fall through /// to the final return block. -fn duplicate_end_returns(blocks: &mut [Block]) { +fn duplicate_end_returns(blocks: &mut Vec) { // Walk the block chain and keep the last non-empty block. let mut last_block = BlockIdx::NULL; let mut current = BlockIdx(0); @@ -3652,14 +3742,17 @@ fn duplicate_end_returns(blocks: &mut [Block]) { // Get the return instructions to clone let return_insts: Vec = last_insts[last_insts.len() - 2..].to_vec(); + let predecessors = compute_predecessors(blocks); - // Find non-cold blocks that fall through to the last block - let mut blocks_to_fix = Vec::new(); + // Find non-cold blocks that reach the last return block either by + // fallthrough or as a jump target that should get its own cloned epilogue. + let mut fallthrough_blocks_to_fix = Vec::new(); + let mut jump_targets_to_fix = Vec::new(); current = BlockIdx(0); while current != BlockIdx::NULL { let block = &blocks[current.idx()]; let next = next_nonempty_block(blocks, block.next); - if current != last_block && next == last_block && !block.cold && !block.except_handler { + if current != last_block && !block.cold && !block.except_handler { let last_ins = block.instructions.last(); let has_fallthrough = last_ins .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump()) @@ -3675,19 +3768,54 @@ fn duplicate_end_returns(blocks: &mut [Block]) { AnyInstruction::Real(Instruction::ReturnValue) ) }; - if has_fallthrough && !already_has_return { - blocks_to_fix.push(current); + if next == last_block && has_fallthrough && !already_has_return { + fallthrough_blocks_to_fix.push(current); + } + if predecessors[last_block.idx()] > 1 { + if let Some(cond_idx) = trailing_conditional_jump_index(block) { + let target = next_nonempty_block(blocks, block.instructions[cond_idx].target); + if target == last_block { + jump_targets_to_fix.push((current, cond_idx)); + } + } else if let Some(last) = block.instructions.last() + && last.instr.is_unconditional_jump() + && last.target != BlockIdx::NULL + && next_nonempty_block(blocks, last.target) == last_block + { + jump_targets_to_fix.push((current, block.instructions.len() - 1)); + } } } current = blocks[current.idx()].next; } // Duplicate the return instructions at the end of fall-through blocks - for block_idx in blocks_to_fix { + for block_idx in fallthrough_blocks_to_fix { blocks[block_idx.idx()] .instructions .extend_from_slice(&return_insts); } + + // Clone the final return block for jump predecessors so their target layout + // matches CPython's duplicated exit blocks. + for (block_idx, instr_idx) in jump_targets_to_fix { + let jump = blocks[block_idx.idx()].instructions[instr_idx]; + let mut cloned_return = return_insts.clone(); + for instr in &mut cloned_return { + maybe_propagate_location(instr, jump.location, jump.end_location); + } + let new_idx = BlockIdx(blocks.len() as u32); + let new_block = Block { + cold: blocks[last_block.idx()].cold, + except_handler: blocks[last_block.idx()].except_handler, + instructions: cloned_return, + next: blocks[block_idx.idx()].next, + ..Block::default() + }; + blocks.push(new_block); + blocks[block_idx.idx()].next = new_idx; + blocks[block_idx.idx()].instructions[instr_idx].target = new_idx; + } } /// Label exception targets: walk CFG with except stack, set per-instruction diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap index 4783c0f2d5..3222bf7bed 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap @@ -1,23 +1,13 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 9553 +assertion_line: 10890 expression: "compile_exec(\"\\\nif True and False and False:\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_FALSE (11) - >> 3 CACHE + >> 1 LOAD_CONST (True) + 2 POP_JUMP_IF_FALSE (1) + 3 CACHE 4 NOT_TAKEN - 5 LOAD_CONST (False) - 6 POP_JUMP_IF_FALSE (7) - >> 7 CACHE - 8 NOT_TAKEN - 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (3) - >> 11 CACHE - 12 NOT_TAKEN - 2 13 LOAD_CONST (None) - 14 RETURN_VALUE - 15 LOAD_CONST (None) - 16 RETURN_VALUE + 2 5 LOAD_CONST (None) + 6 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap index 043bf380af..d56432e971 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap @@ -1,27 +1,23 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 9563 +assertion_line: 10900 expression: "compile_exec(\"\\\nif (True and False) or (False and True):\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) + >> 1 LOAD_CONST (True) 2 POP_JUMP_IF_FALSE (5) - >> 3 CACHE + 3 CACHE 4 NOT_TAKEN >> 5 LOAD_CONST (False) - 6 POP_JUMP_IF_TRUE (9) + 6 POP_JUMP_IF_TRUE (7) >> 7 CACHE 8 NOT_TAKEN - >> 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (7) + 9 LOAD_CONST (False) + 10 POP_JUMP_IF_FALSE (1) 11 CACHE 12 NOT_TAKEN - 13 LOAD_CONST (True) - 14 POP_JUMP_IF_FALSE (3) - 15 CACHE - 16 NOT_TAKEN - 2 17 LOAD_CONST (None) - 18 RETURN_VALUE - 19 LOAD_CONST (None) - 20 RETURN_VALUE + 2 13 LOAD_CONST (None) + 14 RETURN_VALUE + 15 LOAD_CONST (None) + 16 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap index bf4960582c..ef7706dc89 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap @@ -1,20 +1,20 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 9543 +assertion_line: 10880 expression: "compile_exec(\"\\\nif True or False or False:\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_TRUE (9) - >> 3 CACHE + >> 1 LOAD_CONST (True) + 2 POP_JUMP_IF_TRUE (11) + 3 CACHE 4 NOT_TAKEN - >> 5 LOAD_CONST (False) - 6 POP_JUMP_IF_TRUE (5) - 7 CACHE + 5 LOAD_CONST (False) + 6 POP_JUMP_IF_TRUE (7) + >> 7 CACHE 8 NOT_TAKEN - >> 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (3) - 11 CACHE + 9 LOAD_CONST (False) + 10 POP_JUMP_IF_FALSE (1) + >> 11 CACHE 12 NOT_TAKEN 2 13 LOAD_CONST (None) diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap index 4573b7a943..3c82c63ad2 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap @@ -1,6 +1,6 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10960 +assertion_line: 10928 expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')):\n with self.subTest(type=type(stop_exc)):\n try:\n async with egg():\n raise stop_exc\n except Exception as ex:\n self.assertIs(ex, stop_exc)\n else:\n self.fail(f'{stop_exc} was suppressed')\n\")" --- 1 0 RESUME (0) @@ -217,8 +217,8 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 200 POP_TOP 3 201 LOAD_CONST (None) - 202 LOAD_CONST (None) - >> 203 LOAD_CONST (None) + >> 202 LOAD_CONST (None) + 203 LOAD_CONST (None) 204 CALL (3) 205 CACHE 206 CACHE @@ -241,14 +241,13 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 223 POP_TOP 224 POP_TOP 225 POP_TOP - 226 JUMP_BACKWARD (203) - 227 CACHE - 228 COPY (3) - 229 POP_EXCEPT - 230 RERAISE (1) + 226 JUMP_BACKWARD_NO_INTERRUPT(202) + 227 COPY (3) + 228 POP_EXCEPT + 229 RERAISE (1) - 2 231 CALL_INTRINSIC_1 (StopIterationError) - 232 RERAISE (1) + 2 230 CALL_INTRINSIC_1 (StopIterationError) + 231 RERAISE (1) 2 MAKE_FUNCTION 3 STORE_NAME (0, test) diff --git a/scripts/dis_dump.py b/scripts/dis_dump.py index 4a80397267..4b72e37046 100755 --- a/scripts/dis_dump.py +++ b/scripts/dis_dump.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 -"""Dump normalized bytecode for Python source files as JSON. +"""Dump bytecode for Python source files as JSON. -Designed to produce comparable output across different Python implementations. -Normalizes away implementation-specific details (byte offsets, memory addresses) -while preserving semantic instruction content. +Designed to compare raw bytecode streams across different Python +implementations while normalizing only display-only details such as memory +addresses in argument reprs. Usage: python dis_dump.py Lib/ @@ -167,42 +167,60 @@ def _resolve_arg_fallback(code, opname, arg): def _extract_instructions(code): - """Extract normalized instruction list from a code object. - - - Filters out CACHE/PRECALL instructions - - Converts jump targets from byte offsets to instruction indices - - Resolves argument names via fallback when argrepr is missing - - Normalizes argument representations - """ + """Extract a raw code-unit instruction stream from a code object.""" try: raw = list(dis.get_instructions(code)) except Exception as e: return [["ERROR", str(e)]] - # Build filtered list and offset-to-index mapping for the normalized stream. - # This must use post-decomposition indices; otherwise a superinstruction that - # expands into multiple logical ops shifts later jump targets by 1. - filtered = [] + def _metadata_cache_slot_offsets(inst): + cache_offset = getattr(inst, "cache_offset", None) + end_offset = getattr(inst, "end_offset", None) + if ( + isinstance(cache_offset, int) + and isinstance(end_offset, int) + and end_offset >= cache_offset + ): + return range(cache_offset, end_offset, 2) + cache_info = getattr(inst, "cache_info", None) or () + cache_units = sum(size for _, size, _ in cache_info) + return range(inst.offset + 2, inst.offset + 2 + cache_units * 2, 2) + + explicit_offsets = {inst.offset for inst in raw} + cache_counts = {} + stream = [] offset_to_idx = {} - normalized_idx = 0 - for inst in raw: - if inst.opname in SKIP_OPS: - continue - opname = _OPNAME_NORMALIZE.get(inst.opname, inst.opname) - offset_to_idx[inst.offset] = normalized_idx - normalized_idx += len(_SUPER_DECOMPOSE.get(opname, (opname,))) - filtered.append(inst) - - # Map offsets that land on CACHE slots to the next real instruction - for inst in raw: - if inst.offset not in offset_to_idx: - for fi, finst in enumerate(filtered): - if finst.offset >= inst.offset: - offset_to_idx[inst.offset] = fi - break + for i, inst in enumerate(raw): + explicit_cache_count = 0 + next_offset = inst.offset + 2 + j = i + 1 + while ( + j < len(raw) + and raw[j].opname == "CACHE" + and raw[j].offset == next_offset + ): + explicit_cache_count += 1 + next_offset += 2 + j += 1 + cache_counts[inst.offset] = explicit_cache_count + if inst.opname not in SKIP_OPS: + offset_to_idx[inst.offset] = len(stream) + stream.append(("inst", inst)) + if explicit_cache_count == 0: + for cache_offset in _metadata_cache_slot_offsets(inst): + if cache_offset in explicit_offsets: + continue + cache_counts[inst.offset] += 1 + offset_to_idx[cache_offset] = len(stream) + stream.append(("cache", cache_offset)) result = [] - for inst in filtered: + for kind, payload in stream: + if kind == "cache": + result.append(["CACHE"]) + continue + + inst = payload opname = _OPNAME_NORMALIZE.get(inst.opname, inst.opname) # Decompose superinstructions into individual ops @@ -232,32 +250,15 @@ def _extract_instructions(code): if is_backward: # Target = current_offset + INSTR_SIZE + cache # - arg * INSTR_SIZE - # Try different cache sizes (NOT_TAKEN=1 for JUMP_BACKWARD, 0 for NO_INTERRUPT) - if "NO_INTERRUPT" in inst.opname: - cache_order = (0, 1, 2) - else: - cache_order = (1, 0, 2, 3) - for cache in cache_order: - target_off = inst.offset + 2 + cache * 2 - inst.arg * 2 - if target_off >= 0 and target_off in offset_to_idx: - target_idx = offset_to_idx[target_off] - break + cache = cache_counts.get(inst.offset, 0) + target_off = inst.offset + 2 + cache * 2 - inst.arg * 2 + if target_off >= 0 and target_off in offset_to_idx: + target_idx = offset_to_idx[target_off] elif inst.arg is not None: - # Forward jumps: compute target offset using cache entry count. - # POP_JUMP_IF_* have 1 cache entry (NOT_TAKEN), others have 0. - if "POP_JUMP_IF" in inst.opname: - cache_order = (1, 0, 2) - elif inst.opname == "FOR_ITER": - cache_order = (0, 1, 2) - elif inst.opname == "SEND": - cache_order = (1, 0, 2) - else: - cache_order = (0, 1, 2) - for extra in cache_order: - target_off = inst.offset + 2 + extra * 2 + inst.arg * 2 - if target_off in offset_to_idx: - target_idx = offset_to_idx[target_off] - break + cache = cache_counts.get(inst.offset, 0) + target_off = inst.offset + 2 + cache * 2 + inst.arg * 2 + if target_off in offset_to_idx: + target_idx = offset_to_idx[target_off] if target_idx is None: target_idx = inst.argval result.append([opname, "->%d" % target_idx])