diff --git a/Lib/opcode.py b/Lib/opcode.py index f74acf002e9..0e9520b6832 100644 --- a/Lib/opcode.py +++ b/Lib/opcode.py @@ -46,6 +46,75 @@ hascompare = [opmap["COMPARE_OP"]] _cache_format = { + "LOAD_GLOBAL": { + "counter": 1, + "index": 1, + "module_keys_version": 1, + "builtin_keys_version": 1, + }, + "BINARY_OP": { + "counter": 1, + "descr": 4, + }, + "UNPACK_SEQUENCE": { + "counter": 1, + }, + "COMPARE_OP": { + "counter": 1, + }, + "CONTAINS_OP": { + "counter": 1, + }, + "FOR_ITER": { + "counter": 1, + }, + "LOAD_SUPER_ATTR": { + "counter": 1, + }, + "LOAD_ATTR": { + "counter": 1, + "version": 2, + "keys_version": 2, + "descr": 4, + }, + "STORE_ATTR": { + "counter": 1, + "version": 2, + "index": 1, + }, + "CALL": { + "counter": 1, + "func_version": 2, + }, + "CALL_KW": { + "counter": 1, + "func_version": 2, + }, + "STORE_SUBSCR": { + "counter": 1, + }, + "SEND": { + "counter": 1, + }, + "JUMP_BACKWARD": { + "counter": 1, + }, + "TO_BOOL": { + "counter": 1, + "version": 2, + }, + "POP_JUMP_IF_TRUE": { + "counter": 1, + }, + "POP_JUMP_IF_FALSE": { + "counter": 1, + }, + "POP_JUMP_IF_NONE": { + "counter": 1, + }, + "POP_JUMP_IF_NOT_NONE": { + "counter": 1, + }, } _inline_cache_entries = { diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py index 232c1ca6d7c..8ba783bf141 100644 --- a/Lib/test/test_dis.py +++ b/Lib/test/test_dis.py @@ -2236,7 +2236,6 @@ def f(opcode, oparg, offset, *init_args): def get_instructions(self, code): return dis._get_instructions_bytes(code) - @unittest.expectedFailure # TODO: RUSTPYTHON; no inline caches def test_start_offset(self): # When no extended args are present, # start_offset should be equal to offset @@ -2289,7 +2288,6 @@ def last_item(iterable): self.assertEqual(14, instructions[6].offset) self.assertEqual(8, instructions[6].start_offset) - @unittest.expectedFailure # TODO: RUSTPYTHON; no inline caches def test_cache_offset_and_end_offset(self): code = bytes([ opcode.opmap["LOAD_GLOBAL"], 0x01, @@ -2588,7 +2586,6 @@ def f(): with contextlib.redirect_stderr(io.StringIO()): _ = self.invoke_dis('--unknown') - @unittest.expectedFailure # TODO: RUSTPYTHON def test_show_cache(self): # test 'python -m dis -C/--show-caches' source = 'print()' diff --git a/Lib/test/test_inspect/test_inspect.py b/Lib/test/test_inspect/test_inspect.py index 77a160bcf19..3f8093abbf1 100644 --- a/Lib/test/test_inspect/test_inspect.py +++ b/Lib/test/test_inspect/test_inspect.py @@ -2020,7 +2020,6 @@ def function(): _global_ref = object() class TestGetClosureVars(unittest.TestCase): - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: Closu[132 chars]={'print': }, unbound=set()) != Closu[132 chars]={'print': }, unbound={'unbound_ref'}) def test_name_resolution(self): # Basic test of the 4 different resolution mechanisms def f(nonlocal_ref): @@ -2036,7 +2035,6 @@ def g(local_ref): builtin_vars, unbound_names) self.assertEqual(inspect.getclosurevars(f(_arg)), expected) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: Closu[132 chars]={'print': }, unbound=set()) != Closu[132 chars]={'print': }, unbound={'unbound_ref'}) def test_generator_closure(self): def f(nonlocal_ref): def g(local_ref): @@ -2052,7 +2050,6 @@ def g(local_ref): builtin_vars, unbound_names) self.assertEqual(inspect.getclosurevars(f(_arg)), expected) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: Closu[132 chars]={'print': }, unbound=set()) != Closu[132 chars]={'print': }, unbound={'unbound_ref'}) def test_method_closure(self): class C: def f(self, nonlocal_ref): @@ -2068,7 +2065,6 @@ def g(local_ref): builtin_vars, unbound_names) self.assertEqual(inspect.getclosurevars(C().f(_arg)), expected) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: Closu[139 chars]als={}, builtins={'print': }, unbound=set()) != Closu[34 chars]uiltins={'print': }, unbound={'path'}) def test_builtins_fallback(self): f, ns = self._private_globals() ns.pop("__builtins__", None) expected = inspect.ClosureVars({}, {}, {"print":print}, {"path"}) self.assertEqual(inspect.getclosurevars(f), expected) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: ClosureVars(nonlocals={}, globals={}, builtins={}, unbound={'print'}) != ClosureVars(nonlocals={}, globals={}, builtins={'path': 1}, unbound={'print'}) def test_builtins_as_dict(self): f, ns = self._private_globals() ns["__builtins__"] = {"path":1} expected = inspect.ClosureVars({}, {}, {"path":1}, {"print"}) self.assertEqual(inspect.getclosurevars(f), expected) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: Closu[38 chars]ins={}, unbound={'print'}) != Closu[38 chars]ins={'path': POP_JUMP_IF_FALSE to non-jump diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index ebd1e670ac0..72ce4c7350b 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -1173,6 +1173,7 @@ impl Compiler { end_location, except_handler, lineno_override, + cache_entries: 0, }); } @@ -1432,7 +1433,7 @@ impl Compiler { if matches!(info.fb_type, FBlockType::AsyncWith) { emit!(self, Instruction::GetAwaitable { arg: 2 }); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + let _ = self.compile_yield_from_sequence(true)?; } // Pop the __exit__ result @@ -2009,7 +2010,10 @@ impl Compiler { NameOp::Global => { let idx = self.get_global_name_index(&name); let op = match usage { - NameUsage::Load => Instruction::LoadGlobal, + NameUsage::Load => { + self.emit_load_global(idx, false); + return Ok(()); + } NameUsage::Store => Instruction::StoreGlobal, NameUsage::Delete => Instruction::DeleteGlobal, }; @@ -2370,12 +2374,14 @@ impl Compiler { } } ast::Stmt::Break(_) => { + emit!(self, Instruction::Nop); // NOP for line tracing // Unwind fblock stack until we find a loop, emitting cleanup for each fblock self.compile_break_continue(statement.range(), true)?; let dead = self.new_block(); self.switch_to_block(dead); } ast::Stmt::Continue(_) => { + emit!(self, Instruction::Nop); // NOP for line tracing // Unwind fblock stack until we find a loop, emitting cleanup for each fblock self.compile_break_continue(statement.range(), false)?; let dead = self.new_block(); @@ -2448,7 +2454,7 @@ impl Compiler { } } ast::Stmt::Pass(_) => { - // No need to emit any code here :) + emit!(self, Instruction::Nop); // NOP for line tracing } ast::Stmt::TypeAlias(ast::StmtTypeAlias { name, @@ -2967,24 +2973,14 @@ impl Compiler { emit!(self, Instruction::PopExcept); // RERAISE 0: re-raise the original exception to outer handler - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack - } - ); + emit!(self, Instruction::Reraise { depth: 0 }); } if let Some(cleanup) = finally_cleanup_block { self.switch_to_block(cleanup); emit!(self, Instruction::Copy { index: 3_u32 }); emit!(self, Instruction::PopExcept); - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack - } - ); + emit!(self, Instruction::Reraise { depth: 1 }); } self.switch_to_block(end_block); @@ -3122,12 +3118,7 @@ impl Compiler { // Stack at entry: [prev_exc (at handler_depth), lasti, exc] // This RERAISE is within ExceptionHandler scope, so it routes to cleanup_block // which does COPY 3; POP_EXCEPT; RERAISE - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack, - } - ); + emit!(self, Instruction::Reraise { depth: 1 }); } // Switch to normal exit block - this is where handler body success continues @@ -3175,12 +3166,7 @@ impl Compiler { // RERAISE 0 // Stack: [prev_exc, exc] - exception is on stack from PUSH_EXC_INFO // NOTE: We emit RERAISE 0 BEFORE popping fblock so it is within cleanup handler scope - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack, - } - ); + emit!(self, Instruction::Reraise { depth: 0 }); // Pop EXCEPTION_HANDLER fblock // Pop after RERAISE so the instruction has the correct exception handler @@ -3194,12 +3180,7 @@ impl Compiler { self.switch_to_block(cleanup_block); emit!(self, Instruction::Copy { index: 3_u32 }); emit!(self, Instruction::PopExcept); - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack, - } - ); + emit!(self, Instruction::Reraise { depth: 1 }); // We successfully ran the try block: // else: @@ -3271,12 +3252,7 @@ impl Compiler { // RERAISE 0: re-raise the original exception to outer handler // Stack: [lasti, prev_exc, exc] - exception is on top - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack, - } - ); + emit!(self, Instruction::Reraise { depth: 0 }); } // finally cleanup block @@ -3289,12 +3265,7 @@ impl Compiler { // POP_EXCEPT: restore prev_exc as current exception emit!(self, Instruction::PopExcept); // RERAISE 1: reraise with lasti from stack - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack, - } - ); + emit!(self, Instruction::Reraise { depth: 1 }); } // End block - continuation point after try-finally @@ -3667,23 +3638,13 @@ impl Compiler { emit!(self, Instruction::Copy { index: 2_u32 }); emit!(self, Instruction::PopExcept); - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack - } - ); + emit!(self, Instruction::Reraise { depth: 0 }); if let Some(cleanup) = finally_cleanup_block { self.switch_to_block(cleanup); emit!(self, Instruction::Copy { index: 3_u32 }); emit!(self, Instruction::PopExcept); - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack - } - ); + emit!(self, Instruction::Reraise { depth: 1 }); } } @@ -4030,7 +3991,7 @@ impl Compiler { emit!(self, Instruction::LoadDeref(idx)); } else { let cond_annotations_name = self.name("__conditional_annotations__"); - emit!(self, Instruction::LoadGlobal(cond_annotations_name)); + self.emit_load_global(cond_annotations_name, false); } // CONTAINS_OP (in) emit!(self, Instruction::ContainsOp(bytecode::Invert::No)); @@ -4522,7 +4483,7 @@ impl Compiler { // Load (global) __name__ and store as __module__ let dunder_name = self.name("__name__"); - emit!(self, Instruction::LoadGlobal(dunder_name)); + self.emit_load_global(dunder_name, false); let dunder_module = self.name("__module__"); emit!(self, Instruction::StoreName(dunder_module)); @@ -4941,7 +4902,7 @@ impl Compiler { emit!(self, Instruction::Call { nargs: 0 }); // [bound_aexit, awaitable] emit!(self, Instruction::GetAwaitable { arg: 1 }); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + let _ = self.compile_yield_from_sequence(true)?; } else { // Load __exit__ and __enter__, then call __enter__ emit!( @@ -5024,7 +4985,7 @@ impl Compiler { if is_async { emit!(self, Instruction::GetAwaitable { arg: 2 }); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + let _ = self.compile_yield_from_sequence(true)?; } emit!(self, Instruction::PopTop); // Pop __exit__ result emit!( @@ -5062,7 +5023,7 @@ impl Compiler { if is_async { emit!(self, Instruction::GetAwaitable { arg: 2 }); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + let _ = self.compile_yield_from_sequence(true)?; } // TO_BOOL + POP_JUMP_IF_TRUE: check if exception is suppressed @@ -5136,6 +5097,7 @@ impl Compiler { let for_block = self.new_block(); let else_block = self.new_block(); let after_block = self.new_block(); + let mut end_async_for_target = BlockIdx::NULL; // The thing iterated: self.compile_expression(iter)?; @@ -5155,9 +5117,10 @@ impl Compiler { emit!(self, PseudoInstruction::SetupFinally { target: else_block }); emit!(self, Instruction::GetANext); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + end_async_for_target = self.compile_yield_from_sequence(true)?; // POP_BLOCK for SETUP_FINALLY - only GetANext/yield_from are protected emit!(self, PseudoInstruction::PopBlock); + emit!(self, Instruction::NotTaken); // Success block for __anext__ self.compile_store(target)?; @@ -5191,7 +5154,7 @@ impl Compiler { let saved_range = self.current_source_range; self.set_source_range(iter.range()); if is_async { - emit!(self, Instruction::EndAsyncFor); + self.emit_end_async_for(end_async_for_target); } else { emit!(self, Instruction::EndFor); emit!(self, Instruction::PopIter); @@ -6785,7 +6748,7 @@ impl Compiler { /// CLEANUP_THROW /// exit: /// END_SEND - fn compile_yield_from_sequence(&mut self, is_await: bool) -> CompileResult<()> { + fn compile_yield_from_sequence(&mut self, is_await: bool) -> CompileResult { let send_block = self.new_block(); let fail_block = self.new_block(); let exit_block = self.new_block(); @@ -6841,7 +6804,7 @@ impl Compiler { self.switch_to_block(exit_block); emit!(self, Instruction::EndSend); - Ok(()) + Ok(send_block) } fn compile_expression(&mut self, expression: &ast::Expr) -> CompileResult<()> { @@ -6896,7 +6859,11 @@ impl Compiler { if let Some(super_type) = self.can_optimize_super_call(value, attr.as_str()) { // super().attr or super(cls, self).attr optimization // Stack: [global_super, class, self] → LOAD_SUPER_ATTR → [attr] + // Set source range to super() call for arg-loading instructions + let super_range = value.range(); + self.set_source_range(super_range); self.load_args_for_super(&super_type)?; + self.set_source_range(super_range); let idx = self.name(attr.as_str()); match super_type { SuperCallType::TwoArg { .. } => { @@ -6982,7 +6949,7 @@ impl Compiler { self.compile_expression(value)?; emit!(self, Instruction::GetAwaitable { arg: 0 }); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + let _ = self.compile_yield_from_sequence(true)?; } ast::Expr::YieldFrom(ast::ExprYieldFrom { value, .. }) => { match self.ctx.func { @@ -6998,7 +6965,7 @@ impl Compiler { self.compile_expression(value)?; emit!(self, Instruction::GetYieldFromIter); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(false)?; + let _ = self.compile_yield_from_sequence(false)?; } ast::Expr::Name(ast::ExprName { id, .. }) => self.load_name(id.as_str())?, ast::Expr::Lambda(ast::ExprLambda { @@ -7353,7 +7320,11 @@ impl Compiler { if let Some(super_type) = self.can_optimize_super_call(value, attr.as_str()) { // super().method() or super(cls, self).method() optimization // Stack: [global_super, class, self] → LOAD_SUPER_METHOD → [method, self] + // Set source range to the super() call for LOAD_GLOBAL/LOAD_DEREF/etc. + let super_range = value.range(); + self.set_source_range(super_range); self.load_args_for_super(&super_type)?; + self.set_source_range(super_range); let idx = self.name(attr.as_str()); match super_type { SuperCallType::TwoArg { .. } => { @@ -7363,7 +7334,11 @@ impl Compiler { self.emit_load_zero_super_method(idx); } } - self.codegen_call_helper(0, args, call_range)?; + // NOP for line tracking at .method( line + self.set_source_range(attr.range()); + emit!(self, Instruction::Nop); + // CALL at .method( line (not the full expression line) + self.codegen_call_helper(0, args, attr.range())?; } else { // Normal method call: compile object, then LOAD_ATTR with method flag // LOAD_ATTR(method=1) pushes [method, self_or_null] on stack @@ -7653,6 +7628,7 @@ impl Compiler { let mut loop_labels = vec![]; for generator in generators { let loop_block = self.new_block(); + let if_cleanup_block = self.new_block(); let after_block = self.new_block(); if loop_labels.is_empty() { @@ -7670,8 +7646,8 @@ impl Compiler { } } - loop_labels.push((loop_block, after_block, generator.is_async)); self.switch_to_block(loop_block); + let mut end_async_for_target = BlockIdx::NULL; if generator.is_async { emit!( self, @@ -7686,7 +7662,7 @@ impl Compiler { after_block, )?; self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + end_async_for_target = self.compile_yield_from_sequence(true)?; // POP_BLOCK before store: only __anext__/yield_from are // protected by SetupFinally targeting END_ASYNC_FOR. emit!(self, PseudoInstruction::PopBlock); @@ -7701,23 +7677,35 @@ impl Compiler { ); self.compile_store(&generator.target)?; } + loop_labels.push(( + loop_block, + if_cleanup_block, + after_block, + generator.is_async, + end_async_for_target, + )); // Now evaluate the ifs: for if_condition in &generator.ifs { - self.compile_jump_if(if_condition, false, loop_block)? + self.compile_jump_if(if_condition, false, if_cleanup_block)? } } compile_element(self)?; - for (loop_block, after_block, is_async) in loop_labels.iter().rev().copied() { + for (loop_block, if_cleanup_block, after_block, is_async, end_async_for_target) in + loop_labels.iter().rev().copied() + { + emit!(self, PseudoInstruction::Jump { target: loop_block }); + + self.switch_to_block(if_cleanup_block); emit!(self, PseudoInstruction::Jump { target: loop_block }); self.switch_to_block(after_block); if is_async { // EndAsyncFor pops both the exception and the aiter // (handler depth is before GetANext, so aiter is at handler depth) - emit!(self, Instruction::EndAsyncFor); + self.emit_end_async_for(end_async_for_target); } else { // END_FOR + POP_ITER pattern (CPython 3.14) emit!(self, Instruction::EndFor); @@ -7755,7 +7743,7 @@ impl Compiler { if is_async_list_set_dict_comprehension { emit!(self, Instruction::GetAwaitable { arg: 0 }); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + let _ = self.compile_yield_from_sequence(true)?; } Ok(()) @@ -7866,6 +7854,7 @@ impl Compiler { let mut loop_labels = vec![]; for (i, generator) in generators.iter().enumerate() { let loop_block = self.new_block(); + let if_cleanup_block = self.new_block(); let after_block = self.new_block(); if i > 0 { @@ -7878,13 +7867,13 @@ impl Compiler { } } - loop_labels.push((loop_block, after_block, generator.is_async)); self.switch_to_block(loop_block); + let mut end_async_for_target = BlockIdx::NULL; if generator.is_async { emit!(self, Instruction::GetANext); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + end_async_for_target = self.compile_yield_from_sequence(true)?; self.compile_store(&generator.target)?; } else { emit!( @@ -7895,10 +7884,17 @@ impl Compiler { ); self.compile_store(&generator.target)?; } + loop_labels.push(( + loop_block, + if_cleanup_block, + after_block, + generator.is_async, + end_async_for_target, + )); // Evaluate the if conditions for if_condition in &generator.ifs { - self.compile_jump_if(if_condition, false, loop_block)?; + self.compile_jump_if(if_condition, false, if_cleanup_block)?; } } @@ -7906,11 +7902,17 @@ impl Compiler { compile_element(self)?; // Step 7: Close all loops - for (loop_block, after_block, is_async) in loop_labels.iter().rev().copied() { + for (loop_block, if_cleanup_block, after_block, is_async, end_async_for_target) in + loop_labels.iter().rev().copied() + { emit!(self, PseudoInstruction::Jump { target: loop_block }); + + self.switch_to_block(if_cleanup_block); + emit!(self, PseudoInstruction::Jump { target: loop_block }); + self.switch_to_block(after_block); if is_async { - emit!(self, Instruction::EndAsyncFor); + self.emit_end_async_for(end_async_for_target); // Pop the iterator emit!(self, Instruction::PopTop); } else { @@ -7947,12 +7949,7 @@ impl Compiler { emit!(self, Instruction::StoreFast(idx)); } // Re-raise the exception - emit!( - self, - Instruction::RaiseVarargs { - kind: bytecode::RaiseKind::ReraiseFromStack - } - ); + emit!(self, Instruction::Reraise { depth: 0 }); // Normal end path self.switch_to_block(end_block); @@ -8013,6 +8010,7 @@ impl Compiler { end_location, except_handler, lineno_override: None, + cache_entries: 0, }); } @@ -8046,6 +8044,10 @@ impl Compiler { emit!(self, Instruction::ReturnValue) } + fn emit_end_async_for(&mut self, send_target: BlockIdx) { + self._emit(Instruction::EndAsyncFor, OpArg::NULL, send_target); + } + /// Emit LOAD_ATTR for attribute access (method=false). /// Encodes: (name_idx << 1) | 0 fn emit_load_attr(&mut self, name_idx: u32) { @@ -8066,6 +8068,13 @@ impl Compiler { self.emit_arg(encoded, |arg| Instruction::LoadAttr { idx: arg }) } + /// Emit LOAD_GLOBAL. + /// Encodes: (name_idx << 1) | push_null_bit + fn emit_load_global(&mut self, name_idx: u32, push_null: bool) { + let encoded = (name_idx << 1) | u32::from(push_null); + self.emit_arg(encoded, Instruction::LoadGlobal); + } + /// Emit LOAD_SUPER_ATTR for 2-arg super().attr access. /// Encodes: (name_idx << 2) | 0b10 (method=0, class=1) fn emit_load_super_attr(&mut self, name_idx: u32) { @@ -8258,7 +8267,7 @@ impl Compiler { if is_async { emit!(self, Instruction::GetAwaitable { arg: 2 }); self.emit_load_const(ConstantData::None); - self.compile_yield_from_sequence(true)?; + let _ = self.compile_yield_from_sequence(true)?; } emit!(self, Instruction::PopTop); diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 97f4a8c2d8c..4363ffaa768 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -106,6 +106,8 @@ pub struct InstructionInfo { pub except_handler: Option, /// Override line number for linetable (e.g., line 0 for module RESUME) pub lineno_override: Option, + /// Number of CACHE code units emitted after this instruction + pub cache_entries: u32, } /// Exception handler information for an instruction. @@ -211,6 +213,7 @@ impl CodeInfo { label_exception_targets(&mut self.blocks); push_cold_blocks_to_end(&mut self.blocks); normalize_jumps(&mut self.blocks); + self.optimize_load_global_push_null(); let max_stackdepth = self.max_stackdepth()?; let cell2arg = self.cell2arg(); @@ -252,36 +255,99 @@ impl CodeInfo { // Convert pseudo ops and remove resulting NOPs (keep line-marker NOPs) convert_pseudo_ops(&mut blocks, varname_cache.len() as u32); - for block in blocks - .iter_mut() - .filter(|b| b.next != BlockIdx::NULL || !b.instructions.is_empty()) - { - // Collect lines that have non-NOP instructions in this block - let non_nop_lines: IndexSet<_> = block - .instructions - .iter() - .filter(|ins| !matches!(ins.instr.real(), Some(Instruction::Nop))) - .map(|ins| ins.location.line) - .collect(); - let mut kept_nop_lines: IndexSet = IndexSet::default(); - block.instructions.retain(|ins| { - if matches!(ins.instr.real(), Some(Instruction::Nop)) { - let line = ins.location.line; - // Remove if another instruction covers this line, - // or if we already kept a NOP for this line - if non_nop_lines.contains(&line) || kept_nop_lines.contains(&line) { - return false; + // Remove redundant NOPs, keeping line-marker NOPs only when + // they are needed to preserve tracing. + let mut block_order = Vec::new(); + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + block_order.push(current); + current = blocks[current.idx()].next; + } + for block_idx in block_order { + let bi = block_idx.idx(); + let mut src_instructions = core::mem::take(&mut blocks[bi].instructions); + let mut kept = Vec::with_capacity(src_instructions.len()); + let mut prev_lineno = -1i32; + + for src in 0..src_instructions.len() { + let instr = src_instructions[src]; + let lineno = instr + .lineno_override + .unwrap_or_else(|| instr.location.line.get() as i32); + let mut remove = false; + + if matches!(instr.instr.real(), Some(Instruction::Nop)) { + // Remove location-less NOPs. + if lineno < 0 || prev_lineno == lineno { + remove = true; + } + // Remove if the next instruction has same line or no line. + else if src < src_instructions.len() - 1 { + let next_lineno = + src_instructions[src + 1] + .lineno_override + .unwrap_or_else(|| { + src_instructions[src + 1].location.line.get() as i32 + }); + if next_lineno == lineno { + remove = true; + } else if next_lineno < 0 { + src_instructions[src + 1].lineno_override = Some(lineno); + remove = true; + } + } + // Last instruction in block: compare with first real location + // in the next non-empty block. + else { + let mut next = blocks[bi].next; + while next != BlockIdx::NULL && blocks[next.idx()].instructions.is_empty() { + next = blocks[next.idx()].next; + } + if next != BlockIdx::NULL { + let mut next_lineno = None; + for next_instr in &blocks[next.idx()].instructions { + let line = next_instr + .lineno_override + .unwrap_or_else(|| next_instr.location.line.get() as i32); + if matches!(next_instr.instr.real(), Some(Instruction::Nop)) + && line < 0 + { + continue; + } + next_lineno = Some(line); + break; + } + if next_lineno.is_some_and(|line| line == lineno) { + remove = true; + } + } } - kept_nop_lines.insert(line); } - true - }); + + if !remove { + kept.push(instr); + prev_lineno = lineno; + } + } + + blocks[bi].instructions = kept; + } + + // Pre-compute cache_entries for real (non-pseudo) instructions + for block in blocks.iter_mut() { + for instr in &mut block.instructions { + if let AnyInstruction::Real(op) = instr.instr { + instr.cache_entries = op.cache_entries() as u32; + } + } } let mut block_to_offset = vec![Label(0); blocks.len()]; // block_to_index: maps block idx to instruction index (for exception table) - // This is the index into the final instructions array, including EXTENDED_ARG + // This is the index into the final instructions array, including EXTENDED_ARG and CACHE let mut block_to_index = vec![0u32; blocks.len()]; + // The offset (in code units) of END_SEND from SEND in the yield-from sequence. + const END_SEND_OFFSET: u32 = 5; loop { let mut num_instructions = 0; for (idx, block) in iter_blocks(&blocks) { @@ -291,14 +357,14 @@ impl CodeInfo { // and instructions array index == byte offset (each instruction is 1 CodeUnit) block_to_index[idx.idx()] = num_instructions as u32; for instr in &block.instructions { - num_instructions += instr.arg.instr_size(); + num_instructions += instr.arg.instr_size() + instr.cache_entries as usize; } } instructions.reserve_exact(num_instructions); locations.reserve_exact(num_instructions); - let mut recompile_extended_arg = false; + let mut recompile = false; let mut next_block = BlockIdx(0); while next_block != BlockIdx::NULL { let block = &mut blocks[next_block]; @@ -306,43 +372,71 @@ impl CodeInfo { let mut current_offset = block_to_offset[next_block.idx()].0; for info in &mut block.instructions { let target = info.target; - if target != BlockIdx::NULL { - let new_arg = OpArg::new(block_to_offset[target.idx()].0); - recompile_extended_arg |= new_arg.instr_size() != info.arg.instr_size(); - info.arg = new_arg; - } + let mut op = info.instr.expect_real(); + let old_arg_size = info.arg.instr_size(); + let old_cache_entries = info.cache_entries; + // Keep offsets fixed within this pass: changes in jump + // arg/cache sizes only take effect in the next iteration. + let offset_after = current_offset + old_arg_size as u32 + old_cache_entries; - // Convert JUMP pseudo to real instructions (direction depends on offset) - let op = match info.instr { - AnyInstruction::Pseudo(PseudoInstruction::Jump { .. }) - if target != BlockIdx::NULL => - { - let target_offset = block_to_offset[target.idx()].0; - if target_offset > current_offset { - Instruction::JumpForward { + if target != BlockIdx::NULL { + let target_offset = block_to_offset[target.idx()].0; + // Direction must be based on concrete instruction offsets. + // Empty blocks can share offsets, so block-order-based resolution + // may classify some jumps incorrectly. + op = match op { + Instruction::JumpForward { .. } if target_offset <= current_offset => { + Instruction::JumpBackward { target: Arg::marker(), } - } else { - Instruction::JumpBackward { + } + Instruction::JumpBackward { .. } if target_offset > current_offset => { + Instruction::JumpForward { target: Arg::marker(), } } - } - AnyInstruction::Pseudo(PseudoInstruction::JumpNoInterrupt { .. }) - if target != BlockIdx::NULL => - { - // JumpNoInterrupt is always backward (used in yield-from/await loops) - Instruction::JumpBackwardNoInterrupt { - target: Arg::marker(), + Instruction::JumpBackwardNoInterrupt { .. } + if target_offset > current_offset => + { + Instruction::JumpForward { + target: Arg::marker(), + } } - } - other => other.expect_real(), - }; + _ => op, + }; + info.instr = op.into(); + let updated_cache = op.cache_entries() as u32; + recompile |= updated_cache != old_cache_entries; + info.cache_entries = updated_cache; + let new_arg = if matches!(op, Instruction::EndAsyncFor) { + let arg = offset_after + .checked_sub(target_offset + END_SEND_OFFSET) + .expect("END_ASYNC_FOR target must be before instruction"); + OpArg::new(arg) + } else if matches!( + op, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. } + ) { + let arg = offset_after + .checked_sub(target_offset) + .expect("backward jump target must be before instruction"); + OpArg::new(arg) + } else { + let arg = target_offset + .checked_sub(offset_after) + .expect("forward jump target must be after instruction"); + OpArg::new(arg) + }; + recompile |= new_arg.instr_size() != old_arg_size; + info.arg = new_arg; + } + let cache_count = info.cache_entries as usize; let (extras, lo_arg) = info.arg.split(); locations.extend(core::iter::repeat_n( (info.location, info.end_location), - info.arg.instr_size(), + info.arg.instr_size() + cache_count, )); // Collect linetable locations with lineno_override support let lt_loc = LineTableLocation { @@ -354,17 +448,26 @@ impl CodeInfo { end_col: info.end_location.character_offset.to_zero_indexed() as i32, }; linetable_locations.extend(core::iter::repeat_n(lt_loc, info.arg.instr_size())); + // CACHE entries inherit parent instruction's location + if cache_count > 0 { + linetable_locations.extend(core::iter::repeat_n(lt_loc, cache_count)); + } instructions.extend( extras .map(|byte| CodeUnit::new(Instruction::ExtendedArg, byte)) .chain([CodeUnit { op, arg: lo_arg }]), ); - current_offset += info.arg.instr_size() as u32; + // Emit CACHE code units after the instruction + instructions.extend(core::iter::repeat_n( + CodeUnit::new(Instruction::Cache, 0.into()), + cache_count, + )); + current_offset = offset_after; } next_block = block.next; } - if !recompile_extended_arg { + if !recompile { break; } @@ -599,6 +702,33 @@ impl CodeInfo { } } + /// LOAD_GLOBAL + PUSH_NULL -> LOAD_GLOBAL , NOP + fn optimize_load_global_push_null(&mut self) { + for block in &mut self.blocks { + let mut i = 0; + while i + 1 < block.instructions.len() { + let curr = &block.instructions[i]; + let next = &block.instructions[i + 1]; + + let (Some(Instruction::LoadGlobal(_)), Some(Instruction::PushNull)) = + (curr.instr.real(), next.instr.real()) + else { + i += 1; + continue; + }; + + let oparg = u32::from(block.instructions[i].arg); + if (oparg & 1) != 0 { + i += 1; + continue; + } + + block.instructions[i].arg = OpArg::new(oparg | 1); + block.instructions.remove(i + 1); + } + } + } + /// Convert LOAD_CONST for small integers to LOAD_SMALL_INT /// maybe_instr_make_load_smallint fn convert_to_load_small_int(&mut self) { @@ -1005,6 +1135,18 @@ fn generate_linetable( // Get line information let line = loc.line; + + // NO_LOCATION: emit PyCodeLocationInfoKind::None entries (CACHE, etc.) + if line == -1 { + linetable.push( + 0x80 | ((PyCodeLocationInfoKind::None as u8) << 3) | ((entry_length - 1) as u8), + ); + // Do NOT update prev_line + length -= entry_length; + i += entry_length; + continue; + } + let end_line = loc.end_line; let line_delta = line - prev_line; let end_line_delta = end_line - line; @@ -1105,8 +1247,8 @@ fn generate_exception_table(blocks: &[Block], block_to_index: &[u32]) -> Box<[u8 // This matches how frame.rs uses lasti for (_, block) in iter_blocks(blocks) { for instr in &block.instructions { - // instr_size includes EXTENDED_ARG instructions - let instr_size = instr.arg.instr_size() as u32; + // instr_size includes EXTENDED_ARG and CACHE entries + let instr_size = instr.arg.instr_size() as u32 + instr.cache_entries; match (¤t_entry, instr.except_handler) { // No current entry, no handler - nothing to do @@ -1280,6 +1422,7 @@ fn push_cold_blocks_to_end(blocks: &mut Vec) { end_location: end_loc, except_handler: None, lineno_override: None, + cache_entries: 0, }); jump_block.next = blocks[cold_idx.idx()].next; blocks[cold_idx.idx()].next = jump_block_idx; @@ -1356,7 +1499,7 @@ fn normalize_jumps(blocks: &mut [Block]) { visited.fill(false); - for block_idx in visit_order { + for &block_idx in &visit_order { let idx = block_idx.idx(); visited[idx] = true; @@ -1392,6 +1535,7 @@ fn normalize_jumps(blocks: &mut [Block]) { end_location: ins.end_location, except_handler: ins.except_handler, lineno_override: None, + cache_entries: 0, }, )); } @@ -1401,6 +1545,52 @@ fn normalize_jumps(blocks: &mut [Block]) { blocks[idx].instructions.insert(pos, info); } } + + // Resolve JUMP/JUMP_NO_INTERRUPT pseudo instructions before offset fixpoint. + let mut block_order = vec![0u32; blocks.len()]; + for (pos, &block_idx) in visit_order.iter().enumerate() { + block_order[block_idx.idx()] = pos as u32; + } + + for &block_idx in &visit_order { + let source_pos = block_order[block_idx.idx()]; + for info in &mut blocks[block_idx.idx()].instructions { + let target = info.target; + if target == BlockIdx::NULL { + continue; + } + let target_pos = block_order[target.idx()]; + info.instr = match info.instr { + AnyInstruction::Pseudo(PseudoInstruction::Jump { .. }) => { + if target_pos > source_pos { + Instruction::JumpForward { + target: Arg::marker(), + } + .into() + } else { + Instruction::JumpBackward { + target: Arg::marker(), + } + .into() + } + } + AnyInstruction::Pseudo(PseudoInstruction::JumpNoInterrupt { .. }) => { + if target_pos > source_pos { + Instruction::JumpForward { + target: Arg::marker(), + } + .into() + } else { + Instruction::JumpBackwardNoInterrupt { + target: Arg::marker(), + } + .into() + } + } + other => other, + }; + } + } } /// 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 920754e8ec0..9dd78c6b7b2 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,17 +1,21 @@ --- source: crates/codegen/src/compile.rs +assertion_line: 9100 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 (10) - 3 NOT_TAKEN - 4 LOAD_CONST (False) - 5 POP_JUMP_IF_FALSE (10) - 6 NOT_TAKEN - 7 LOAD_CONST (False) - 8 POP_JUMP_IF_FALSE (10) - 9 NOT_TAKEN + >> 1 LOAD_CONST (True) + 2 POP_JUMP_IF_FALSE (9) + 3 CACHE + 4 NOT_TAKEN + >> 5 LOAD_CONST (False) + 6 POP_JUMP_IF_FALSE (5) + 7 CACHE + 8 NOT_TAKEN + >> 9 LOAD_CONST (False) + 10 POP_JUMP_IF_FALSE (1) + 11 CACHE + 12 NOT_TAKEN - 2 >> 10 LOAD_CONST (None) - 11 RETURN_VALUE + 2 13 LOAD_CONST (None) + 14 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 f8e18abf2cf..e9c3ad8a3c6 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,20 +1,25 @@ --- source: crates/codegen/src/compile.rs +assertion_line: 9110 expression: "compile_exec(\"\\\nif (True and False) or (False and True):\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_FALSE (7) - 3 NOT_TAKEN - 4 LOAD_CONST (False) - 5 POP_JUMP_IF_TRUE (13) - 6 NOT_TAKEN - >> 7 LOAD_CONST (False) - 8 POP_JUMP_IF_FALSE (13) - 9 NOT_TAKEN - 10 LOAD_CONST (True) - 11 POP_JUMP_IF_FALSE (13) + >> 1 LOAD_CONST (True) + 2 POP_JUMP_IF_FALSE (5) + 3 CACHE + 4 NOT_TAKEN + >> 5 LOAD_CONST (False) + 6 POP_JUMP_IF_TRUE (9) + 7 CACHE + 8 NOT_TAKEN + >> 9 LOAD_CONST (False) + 10 POP_JUMP_IF_FALSE (5) + 11 CACHE 12 NOT_TAKEN + 13 LOAD_CONST (True) + 14 POP_JUMP_IF_FALSE (1) + 15 CACHE + 16 NOT_TAKEN - 2 >> 13 LOAD_CONST (None) - 14 RETURN_VALUE + 2 17 LOAD_CONST (None) + 18 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 d0b7028e3d8..83212144b99 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,17 +1,21 @@ --- source: crates/codegen/src/compile.rs +assertion_line: 9090 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 (10) - 3 NOT_TAKEN - 4 LOAD_CONST (False) - 5 POP_JUMP_IF_TRUE (10) - 6 NOT_TAKEN - 7 LOAD_CONST (False) - 8 POP_JUMP_IF_FALSE (10) - 9 NOT_TAKEN + >> 1 LOAD_CONST (True) + 2 POP_JUMP_IF_TRUE (9) + 3 CACHE + 4 NOT_TAKEN + >> 5 LOAD_CONST (False) + 6 POP_JUMP_IF_TRUE (5) + 7 CACHE + 8 NOT_TAKEN + >> 9 LOAD_CONST (False) + 10 POP_JUMP_IF_FALSE (1) + 11 CACHE + 12 NOT_TAKEN - 2 >> 10 LOAD_CONST (None) - 11 RETURN_VALUE + 2 13 LOAD_CONST (None) + 14 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_bool_op.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_bool_op.snap index dea5f5386ff..a6db9ca4bdb 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_bool_op.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_bool_op.snap @@ -5,17 +5,22 @@ expression: "compile_exec(\"\\\nx = Test() and False or False\n\")" 1 0 RESUME (0) 1 LOAD_NAME (0, Test) 2 PUSH_NULL - 3 CALL (0) - 4 COPY (1) - 5 POP_JUMP_IF_FALSE (12) - 6 NOT_TAKEN - 7 POP_TOP - 8 LOAD_CONST (False) - 9 COPY (1) - 10 POP_JUMP_IF_TRUE (14) - 11 NOT_TAKEN - >> 12 POP_TOP - 13 LOAD_CONST (False) - >> 14 STORE_NAME (1, x) - 15 LOAD_CONST (None) - 16 RETURN_VALUE + >> 3 CALL (0) + 4 CACHE + 5 CACHE + 6 CACHE + >> 7 COPY (1) + 8 POP_JUMP_IF_FALSE (7) + 9 CACHE + 10 NOT_TAKEN + 11 POP_TOP + 12 LOAD_CONST (False) + 13 COPY (1) + 14 POP_JUMP_IF_TRUE (3) + 15 CACHE + 16 NOT_TAKEN + 17 POP_TOP + 18 LOAD_CONST (False) + 19 STORE_NAME (1, x) + 20 LOAD_CONST (None) + 21 RETURN_VALUE 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 0f48c2a90ea..7a1db8e7b8c 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,173 +1,276 @@ --- source: crates/codegen/src/compile.rs +assertion_line: 9089 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) 1 LOAD_CONST (): 1 0 RETURN_GENERATOR 1 POP_TOP - 2 RESUME (0) + >> 2 RESUME (0) - 2 3 LOAD_GLOBAL (0, StopIteration) - 4 PUSH_NULL - 5 LOAD_CONST ("spam") - 6 CALL (1) - 7 LOAD_GLOBAL (1, StopAsyncIteration) - 8 PUSH_NULL - 9 LOAD_CONST ("ham") - 10 CALL (1) - 11 BUILD_TUPLE (2) - 12 GET_ITER - >> 13 FOR_ITER (50) - 14 STORE_FAST (0, stop_exc) + 2 >> 3 LOAD_GLOBAL (1, NULL + StopIteration) + >> 4 CACHE + >> 5 CACHE + 6 CACHE + 7 CACHE + 8 LOAD_CONST ("spam") + 9 CALL (1) + >> 10 CACHE + 11 CACHE + 12 CACHE + 13 LOAD_GLOBAL (3, NULL + StopAsyncIteration) + 14 CACHE + 15 CACHE + 16 CACHE + 17 CACHE + >> 18 LOAD_CONST ("ham") + 19 CALL (1) + 20 CACHE + 21 CACHE + 22 CACHE + 23 BUILD_TUPLE (2) + 24 GET_ITER + 25 FOR_ITER (71) + 26 CACHE + >> 27 STORE_FAST (0, stop_exc) - 3 15 LOAD_GLOBAL (2, self) - 16 LOAD_ATTR (7, subTest, method=true) - 17 LOAD_GLOBAL (4, type) - 18 PUSH_NULL - 19 LOAD_FAST (0, stop_exc) - 20 CALL (1) - 21 LOAD_CONST (("type")) - 22 CALL_KW (1) - 23 COPY (1) - 24 LOAD_SPECIAL (__exit__) - 25 SWAP (2) - 26 LOAD_SPECIAL (__enter__) - 27 PUSH_NULL - 28 CALL (0) - 29 POP_TOP + 3 >> 28 LOAD_GLOBAL (4, self) + 29 CACHE + 30 CACHE + 31 CACHE + 32 CACHE + 33 LOAD_ATTR (7, subTest, method=true) + >> 34 CACHE + 35 CACHE + 36 CACHE + 37 CACHE + 38 CACHE + 39 CACHE + 40 CACHE + 41 CACHE + 42 CACHE + 43 LOAD_GLOBAL (9, NULL + type) + 44 CACHE + 45 CACHE + 46 CACHE + >> 47 CACHE + 48 LOAD_FAST (0, stop_exc) + 49 CALL (1) + 50 CACHE + 51 CACHE + 52 CACHE + 53 LOAD_CONST (("type")) + 54 CALL_KW (1) + 55 CACHE + 56 CACHE + 57 CACHE + 58 COPY (1) + 59 LOAD_SPECIAL (__exit__) + 60 SWAP (2) + 61 LOAD_SPECIAL (__enter__) + 62 PUSH_NULL + 63 CALL (0) + 64 CACHE + 65 CACHE + 66 CACHE + 67 POP_TOP - 4 30 NOP + 4 68 NOP - 5 31 LOAD_GLOBAL (5, egg) - 32 PUSH_NULL - 33 CALL (0) - 34 COPY (1) - 35 LOAD_SPECIAL (__aexit__) - 36 SWAP (2) - 37 LOAD_SPECIAL (__aenter__) - 38 PUSH_NULL - 39 CALL (0) - 40 GET_AWAITABLE (1) - 41 LOAD_CONST (None) - >> 42 SEND (46) - 43 YIELD_VALUE (1) - 44 RESUME (3) - 45 JUMP_BACKWARD_NO_INTERRUPT(42) - >> 46 END_SEND - 47 POP_TOP + 5 69 LOAD_GLOBAL (11, NULL + egg) + 70 CACHE + >> 71 CACHE + 72 CACHE + 73 CACHE + 74 CALL (0) + 75 CACHE + 76 CACHE + 77 CACHE + 78 COPY (1) + 79 LOAD_SPECIAL (__aexit__) + 80 SWAP (2) + 81 LOAD_SPECIAL (__aenter__) + 82 PUSH_NULL + 83 CALL (0) + 84 CACHE + 85 CACHE + 86 CACHE + 87 GET_AWAITABLE (1) + 88 LOAD_CONST (None) + 89 SEND (3) + 90 CACHE + 91 YIELD_VALUE (1) + 92 RESUME (3) + 93 JUMP_BACKWARD_NO_INTERRUPT(5) + 94 END_SEND + 95 POP_TOP - 6 48 LOAD_FAST (0, stop_exc) - 49 RAISE_VARARGS (Raise) + 6 96 LOAD_FAST (0, stop_exc) + 97 RAISE_VARARGS (Raise) - 2 >> 50 END_FOR - 51 POP_ITER - 52 LOAD_CONST (None) - 53 RETURN_VALUE + 2 98 END_FOR + 99 POP_ITER + 100 LOAD_CONST (None) + 101 RETURN_VALUE - 5 54 CLEANUP_THROW - 55 JUMP_BACKWARD_NO_INTERRUPT(46) + 5 102 CLEANUP_THROW + 103 JUMP_BACKWARD_NO_INTERRUPT(10) - 6 56 NOP + 6 104 NOP - 5 57 PUSH_NULL - 58 LOAD_CONST (None) - 59 LOAD_CONST (None) - 60 LOAD_CONST (None) - 61 CALL (3) - 62 GET_AWAITABLE (2) - 63 LOAD_CONST (None) - >> 64 SEND (69) - 65 YIELD_VALUE (1) - 66 RESUME (3) - 67 JUMP_BACKWARD_NO_INTERRUPT(64) - 68 CLEANUP_THROW - >> 69 END_SEND - 70 POP_TOP - 71 JUMP_FORWARD (94) - 72 PUSH_EXC_INFO - 73 WITH_EXCEPT_START - 74 GET_AWAITABLE (2) - 75 LOAD_CONST (None) - >> 76 SEND (81) - 77 YIELD_VALUE (1) - 78 RESUME (3) - 79 JUMP_BACKWARD_NO_INTERRUPT(76) - 80 CLEANUP_THROW - >> 81 END_SEND - 82 TO_BOOL - 83 POP_JUMP_IF_TRUE (86) - 84 NOT_TAKEN - 85 RERAISE (2) - >> 86 POP_TOP - 87 POP_EXCEPT - 88 POP_TOP - 89 POP_TOP - 90 JUMP_FORWARD (94) - 91 COPY (3) - 92 POP_EXCEPT - 93 RERAISE (1) - >> 94 JUMP_FORWARD (121) - 95 PUSH_EXC_INFO - - 7 96 LOAD_GLOBAL (6, Exception) - 97 CHECK_EXC_MATCH - 98 POP_JUMP_IF_FALSE (117) - 99 NOT_TAKEN - 100 STORE_FAST (1, ex) - - 8 101 LOAD_GLOBAL (2, self) - 102 LOAD_ATTR (15, assertIs, method=true) - 103 LOAD_FAST (1, ex) - 104 LOAD_FAST (0, stop_exc) - 105 CALL (2) - 106 POP_TOP - 107 JUMP_BACKWARD_NO_INTERRUPT(112) + 5 105 PUSH_NULL + 106 LOAD_CONST (None) + 107 LOAD_CONST (None) 108 LOAD_CONST (None) - 109 STORE_FAST (1, ex) - 110 DELETE_FAST (1, ex) - 111 RAISE_VARARGS (ReraiseFromStack) - >> 112 POP_EXCEPT - 113 LOAD_CONST (None) - 114 STORE_FAST (1, ex) - 115 DELETE_FAST (1, ex) - 116 JUMP_BACKWARD_NO_INTERRUPT(129) - >> 117 RAISE_VARARGS (ReraiseFromStack) - 118 COPY (3) - 119 POP_EXCEPT - 120 RAISE_VARARGS (ReraiseFromStack) - - 10 >> 121 LOAD_GLOBAL (2, self) - 122 LOAD_ATTR (17, fail, method=true) - 123 LOAD_FAST_BORROW (0, stop_exc) - 124 FORMAT_SIMPLE - 125 LOAD_CONST (" was suppressed") - 126 BUILD_STRING (2) - 127 CALL (1) - 128 POP_TOP - >> 129 NOP - - 3 130 PUSH_NULL - 131 LOAD_CONST (None) - 132 LOAD_CONST (None) - 133 LOAD_CONST (None) - 134 CALL (3) - 135 POP_TOP - 136 JUMP_FORWARD (151) - 137 PUSH_EXC_INFO - 138 WITH_EXCEPT_START - 139 TO_BOOL - 140 POP_JUMP_IF_TRUE (143) + 109 CALL (3) + 110 CACHE + 111 CACHE + 112 CACHE + 113 GET_AWAITABLE (2) + 114 LOAD_CONST (None) + 115 SEND (4) + 116 CACHE + 117 YIELD_VALUE (1) + 118 RESUME (3) + 119 JUMP_BACKWARD_NO_INTERRUPT(5) + 120 CLEANUP_THROW + 121 END_SEND + 122 POP_TOP + 123 JUMP_FORWARD (27) + 124 PUSH_EXC_INFO + 125 WITH_EXCEPT_START + 126 GET_AWAITABLE (2) + 127 LOAD_CONST (None) + 128 SEND (4) + 129 CACHE + 130 YIELD_VALUE (1) + 131 RESUME (3) + 132 JUMP_BACKWARD_NO_INTERRUPT(5) + 133 CLEANUP_THROW + 134 END_SEND + 135 TO_BOOL + 136 CACHE + 137 CACHE + 138 CACHE + 139 POP_JUMP_IF_TRUE (2) + 140 CACHE 141 NOT_TAKEN 142 RERAISE (2) - >> 143 POP_TOP + 143 POP_TOP 144 POP_EXCEPT 145 POP_TOP 146 POP_TOP - 147 JUMP_FORWARD (151) + 147 JUMP_FORWARD (3) 148 COPY (3) 149 POP_EXCEPT 150 RERAISE (1) - >> 151 JUMP_BACKWARD (13) + 151 JUMP_FORWARD (47) + 152 PUSH_EXC_INFO + + 7 153 LOAD_GLOBAL (12, Exception) + 154 CACHE + 155 CACHE + 156 CACHE + 157 CACHE + 158 CHECK_EXC_MATCH + 159 POP_JUMP_IF_FALSE (34) + 160 CACHE + 161 NOT_TAKEN + 162 STORE_FAST (1, ex) + + 8 163 LOAD_GLOBAL (4, self) + 164 CACHE + 165 CACHE + 166 CACHE + 167 CACHE + 168 LOAD_ATTR (15, assertIs, method=true) + 169 CACHE + 170 CACHE + 171 CACHE + 172 CACHE + 173 CACHE + 174 CACHE + 175 CACHE + 176 CACHE + 177 CACHE + 178 LOAD_FAST (1, ex) + 179 LOAD_FAST (0, stop_exc) + 180 CALL (2) + 181 CACHE + 182 CACHE + 183 CACHE + 184 POP_TOP + 185 JUMP_FORWARD (4) + 186 LOAD_CONST (None) + 187 STORE_FAST (1, ex) + 188 DELETE_FAST (1, ex) + 189 RERAISE (1) + 190 POP_EXCEPT + 191 LOAD_CONST (None) + 192 STORE_FAST (1, ex) + 193 DELETE_FAST (1, ex) + 194 JUMP_FORWARD (28) + 195 RERAISE (0) + 196 COPY (3) + 197 POP_EXCEPT + 198 RERAISE (1) + + 10 199 LOAD_GLOBAL (4, self) + 200 CACHE + 201 CACHE + 202 CACHE + 203 CACHE + 204 LOAD_ATTR (17, fail, method=true) + 205 CACHE + 206 CACHE + 207 CACHE + 208 CACHE + 209 CACHE + 210 CACHE + 211 CACHE + 212 CACHE + 213 CACHE + 214 LOAD_FAST_BORROW (0, stop_exc) + 215 FORMAT_SIMPLE + 216 LOAD_CONST (" was suppressed") + 217 BUILD_STRING (2) + 218 CALL (1) + 219 CACHE + 220 CACHE + 221 CACHE + 222 POP_TOP + 223 NOP + + 3 224 PUSH_NULL + 225 LOAD_CONST (None) + 226 LOAD_CONST (None) + 227 LOAD_CONST (None) + 228 CALL (3) + >> 229 CACHE + 230 CACHE + 231 CACHE + 232 POP_TOP + 233 JUMP_FORWARD (18) + 234 PUSH_EXC_INFO + 235 WITH_EXCEPT_START + 236 TO_BOOL + 237 CACHE + 238 CACHE + 239 CACHE + 240 POP_JUMP_IF_TRUE (2) + 241 CACHE + 242 NOT_TAKEN + 243 RERAISE (2) + 244 POP_TOP + 245 POP_EXCEPT + 246 POP_TOP + 247 POP_TOP + 248 JUMP_FORWARD (3) + 249 COPY (3) + 250 POP_EXCEPT + 251 RERAISE (1) + 252 JUMP_BACKWARD (229) + 253 CACHE 2 MAKE_FUNCTION 3 STORE_NAME (0, test) diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index 8ba71381dbc..c1c5e8cd847 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -511,6 +511,125 @@ impl Instruction { _ => return None, }) } + + /// Number of CACHE code units that follow this instruction. + /// _PyOpcode_Caches + pub fn cache_entries(self) -> usize { + match self { + // LOAD_ATTR: 9 cache entries + Self::LoadAttr { .. } + | Self::LoadAttrClass + | Self::LoadAttrClassWithMetaclassCheck + | Self::LoadAttrGetattributeOverridden + | Self::LoadAttrInstanceValue + | Self::LoadAttrMethodLazyDict + | Self::LoadAttrMethodNoDict + | Self::LoadAttrMethodWithValues + | Self::LoadAttrModule + | Self::LoadAttrNondescriptorNoDict + | Self::LoadAttrNondescriptorWithValues + | Self::LoadAttrProperty + | Self::LoadAttrSlot + | Self::LoadAttrWithHint => 9, + + // BINARY_OP: 5 cache entries + Self::BinaryOp { .. } + | Self::BinaryOpAddFloat + | Self::BinaryOpAddInt + | Self::BinaryOpAddUnicode + | Self::BinaryOpExtend + | Self::BinaryOpInplaceAddUnicode + | Self::BinaryOpMultiplyFloat + | Self::BinaryOpMultiplyInt + | Self::BinaryOpSubscrDict + | Self::BinaryOpSubscrGetitem + | Self::BinaryOpSubscrListInt + | Self::BinaryOpSubscrListSlice + | Self::BinaryOpSubscrStrInt + | Self::BinaryOpSubscrTupleInt + | Self::BinaryOpSubtractFloat + | Self::BinaryOpSubtractInt => 5, + + // LOAD_GLOBAL / STORE_ATTR: 4 cache entries + Self::LoadGlobal(_) + | Self::LoadGlobalBuiltin + | Self::LoadGlobalModule + | Self::StoreAttr { .. } + | Self::StoreAttrInstanceValue + | Self::StoreAttrSlot + | Self::StoreAttrWithHint => 4, + + // CALL / CALL_KW / TO_BOOL: 3 cache entries + Self::Call { .. } + | Self::CallAllocAndEnterInit + | Self::CallBoundMethodExactArgs + | Self::CallBoundMethodGeneral + | Self::CallBuiltinClass + | Self::CallBuiltinFast + | Self::CallBuiltinFastWithKeywords + | Self::CallBuiltinO + | Self::CallIsinstance + | Self::CallLen + | Self::CallListAppend + | Self::CallMethodDescriptorFast + | Self::CallMethodDescriptorFastWithKeywords + | Self::CallMethodDescriptorNoargs + | Self::CallMethodDescriptorO + | Self::CallNonPyGeneral + | Self::CallPyExactArgs + | Self::CallPyGeneral + | Self::CallStr1 + | Self::CallTuple1 + | Self::CallType1 + | Self::CallKw { .. } + | Self::CallKwBoundMethod + | Self::CallKwNonPy + | Self::CallKwPy + | Self::ToBool + | Self::ToBoolAlwaysTrue + | Self::ToBoolBool + | Self::ToBoolInt + | Self::ToBoolList + | Self::ToBoolNone + | Self::ToBoolStr => 3, + + // 1 cache entry + Self::CompareOp { .. } + | Self::CompareOpFloat + | Self::CompareOpInt + | Self::CompareOpStr + | Self::ContainsOp(_) + | Self::ContainsOpDict + | Self::ContainsOpSet + | Self::ForIter { .. } + | Self::ForIterGen + | Self::ForIterList + | Self::ForIterRange + | Self::ForIterTuple + | Self::JumpBackward { .. } + | Self::JumpBackwardJit + | Self::JumpBackwardNoJit + | Self::LoadSuperAttr { .. } + | Self::LoadSuperAttrAttr + | Self::LoadSuperAttrMethod + | Self::PopJumpIfTrue { .. } + | Self::PopJumpIfFalse { .. } + | Self::PopJumpIfNone { .. } + | Self::PopJumpIfNotNone { .. } + | Self::Send { .. } + | Self::SendGen + | Self::StoreSubscr + | Self::StoreSubscrDict + | Self::StoreSubscrListInt + | Self::UnpackSequence { .. } + | Self::UnpackSequenceList + | Self::UnpackSequenceTuple + | Self::UnpackSequenceTwoTuple => 1, + + // Everything else: 0 cache entries + _ => 0, + } + } } impl InstructionMetadata for Instruction { @@ -708,10 +827,7 @@ impl InstructionMetadata for Instruction { Self::LoadFastLoadFast { .. } => (2, 0), Self::LoadFromDictOrDeref(_) => (1, 1), Self::LoadFromDictOrGlobals(_) => (1, 1), - Self::LoadGlobal(_) => ( - 1, // TODO: Differs from CPython `1 + (oparg & 1)` - 0, - ), + Self::LoadGlobal(_) => (1 + (oparg & 1), 0), Self::LoadGlobalBuiltin => (1 + (oparg & 1), 0), Self::LoadGlobalModule => (1 + (oparg & 1), 0), Self::LoadLocals => (1, 0), @@ -739,28 +855,13 @@ impl InstructionMetadata for Instruction { Self::PopTop => (0, 1), Self::PushExcInfo => (2, 1), Self::PushNull => (1, 0), - Self::RaiseVarargs { kind } => ( - 0, - // TODO: Differs from CPython: `oparg` - match kind.get((oparg as u32).into()) { - RaiseKind::BareRaise => 0, - RaiseKind::Raise => 1, - RaiseKind::RaiseCause => 2, - RaiseKind::ReraiseFromStack => 1, - }, - ), - Self::Reraise { .. } => ( - 1 + oparg, // TODO: Differs from CPython: `oparg` - 1 + oparg, - ), + Self::RaiseVarargs { .. } => (0, oparg), + Self::Reraise { .. } => (oparg, 1 + oparg), Self::Reserved => (0, 0), Self::Resume { .. } => (0, 0), Self::ResumeCheck => (0, 0), Self::ReturnGenerator => (1, 0), - Self::ReturnValue => ( - 0, // TODO: Differs from CPython: `1` - 1, - ), + Self::ReturnValue => (1, 1), Self::Send { .. } => (2, 2), Self::SendGen => (1, 2), Self::SetAdd { .. } => (1 + (oparg - 1), 2 + (oparg - 1)), @@ -860,7 +961,9 @@ impl InstructionMetadata for Instruction { }; match self { + Self::BinarySlice => w!(BINARY_SLICE), Self::BinaryOp { op } => write!(f, "{:pad$}({})", "BINARY_OP", op.get(arg)), + Self::BinaryOpInplaceAddUnicode => w!(BINARY_OP_INPLACE_ADD_UNICODE), Self::BuildList { size } => w!(BUILD_LIST, size), Self::BuildMap { size } => w!(BUILD_MAP, size), Self::BuildSet { size } => w!(BUILD_SET, size), @@ -872,6 +975,7 @@ impl InstructionMetadata for Instruction { Self::CallKw { nargs } => w!(CALL_KW, nargs), Self::CallIntrinsic1 { func } => w!(CALL_INTRINSIC_1, ?func), Self::CallIntrinsic2 { func } => w!(CALL_INTRINSIC_2, ?func), + Self::Cache => w!(CACHE), Self::CheckEgMatch => w!(CHECK_EG_MATCH), Self::CheckExcMatch => w!(CHECK_EXC_MATCH), Self::CleanupThrow => w!(CLEANUP_THROW), @@ -879,6 +983,7 @@ impl InstructionMetadata for Instruction { Self::ContainsOp(inv) => w!(CONTAINS_OP, ?inv), Self::ConvertValue { oparg } => write!(f, "{:pad$}{}", "CONVERT_VALUE", oparg.get(arg)), Self::Copy { index } => w!(COPY, index), + Self::CopyFreeVars { count } => w!(COPY_FREE_VARS, count), Self::DeleteAttr { idx } => w!(DELETE_ATTR, name = idx), Self::DeleteDeref(idx) => w!(DELETE_DEREF, cell_name = idx), Self::DeleteFast(idx) => w!(DELETE_FAST, varname = idx), @@ -890,6 +995,7 @@ impl InstructionMetadata for Instruction { Self::EndAsyncFor => w!(END_ASYNC_FOR), Self::EndSend => w!(END_SEND), Self::ExtendedArg => w!(EXTENDED_ARG, Arg::::marker()), + Self::ExitInitCheck => w!(EXIT_INIT_CHECK), Self::ForIter { target } => w!(FOR_ITER, target), Self::FormatSimple => w!(FORMAT_SIMPLE), Self::FormatWithSpec => w!(FORMAT_WITH_SPEC), @@ -901,6 +1007,7 @@ impl InstructionMetadata for Instruction { Self::GetLen => w!(GET_LEN), Self::ImportFrom { idx } => w!(IMPORT_FROM, name = idx), Self::ImportName { idx } => w!(IMPORT_NAME, name = idx), + Self::InterpreterExit => w!(INTERPRETER_EXIT), Self::IsOp(inv) => w!(IS_OP, ?inv), Self::JumpBackward { target } => w!(JUMP_BACKWARD, target), Self::JumpBackwardNoInterrupt { target } => w!(JUMP_BACKWARD_NO_INTERRUPT, target), @@ -922,6 +1029,7 @@ impl InstructionMetadata for Instruction { } } Self::LoadBuildClass => w!(LOAD_BUILD_CLASS), + Self::LoadCommonConstant { idx } => w!(LOAD_COMMON_CONSTANT, ?idx), Self::LoadFromDictOrDeref(i) => w!(LOAD_FROM_DICT_OR_DEREF, cell_name = i), Self::LoadConst { idx } => fmt_const("LOAD_CONST", arg, f, idx), Self::LoadSmallInt { idx } => w!(LOAD_SMALL_INT, idx), @@ -951,7 +1059,64 @@ impl InstructionMetadata for Instruction { ) } Self::LoadFromDictOrGlobals(idx) => w!(LOAD_FROM_DICT_OR_GLOBALS, name = idx), - Self::LoadGlobal(idx) => w!(LOAD_GLOBAL, name = idx), + Self::LoadGlobal(idx) => { + let oparg = idx.get(arg); + let name_idx = oparg >> 1; + if (oparg & 1) != 0 { + write!( + f, + "{:pad$}({}, NULL + {})", + "LOAD_GLOBAL", + oparg, + name(name_idx) + ) + } else { + write!(f, "{:pad$}({}, {})", "LOAD_GLOBAL", oparg, name(name_idx)) + } + } + Self::LoadGlobalBuiltin => { + let oparg = u32::from(arg); + let name_idx = oparg >> 1; + if (oparg & 1) != 0 { + write!( + f, + "{:pad$}({}, NULL + {})", + "LOAD_GLOBAL_BUILTIN", + oparg, + name(name_idx) + ) + } else { + write!( + f, + "{:pad$}({}, {})", + "LOAD_GLOBAL_BUILTIN", + oparg, + name(name_idx) + ) + } + } + Self::LoadGlobalModule => { + let oparg = u32::from(arg); + let name_idx = oparg >> 1; + if (oparg & 1) != 0 { + write!( + f, + "{:pad$}({}, NULL + {})", + "LOAD_GLOBAL_MODULE", + oparg, + name(name_idx) + ) + } else { + write!( + f, + "{:pad$}({}, {})", + "LOAD_GLOBAL_MODULE", + oparg, + name(name_idx) + ) + } + } + Self::LoadLocals => w!(LOAD_LOCALS), Self::LoadName(idx) => w!(LOAD_NAME, name = idx), Self::LoadSpecial { method } => w!(LOAD_SPECIAL, method), Self::LoadSuperAttr { arg: idx } => { @@ -966,6 +1131,7 @@ impl InstructionMetadata for Instruction { oparg.has_class() ) } + Self::MakeCell(idx) => w!(MAKE_CELL, cell_name = idx), Self::MakeFunction => w!(MAKE_FUNCTION), Self::MapAdd { i } => w!(MAP_ADD, i), Self::MatchClass(arg) => w!(MATCH_CLASS, arg), @@ -976,6 +1142,8 @@ impl InstructionMetadata for Instruction { Self::NotTaken => w!(NOT_TAKEN), Self::PopExcept => w!(POP_EXCEPT), Self::PopJumpIfFalse { target } => w!(POP_JUMP_IF_FALSE, target), + Self::PopJumpIfNone { target } => w!(POP_JUMP_IF_NONE, target), + Self::PopJumpIfNotNone { target } => w!(POP_JUMP_IF_NOT_NONE, target), Self::PopJumpIfTrue { target } => w!(POP_JUMP_IF_TRUE, target), Self::PopTop => w!(POP_TOP), Self::EndFor => w!(END_FOR), @@ -1002,8 +1170,21 @@ impl InstructionMetadata for Instruction { write!(f, "STORE_FAST_LOAD_FAST")?; write!(f, " ({}, {})", store_idx, load_idx) } + Self::StoreFastStoreFast { arg: packed } => { + let oparg = packed.get(arg); + let idx1 = oparg >> 4; + let idx2 = oparg & 15; + write!( + f, + "{:pad$}({}, {})", + "STORE_FAST_STORE_FAST", + varname(idx1), + varname(idx2) + ) + } Self::StoreGlobal(idx) => w!(STORE_GLOBAL, name = idx), Self::StoreName(idx) => w!(STORE_NAME, name = idx), + Self::StoreSlice => w!(STORE_SLICE), Self::StoreSubscr => w!(STORE_SUBSCR), Self::Swap { index } => w!(SWAP, index), Self::ToBool => w!(TO_BOOL), diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 5e760d05449..2fa72c6e375 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -1,5 +1,6 @@ // spell-checker: disable use super::{JitCompileError, JitSig, JitType}; +use alloc::collections::BTreeSet; use cranelift::codegen::ir::FuncRef; use cranelift::prelude::*; use num_traits::cast::ToPrimitive; @@ -154,12 +155,69 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { .or_insert_with(|| builder.create_block()) } + fn jump_target_forward(offset: u32, caches: u32, arg: OpArg) -> Result { + let after = offset + .checked_add(1) + .and_then(|i| i.checked_add(caches)) + .ok_or(JitCompileError::BadBytecode)?; + let target = after + .checked_add(u32::from(arg)) + .ok_or(JitCompileError::BadBytecode)?; + Ok(Label(target)) + } + + fn jump_target_backward( + offset: u32, + caches: u32, + arg: OpArg, + ) -> Result { + let after = offset + .checked_add(1) + .and_then(|i| i.checked_add(caches)) + .ok_or(JitCompileError::BadBytecode)?; + let target = after + .checked_sub(u32::from(arg)) + .ok_or(JitCompileError::BadBytecode)?; + Ok(Label(target)) + } + + fn instruction_target( + offset: u32, + instruction: Instruction, + arg: OpArg, + ) -> Result, JitCompileError> { + let caches = instruction.cache_entries() as u32; + let target = match instruction { + Instruction::JumpForward { .. } => { + Some(Self::jump_target_forward(offset, caches, arg)?) + } + Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } => { + Some(Self::jump_target_backward(offset, caches, arg)?) + } + Instruction::PopJumpIfFalse { .. } + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfNone { .. } + | Instruction::PopJumpIfNotNone { .. } + | Instruction::ForIter { .. } + | Instruction::Send { .. } => Some(Self::jump_target_forward(offset, caches, arg)?), + _ => None, + }; + Ok(target) + } + pub fn compile( &mut self, func_ref: FuncRef, bytecode: &CodeObject, ) -> Result<(), JitCompileError> { - let label_targets = bytecode.label_targets(); + let mut label_targets = BTreeSet::new(); + let mut target_arg_state = OpArgState::default(); + for (offset, &raw_instr) in bytecode.instructions.iter().enumerate() { + let (instruction, arg) = target_arg_state.get(raw_instr); + if let Some(target) = Self::instruction_target(offset as u32, instruction, arg)? { + label_targets.insert(target); + } + } let mut arg_state = OpArgState::default(); // Track whether we have "returned" in the current block @@ -206,7 +264,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } // Actually compile this instruction: - self.add_instruction(func_ref, bytecode, instruction, arg)?; + self.add_instruction(func_ref, bytecode, offset as u32, instruction, arg)?; // If that was an unconditional branch or return, mark future instructions unreachable match instruction { @@ -288,6 +346,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { &mut self, func_ref: FuncRef, bytecode: &CodeObject, + offset: u32, instruction: Instruction, arg: OpArg, ) -> Result<(), JitCompileError> { @@ -557,12 +616,14 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { _ => Err(JitCompileError::NotSupported), } } - Instruction::ExtendedArg => Ok(()), + Instruction::ExtendedArg | Instruction::Cache => Ok(()), - Instruction::JumpBackward { target } - | Instruction::JumpBackwardNoInterrupt { target } - | Instruction::JumpForward { target } => { - let target_block = self.get_or_create_block(target.get(arg)); + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. } + | Instruction::JumpForward { .. } => { + let target = Self::instruction_target(offset, instruction, arg)? + .ok_or(JitCompileError::BadBytecode)?; + let target_block = self.get_or_create_block(target); self.builder.ins().jump(target_block, &[]); Ok(()) } @@ -605,20 +666,26 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Ok(()) } Instruction::LoadGlobal(idx) => { - let name = &bytecode.names[idx.get(arg) as usize]; + let oparg = idx.get(arg); + let name = &bytecode.names[(oparg >> 1) as usize]; if name.as_ref() != bytecode.obj_name.as_ref() { Err(JitCompileError::NotSupported) } else { self.stack.push(JitValue::FuncRef(func_ref)); + if (oparg & 1) != 0 { + self.stack.push(JitValue::Null); + } Ok(()) } } Instruction::Nop | Instruction::NotTaken => Ok(()), - Instruction::PopJumpIfFalse { target } => { + Instruction::PopJumpIfFalse { .. } => { let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; let val = self.boolean_val(cond)?; - let then_block = self.get_or_create_block(target.get(arg)); + let then_label = Self::instruction_target(offset, instruction, arg)? + .ok_or(JitCompileError::BadBytecode)?; + let then_block = self.get_or_create_block(then_label); let else_block = self.builder.create_block(); self.builder @@ -628,10 +695,12 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Ok(()) } - Instruction::PopJumpIfTrue { target } => { + Instruction::PopJumpIfTrue { .. } => { let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; let val = self.boolean_val(cond)?; - let then_block = self.get_or_create_block(target.get(arg)); + let then_label = Self::instruction_target(offset, instruction, arg)? + .ok_or(JitCompileError::BadBytecode)?; + let then_block = self.get_or_create_block(then_label); let else_block = self.builder.create_block(); self.builder diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index d47aac60aec..79b47c1bb16 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -79,9 +79,12 @@ fn extract_annotations_from_annotate_code(code: &CodeObject) -> HashMap { stack.push((true, idx.get(arg) as usize)); } - Instruction::LoadName(idx) | Instruction::LoadGlobal(idx) => { + Instruction::LoadName(idx) => { stack.push((false, idx.get(arg) as usize)); } + Instruction::LoadGlobal(idx) => { + stack.push((false, (idx.get(arg) >> 1) as usize)); + } Instruction::BuildMap { size, .. } => { let count = size.get(arg) as usize; // Stack has key-value pairs in order: k1, v1, k2, v2, ... @@ -139,6 +142,7 @@ fn extract_annotations_from_annotate_code(code: &CodeObject) -> HashMap { // Ignore these instructions for annotation extraction } @@ -185,7 +189,7 @@ impl StackMachine { names: &[String], ) -> ControlFlow<()> { match instruction { - Instruction::Resume { .. } | Instruction::NotTaken => { + Instruction::Resume { .. } | Instruction::Cache | Instruction::NotTaken => { // No-op for JIT tests } Instruction::LoadConst { idx } => { diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index 18d40a818f2..1708477004e 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -932,11 +932,14 @@ impl PyCode { let oparg = extended_arg | raw_arg; extended_arg = 0; + let caches = op.cache_entries(); let (src, left, right) = match op { Instruction::ForIter { .. } => { - // left = fall-through (continue iteration) + // left = fall-through past CACHE entries (continue iteration) // right = past END_FOR (iterator exhausted, skip cleanup) - let target = oparg as usize; + // arg is relative forward from after instruction+caches + let after_cache = i + 1 + caches; + let target = after_cache + oparg as usize; let right = if matches!( instructions.get(target).map(|u| u.op), Some(Instruction::EndFor) | Some(Instruction::InstrumentedEndFor) @@ -945,23 +948,25 @@ impl PyCode { } else { target * 2 }; - (i * 2, (i + 1) * 2, right) + (i * 2, after_cache * 2, right) } Instruction::PopJumpIfFalse { .. } | Instruction::PopJumpIfTrue { .. } | Instruction::PopJumpIfNone { .. } | Instruction::PopJumpIfNotNone { .. } => { - // left = fall-through (skip NOT_TAKEN if present) - // right = jump target (condition met) + // left = fall-through past CACHE entries (skip NOT_TAKEN if present) + // right = jump target (relative forward from after instruction+caches) + let after_cache = i + 1 + caches; let next_op = instructions - .get(i + 1) + .get(after_cache) .map(|u| u.op.to_base().unwrap_or(u.op)); let fallthrough = if matches!(next_op, Some(Instruction::NotTaken)) { - (i + 2) * 2 + (after_cache + 1) * 2 } else { - (i + 1) * 2 + after_cache * 2 }; - (i * 2, fallthrough, oparg as usize * 2) + let right_target = after_cache + oparg as usize; + (i * 2, fallthrough, right_target * 2) } Instruction::EndAsyncFor => { // src = END_SEND position (next_i - oparg) @@ -969,6 +974,7 @@ impl PyCode { let Some(src_i) = next_i.checked_sub(oparg as usize) else { continue; }; + // left = fall-through past NOT_TAKEN (src_i * 2, (src_i + 2) * 2, next_i * 2) } _ => continue, diff --git a/crates/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index 67c49962a91..9b057b5b806 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -185,7 +185,8 @@ pub(crate) mod stack_analysis { // De-instrument: get the underlying real instruction let opcode = opcode.to_base().unwrap_or(opcode); - let next_i = i + 1; // No inline caches in RustPython + let caches = opcode.cache_entries(); + let next_i = i + 1 + caches; if next_stack == UNINITIALIZED { i = next_i; @@ -197,8 +198,8 @@ pub(crate) mod stack_analysis { | Instruction::PopJumpIfTrue { .. } | Instruction::PopJumpIfNone { .. } | Instruction::PopJumpIfNotNone { .. } => { - // Jump target is absolute instruction index - let j = oparg as usize; + // Relative forward: target = after_caches + delta + let j = next_i + oparg as usize; next_stack = pop_value(next_stack); let target_stack = next_stack; if j < stacks.len() && stacks[j] == UNINITIALIZED { @@ -209,8 +210,8 @@ pub(crate) mod stack_analysis { } } Instruction::Send { .. } => { - // target is absolute - let j = oparg as usize; + // Relative forward: target = after_caches + delta + let j = next_i + oparg as usize; if j < stacks.len() && stacks[j] == UNINITIALIZED { stacks[j] = next_stack; } @@ -219,16 +220,16 @@ pub(crate) mod stack_analysis { } } Instruction::JumpForward { .. } => { - // target is absolute in RustPython - let j = oparg as usize; + // Relative forward: target = after_caches + delta + let j = next_i + oparg as usize; if j < stacks.len() && stacks[j] == UNINITIALIZED { stacks[j] = next_stack; } } Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } => { - // target is absolute in RustPython - let j = oparg as usize; + // Relative backward: target = after_caches - delta + let j = next_i - oparg as usize; if j < stacks.len() && stacks[j] == UNINITIALIZED { stacks[j] = next_stack; if j < i { @@ -248,10 +249,8 @@ pub(crate) mod stack_analysis { if next_i < stacks.len() { stacks[next_i] = body_stack; } - // Exhaustion path: execute_for_iter skips END_FOR and - // jumps directly to POP_ITER. The iterator stays on - // the stack and POP_ITER removes it. - let mut j = oparg as usize; + // Exhaustion path: relative forward from after_caches + let mut j = next_i + oparg as usize; if j < instructions.len() { let target_op = instructions[j].op.to_base().unwrap_or(instructions[j].op); @@ -297,9 +296,10 @@ pub(crate) mod stack_analysis { } } Instruction::LoadGlobal(_) => { - // RustPython's LOAD_GLOBAL doesn't encode push_null in oparg - // (separate PUSH_NULL instructions are used instead) next_stack = push_value(next_stack, Kind::Object as i64); + if oparg & 1 != 0 { + next_stack = push_value(next_stack, Kind::Null as i64); + } if next_i < stacks.len() { stacks[next_i] = next_stack; } diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 5224a54ed54..13825553fb2 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -578,7 +578,9 @@ impl ExecutingFrame<'_> { self.prev_line = loc.line.get() as u32; } + let lasti_before = self.lasti(); let result = self.execute_instruction(op, arg, &mut do_extend_arg, vm); + self.skip_caches_if_fallthrough(op, lasti_before); match result { Ok(None) => {} Ok(Some(value)) => { @@ -747,12 +749,15 @@ impl ExecutingFrame<'_> { if let Some(unit) = self.code.instructions.get(lasti) { match &unit.op { Instruction::Send { .. } => return Some(self.top_value()), - Instruction::Resume { .. } => { + Instruction::Resume { .. } | Instruction::InstrumentedResume => { // Check if previous instruction was YIELD_VALUE with arg >= 1 // This indicates yield-from/await context if lasti > 0 && let Some(prev_unit) = self.code.instructions.get(lasti - 1) - && let Instruction::YieldValue { .. } = &prev_unit.op + && matches!( + &prev_unit.op, + Instruction::YieldValue { .. } | Instruction::InstrumentedYieldValue + ) { // YIELD_VALUE arg: 0 = direct yield, >= 1 = yield-from/await // OpArgByte.0 is the raw byte value @@ -1341,8 +1346,10 @@ impl ExecutingFrame<'_> { *extend_arg = true; Ok(None) } - Instruction::ForIter { target } => { - self.execute_for_iter(vm, target.get(arg))?; + Instruction::ForIter { .. } => { + // Relative forward jump: target = lasti + caches + delta + let target = bytecode::Label(self.lasti() + 1 + u32::from(arg)); + self.execute_for_iter(vm, target)?; Ok(None) } Instruction::FormatSimple => { @@ -1515,16 +1522,16 @@ impl ExecutingFrame<'_> { self.push_value(vm.ctx.new_bool(value).into()); Ok(None) } - Instruction::JumpForward { target } => { - self.jump(target.get(arg)); + Instruction::JumpForward { .. } => { + self.jump_relative_forward(u32::from(arg), 0); Ok(None) } - Instruction::JumpBackward { target } => { - self.jump(target.get(arg)); + Instruction::JumpBackward { .. } => { + self.jump_relative_backward(u32::from(arg), 1); Ok(None) } - Instruction::JumpBackwardNoInterrupt { target } => { - self.jump(target.get(arg)); + Instruction::JumpBackwardNoInterrupt { .. } => { + self.jump_relative_backward(u32::from(arg), 0); Ok(None) } Instruction::ListAppend { i } => { @@ -1815,9 +1822,13 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::LoadGlobal(idx) => { - let name = &self.code.names[idx.get(arg) as usize]; + let oparg = idx.get(arg); + let name = &self.code.names[(oparg >> 1) as usize]; let x = self.load_global_or_builtin(name, vm)?; self.push_value(x); + if (oparg & 1) != 0 { + self.push_value_opt(None); + } Ok(None) } Instruction::LoadName(idx) => { @@ -2119,19 +2130,19 @@ impl ExecutingFrame<'_> { Ok(None) } - Instruction::PopJumpIfFalse { target } => self.pop_jump_if(vm, target.get(arg), false), - Instruction::PopJumpIfTrue { target } => self.pop_jump_if(vm, target.get(arg), true), - Instruction::PopJumpIfNone { target } => { + Instruction::PopJumpIfFalse { .. } => self.pop_jump_if_relative(vm, arg, 1, false), + Instruction::PopJumpIfTrue { .. } => self.pop_jump_if_relative(vm, arg, 1, true), + Instruction::PopJumpIfNone { .. } => { let value = self.pop_value(); if vm.is_none(&value) { - self.jump(target.get(arg)); + self.jump_relative_forward(u32::from(arg), 1); } Ok(None) } - Instruction::PopJumpIfNotNone { target } => { + Instruction::PopJumpIfNotNone { .. } => { let value = self.pop_value(); if !vm.is_none(&value) { - self.jump(target.get(arg)); + self.jump_relative_forward(u32::from(arg), 1); } Ok(None) } @@ -2409,12 +2420,13 @@ impl ExecutingFrame<'_> { }; Ok(Some(ExecutionResult::Yield(value))) } - Instruction::Send { target } => { + Instruction::Send { .. } => { // (receiver, v -- receiver, retval) // Pops v, sends it to receiver. On yield, pushes retval // (so stack = [..., receiver, retval]). On return/StopIteration, // also pushes retval and jumps to END_SEND which will pop receiver. - let exit_label = target.get(arg); + // Relative forward: target = lasti + caches(1) + delta + let exit_label = bytecode::Label(self.lasti() + 1 + u32::from(arg)); let val = self.pop_value(); let receiver = self.top_value(); @@ -2637,9 +2649,20 @@ impl ExecutingFrame<'_> { } } } - Instruction::InstrumentedJumpForward | Instruction::InstrumentedJumpBackward => { + Instruction::InstrumentedJumpForward => { + let src_offset = (self.lasti() - 1) * 2; + let target_idx = self.lasti() + u32::from(arg); + let target = bytecode::Label(target_idx); + self.jump(target); + if self.monitoring_mask & monitoring::EVENT_JUMP != 0 { + monitoring::fire_jump(vm, self.code, src_offset, target.0 * 2)?; + } + Ok(None) + } + Instruction::InstrumentedJumpBackward => { let src_offset = (self.lasti() - 1) * 2; - let target = bytecode::Label::from(u32::from(arg)); + let target_idx = self.lasti() + 1 - u32::from(arg); + let target = bytecode::Label(target_idx); self.jump(target); if self.monitoring_mask & monitoring::EVENT_JUMP != 0 { monitoring::fire_jump(vm, self.code, src_offset, target.0 * 2)?; @@ -2648,11 +2671,11 @@ impl ExecutingFrame<'_> { } Instruction::InstrumentedForIter => { let src_offset = (self.lasti() - 1) * 2; - let target = bytecode::Label::from(u32::from(arg)); + let target = bytecode::Label(self.lasti() + 1 + u32::from(arg)); let continued = self.execute_for_iter(vm, target)?; if continued { if self.monitoring_mask & monitoring::EVENT_BRANCH_LEFT != 0 { - let dest_offset = self.lasti() * 2; + let dest_offset = (self.lasti() + 1) * 2; // after caches monitoring::fire_branch_left(vm, self.code, src_offset, dest_offset)?; } } else if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { @@ -2694,64 +2717,67 @@ impl ExecutingFrame<'_> { } Instruction::InstrumentedPopJumpIfTrue => { let src_offset = (self.lasti() - 1) * 2; - let target = bytecode::Label::from(u32::from(arg)); + let target_idx = self.lasti() + 1 + u32::from(arg); let obj = self.pop_value(); let value = obj.try_to_bool(vm)?; if value { - self.jump(target); + self.jump(bytecode::Label(target_idx)); if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { - monitoring::fire_branch_right(vm, self.code, src_offset, target.0 * 2)?; + monitoring::fire_branch_right(vm, self.code, src_offset, target_idx * 2)?; } } Ok(None) } Instruction::InstrumentedPopJumpIfFalse => { let src_offset = (self.lasti() - 1) * 2; - let target = bytecode::Label::from(u32::from(arg)); + let target_idx = self.lasti() + 1 + u32::from(arg); let obj = self.pop_value(); let value = obj.try_to_bool(vm)?; if !value { - self.jump(target); + self.jump(bytecode::Label(target_idx)); if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { - monitoring::fire_branch_right(vm, self.code, src_offset, target.0 * 2)?; + monitoring::fire_branch_right(vm, self.code, src_offset, target_idx * 2)?; } } Ok(None) } Instruction::InstrumentedPopJumpIfNone => { let src_offset = (self.lasti() - 1) * 2; + let target_idx = self.lasti() + 1 + u32::from(arg); let value = self.pop_value(); - let target = bytecode::Label::from(u32::from(arg)); if vm.is_none(&value) { - self.jump(target); + self.jump(bytecode::Label(target_idx)); if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { - monitoring::fire_branch_right(vm, self.code, src_offset, target.0 * 2)?; + monitoring::fire_branch_right(vm, self.code, src_offset, target_idx * 2)?; } } Ok(None) } Instruction::InstrumentedPopJumpIfNotNone => { let src_offset = (self.lasti() - 1) * 2; + let target_idx = self.lasti() + 1 + u32::from(arg); let value = self.pop_value(); - let target = bytecode::Label::from(u32::from(arg)); if !vm.is_none(&value) { - self.jump(target); + self.jump(bytecode::Label(target_idx)); if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { - monitoring::fire_branch_right(vm, self.code, src_offset, target.0 * 2)?; + monitoring::fire_branch_right(vm, self.code, src_offset, target_idx * 2)?; } } Ok(None) } Instruction::InstrumentedNotTaken => { if self.monitoring_mask & monitoring::EVENT_BRANCH_LEFT != 0 { - let offset = (self.lasti() - 1) * 2; + let not_taken_idx = self.lasti() as usize - 1; + // Scan backwards past CACHE entries to find the branch instruction + let mut branch_idx = not_taken_idx.saturating_sub(1); + while branch_idx > 0 + && matches!(self.code.instructions[branch_idx].op, Instruction::Cache) + { + branch_idx -= 1; + } + let src_offset = (branch_idx as u32) * 2; let dest_offset = self.lasti() * 2; - monitoring::fire_branch_left( - vm, - self.code, - offset.saturating_sub(2), - dest_offset, - )?; + monitoring::fire_branch_left(vm, self.code, src_offset, dest_offset)?; } Ok(None) } @@ -2823,12 +2849,15 @@ impl ExecutingFrame<'_> { // Re-dispatch to the real original opcode let original_op = Instruction::try_from(real_op_byte) .expect("invalid opcode in side-table chain"); - if original_op.to_base().is_some() { + let lasti_before_dispatch = self.lasti(); + let result = if original_op.to_base().is_some() { self.execute_instrumented(original_op, arg, vm) } else { let mut do_extend_arg = false; self.execute_instruction(original_op, arg, &mut do_extend_arg, vm) - } + }; + self.skip_caches_if_fallthrough(original_op, lasti_before_dispatch); + result } Instruction::InstrumentedInstruction => { let idx = self.lasti() as usize - 1; @@ -2852,12 +2881,15 @@ impl ExecutingFrame<'_> { // Re-dispatch to original opcode let original_op = Instruction::try_from(original_op_byte) .expect("invalid opcode in instruction side-table"); - if original_op.to_base().is_some() { + let lasti_before_dispatch = self.lasti(); + let result = if original_op.to_base().is_some() { self.execute_instrumented(original_op, arg, vm) } else { let mut do_extend_arg = false; self.execute_instruction(original_op, arg, &mut do_extend_arg, vm) - } + }; + self.skip_caches_if_fallthrough(original_op, lasti_before_dispatch); + result } _ => { unreachable!("{instruction:?} instruction should not be executed") @@ -3499,17 +3531,52 @@ impl ExecutingFrame<'_> { self.update_lasti(|i| *i = target_pc); } + /// Jump forward by `delta` code units from after instruction + caches. + /// lasti is already at instruction_index + 1, so after = lasti + caches. + /// + /// Unchecked arithmetic is intentional: the compiler guarantees valid + /// targets, and debug builds will catch overflow via Rust's default checks. + #[inline] + fn jump_relative_forward(&mut self, delta: u32, caches: u32) { + let target = self.lasti() + caches + delta; + self.update_lasti(|i| *i = target); + } + + /// Jump backward by `delta` code units from after instruction + caches. + /// + /// Unchecked arithmetic is intentional: the compiler guarantees valid + /// targets, and debug builds will catch underflow via Rust's default checks. + #[inline] + fn jump_relative_backward(&mut self, delta: u32, caches: u32) { + let target = self.lasti() + caches - delta; + self.update_lasti(|i| *i = target); + } + + /// Skip past CACHE code units after an instruction, but only if the + /// instruction did not modify lasti (i.e., it did not jump). #[inline] - fn pop_jump_if( + fn skip_caches_if_fallthrough(&mut self, op: Instruction, lasti_before: u32) { + if self.lasti() == lasti_before { + let base = op.to_base().unwrap_or(op); + let caches = base.cache_entries(); + if caches > 0 { + self.update_lasti(|i| *i += caches as u32); + } + } + } + + #[inline] + fn pop_jump_if_relative( &mut self, vm: &VirtualMachine, - target: bytecode::Label, + arg: bytecode::OpArg, + caches: u32, flag: bool, ) -> FrameResult { let obj = self.pop_value(); let value = obj.try_to_bool(vm)?; if value == flag { - self.jump(target); + self.jump_relative_forward(u32::from(arg), caches); } Ok(None) } diff --git a/crates/vm/src/stdlib/sys/monitoring.rs b/crates/vm/src/stdlib/sys/monitoring.rs index 05c427357dd..7b33f2eaf67 100644 --- a/crates/vm/src/stdlib/sys/monitoring.rs +++ b/crates/vm/src/stdlib/sys/monitoring.rs @@ -313,9 +313,12 @@ pub fn instrument_code(code: &PyCode, events: u32) { if matches!(op, Instruction::ExtendedArg) { continue; } - // Excluded: RESUME and END_FOR (and their instrumented variants) + // Excluded: RESUME, END_FOR, CACHE (and their instrumented variants) let base = op.to_base().map_or(op, |b| b); - if matches!(base, Instruction::Resume { .. } | Instruction::EndFor) { + if matches!( + base, + Instruction::Resume { .. } | Instruction::EndFor | Instruction::Cache + ) { continue; } // Store current opcode (may already be INSTRUMENTED_*) and replace @@ -358,6 +361,7 @@ pub fn instrument_code(code: &PyCode, events: u32) { | Instruction::EndSend | Instruction::PopIter | Instruction::EndAsyncFor + | Instruction::Cache ) { continue; } @@ -376,25 +380,34 @@ pub fn instrument_code(code: &PyCode, events: u32) { // same source line as the preceding instruction. Critical for loops // (JUMP_BACKWARD → FOR_ITER). let mut arg_state = bytecode::OpArgState::default(); + let mut instr_idx = first_traceable; for unit in code.code.instructions[first_traceable..len].iter().copied() { let (op, arg) = arg_state.get(unit); let base = op.to_base().map_or(op, |b| b); - if matches!(base, Instruction::ExtendedArg) { + if matches!(base, Instruction::ExtendedArg) || matches!(base, Instruction::Cache) { + instr_idx += 1; continue; } + let caches = base.cache_entries(); + let after_caches = instr_idx + 1 + caches; + let delta = u32::from(arg) as usize; + let target: Option = match base { + // Forward relative jumps Instruction::PopJumpIfFalse { .. } | Instruction::PopJumpIfTrue { .. } | Instruction::PopJumpIfNone { .. } | Instruction::PopJumpIfNotNone { .. } - | Instruction::JumpForward { .. } - | Instruction::JumpBackward { .. } - | Instruction::JumpBackwardNoInterrupt { .. } => Some(u32::from(arg) as usize), + | Instruction::JumpForward { .. } => Some(after_caches + delta), + // Backward relative jumps + Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } => { + Some(after_caches.wrapping_sub(delta)) + } Instruction::ForIter { .. } | Instruction::Send { .. } => { // Skip over END_FOR/END_SEND - Some(u32::from(arg) as usize + 1) + Some(after_caches + delta + 1) } _ => None, }; @@ -407,6 +420,7 @@ pub fn instrument_code(code: &PyCode, events: u32) { let target_base = target_op.to_base().map_or(target_op, |b| b); // Skip POP_ITER targets if matches!(target_base, Instruction::PopIter) { + instr_idx += 1; continue; } if let Some((loc, _)) = code.code.locations.get(target_idx) @@ -415,6 +429,7 @@ pub fn instrument_code(code: &PyCode, events: u32) { is_line_start[target_idx] = true; } } + instr_idx += 1; } // Third pass: mark exception handler targets as line starts.