From d201c48e1ca4e39fa58ffc981de1bdf8c181d226 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" <69878+youknowone@users.noreply.github.com> Date: Sun, 12 Apr 2026 22:43:35 +0900 Subject: [PATCH] Bytecode parity - for-loop cleanup, return duplication, late jump threading fix (#7580) - Use POP_TOP instead of POP_ITER for for-loop break/return cleanup - Expand duplicate_end_returns to clone final return for jump predecessors - Restrict late jump threading pass to unconditional jumps only - Skip exception blocks in inline/reorder passes - Simplify threaded_jump_instr NoInterrupt handling --- Lib/test/test_dis.py | 1 - Lib/test/test_yield_from.py | 2 - crates/codegen/src/compile.rs | 332 ++++++++++++++- crates/codegen/src/ir.rs | 387 +++++++++++++++--- ...thon_codegen__compile__tests__if_ands.snap | 2 +- ...hon_codegen__compile__tests__if_mixed.snap | 2 +- ...ython_codegen__compile__tests__if_ors.snap | 2 +- ...pile__tests__nested_double_async_with.snap | 33 +- .../compiler-core/src/bytecode/instruction.rs | 9 +- crates/compiler-core/src/bytecode/oparg.rs | 37 +- crates/vm/src/coroutine.rs | 57 ++- crates/vm/src/frame.rs | 74 ++-- crates/vm/src/stdlib/_ast/python.rs | 77 +++- crates/vm/src/stdlib/_thread.rs | 36 +- crates/vm/src/stdlib/os.rs | 19 +- crates/vm/src/stdlib/sys/monitoring.rs | 13 +- crates/vm/src/vm/thread.rs | 4 + extra_tests/snippets/stdlib_threading.py | 77 ++++ scripts/dis_dump.py | 130 +++--- 19 files changed, 1031 insertions(+), 263 deletions(-) create mode 100644 extra_tests/snippets/stdlib_threading.py diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py index 15455a0fc..a9533f5a3 100644 --- a/Lib/test/test_dis.py +++ b/Lib/test/test_dis.py @@ -1187,7 +1187,6 @@ class DisTests(DisTestBase): def test_disassemble_instance_method(self): self.do_disassembly_test(_C(1).__init__, dis_c_instance_method) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_disassemble_instance_method_bytes(self): method_bytecode = _C(1).__init__.__code__.co_code self.do_disassembly_test(method_bytecode, dis_c_instance_method_bytes) diff --git a/Lib/test/test_yield_from.py b/Lib/test/test_yield_from.py index 7028a6062..6da48253b 100644 --- a/Lib/test/test_yield_from.py +++ b/Lib/test/test_yield_from.py @@ -882,7 +882,6 @@ class TestPEP380Operation(unittest.TestCase): yield from () self.assertRaises(StopIteration, next, g()) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_delegating_generators_claim_to_be_running(self): # Check with basic iteration def one(): @@ -909,7 +908,6 @@ class TestPEP380Operation(unittest.TestCase): pass self.assertEqual(res, [0, 1, 2, 3]) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: Lists differ: [0, 1, 2] != [0, 1, 2, 3] def test_delegating_generators_claim_to_be_running_with_throw(self): # Check with throw class MyErr(Exception): diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 51b513c75..1a35d2c23 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -1513,13 +1513,11 @@ 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); } FBlockType::TryExcept => { @@ -3442,7 +3440,11 @@ impl Compiler { let handler_block = self.new_block(); let cleanup_block = self.new_block(); let end_block = self.new_block(); - let orelse_block = self.new_block(); + let orelse_block = if orelse.is_empty() { + end_block + } else { + self.new_block() + }; emit!(self, Instruction::Nop); emit!( @@ -3589,16 +3591,16 @@ impl Compiler { emit!(self, Instruction::Reraise { depth: 1 }); self.set_no_location(); - self.switch_to_block(orelse_block); - self.set_no_location(); if !orelse.is_empty() { + self.switch_to_block(orelse_block); + self.set_no_location(); self.compile_statements(orelse)?; + emit!( + self, + PseudoInstruction::JumpNoInterrupt { delta: end_block } + ); + self.set_no_location(); } - emit!( - self, - PseudoInstruction::JumpNoInterrupt { delta: end_block } - ); - self.set_no_location(); self.switch_to_block(end_block); Ok(()) @@ -5809,7 +5811,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] @@ -7739,6 +7744,7 @@ impl Compiler { /// JUMP send /// fail: /// CLEANUP_THROW + /// JUMP exit /// exit: /// END_SEND fn compile_yield_from_sequence(&mut self, is_await: bool) -> CompileResult { @@ -7786,10 +7792,13 @@ impl Compiler { // fail: CLEANUP_THROW // Stack when exception: [receiver, yielded_value, exc] // CLEANUP_THROW: [sub_iter, last_sent_val, exc] -> [None, value] - // After: stack is [None, value], fall through to exit + // Jump back to the shared END_SEND block like CPython's fail path. self.switch_to_block(fail_block); emit!(self, Instruction::CleanupThrow); - // Fall through to exit block + emit!( + self, + PseudoInstruction::JumpNoInterrupt { delta: exit_block } + ); // exit: END_SEND // Stack: [receiver, value] (from SEND) or [None, value] (from CLEANUP_THROW) @@ -9920,9 +9929,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 +11164,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 +11253,210 @@ def f(obj): ); } + #[test] + fn test_nested_try_finally_cleanup_reorder_does_not_invert_forward_jumps() { + compile_exec(include_str!("../../../Lib/poplib.py")); + } + + #[test] + fn test_conditional_body_is_preserved_before_final_return() { + let code = compile_exec( + "\ +def f(x, y): + if x == y: + print('then', flush=True) +", + ); + 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 cond_idx = ops + .iter() + .position(|op| matches!(op, Instruction::PopJumpIfFalse { .. })) + .expect("missing POP_JUMP_IF_FALSE"); + let first_return_idx = ops + .iter() + .position(|op| matches!(op, Instruction::ReturnValue)) + .expect("missing RETURN_VALUE"); + + assert!( + ops[cond_idx..first_return_idx] + .iter() + .any(|op| matches!(op, Instruction::CallKw { .. })), + "expected conditional body call before final return, got ops={ops:?}" + ); + } + + #[test] + fn test_nested_conditional_body_is_preserved_before_final_return() { + let code = compile_exec( + "\ +def outer(): + def side(): + print('side', flush=True) + def cb(): + flag = True + if flag: + side() + return cb +", + ); + let cb = find_code(&code, "cb").expect("missing nested cb code"); + let ops: Vec<_> = cb + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let cond_idx = ops + .iter() + .position(|op| matches!(op, Instruction::PopJumpIfFalse { .. })) + .expect("missing POP_JUMP_IF_FALSE"); + let first_return_idx = ops + .iter() + .position(|op| matches!(op, Instruction::ReturnValue)) + .expect("missing RETURN_VALUE"); + + assert!( + ops[cond_idx..first_return_idx] + .iter() + .any(|op| matches!(op, Instruction::Call { .. })), + "expected nested conditional body call before final return, got ops={ops:?}" + ); + } + + #[test] + fn test_conditional_compare_uses_bool_compare_oparg() { + let code = compile_exec( + "\ +def f(x, y): + if x == y: + return 1 + return 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let compare = f + .instructions + .iter() + .find(|unit| matches!(unit.op, Instruction::CompareOp { .. })) + .expect("missing COMPARE_OP"); + + assert_eq!(u8::from(compare.arg), 88); + } + + #[test] + fn test_chained_conditional_compares_use_bool_compare_oparg() { + let code = compile_exec( + "\ +def f(a, b, c): + if a < b < c: + return 1 + return 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let compare_args: Vec<_> = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::CompareOp { .. })) + .map(|unit| u8::from(unit.arg)) + .collect(); + + assert_eq!(compare_args, vec![18, 18]); + } + + #[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( @@ -11293,6 +11506,83 @@ def f(code): assert_eq!(pop_top_count, 1); } + #[test] + fn test_yield_from_cleanup_jumps_to_shared_end_send() { + let code = compile_exec( + "\ +def outer(): + def inner(): + yield from outer_gen + return inner +", + ); + let inner = find_code(&code, "inner").expect("missing inner code"); + let ops: Vec<_> = inner + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let cleanup_idx = ops + .iter() + .position(|op| matches!(op, Instruction::CleanupThrow)) + .expect("missing CLEANUP_THROW"); + assert!( + matches!( + ops.get(cleanup_idx + 1), + Some(Instruction::JumpBackwardNoInterrupt { .. }) + | Some(Instruction::JumpForward { .. }) + ), + "expected CLEANUP_THROW to jump to shared END_SEND block, got ops={ops:?}" + ); + assert!( + !matches!(ops.get(cleanup_idx + 1), Some(Instruction::EndSend)), + "CLEANUP_THROW should not inline END_SEND directly, got ops={ops:?}" + ); + } + + #[test] + fn test_try_except_falls_through_to_post_handler_code() { + let code = compile_exec( + "\ +def f(): + try: + line = 2 + raise KeyError + except: + line = 5 + line = 6 +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let first_pop_except = ops + .iter() + .position(|op| matches!(op, Instruction::PopExcept)) + .expect("missing POP_EXCEPT"); + assert!( + !matches!( + ops.get(first_pop_except + 1), + Some(Instruction::JumpForward { .. }) + ), + "expected except body to fall through to post-handler code, got ops={ops:?}" + ); + assert!( + matches!( + ops.get(first_pop_except + 1), + Some(Instruction::LoadSmallInt { .. }) | Some(Instruction::LoadConst { .. }) + ), + "expected line-after-except code immediately after POP_EXCEPT, got ops={ops:?}" + ); + } + #[test] fn test_constant_slice_folding_handles_string_and_bigint_bounds() { let code = compile_exec( @@ -11335,7 +11625,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 9b48cab08..84a947b7c 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -242,6 +242,7 @@ impl CodeInfo { normalize_jumps(&mut self.blocks); reorder_conditional_exit_and_jump_blocks(&mut self.blocks); reorder_conditional_jump_and_exit_blocks(&mut self.blocks); + reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); inline_small_or_no_lineno_blocks(&mut self.blocks); self.dce(); // re-run within-block DCE after normalize_jumps creates new instructions self.eliminate_unreachable_blocks(); @@ -256,6 +257,7 @@ impl CodeInfo { jump_threading_unconditional(&mut self.blocks); reorder_conditional_exit_and_jump_blocks(&mut self.blocks); reorder_conditional_jump_and_exit_blocks(&mut self.blocks); + reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); self.eliminate_unreachable_blocks(); remove_redundant_nops_and_jumps(&mut self.blocks); self.add_checks_for_loads_of_uninitialized_variables(); @@ -1359,32 +1361,33 @@ impl CodeInfo { /// - n == 2 or 3: BUILD_TUPLE → NOP, UNPACK_SEQUENCE → SWAP fn optimize_build_tuple_unpack(&mut self) { for block in &mut self.blocks { - let instrs = &mut block.instructions; - let len = instrs.len(); + let instructions = &mut block.instructions; + let len = instructions.len(); for i in 0..len.saturating_sub(1) { - let Some(Instruction::BuildTuple { .. }) = instrs[i].instr.real() else { + let Some(Instruction::BuildTuple { .. }) = instructions[i].instr.real() else { continue; }; - let n = u32::from(instrs[i].arg); - let Some(Instruction::UnpackSequence { .. }) = instrs[i + 1].instr.real() else { + let n = u32::from(instructions[i].arg); + let Some(Instruction::UnpackSequence { .. }) = instructions[i + 1].instr.real() + else { continue; }; - if u32::from(instrs[i + 1].arg) != n { + if u32::from(instructions[i + 1].arg) != n { continue; } match n { 1 => { - instrs[i].instr = AnyInstruction::Real(Instruction::Nop); - instrs[i].arg = OpArg::new(0); - instrs[i + 1].instr = AnyInstruction::Real(Instruction::Nop); - instrs[i + 1].arg = OpArg::new(0); + instructions[i].instr = AnyInstruction::Real(Instruction::Nop); + instructions[i].arg = OpArg::new(0); + instructions[i + 1].instr = AnyInstruction::Real(Instruction::Nop); + instructions[i + 1].arg = OpArg::new(0); } 2 | 3 => { - instrs[i].instr = AnyInstruction::Real(Instruction::Nop); - instrs[i].arg = OpArg::new(0); - instrs[i + 1].instr = + instructions[i].instr = AnyInstruction::Real(Instruction::Nop); + instructions[i].arg = OpArg::new(0); + instructions[i + 1].instr = AnyInstruction::Real(Instruction::Swap { i: Arg::marker() }); - instrs[i + 1].arg = OpArg::new(n); + instructions[i + 1].arg = OpArg::new(n); } _ => {} } @@ -1419,16 +1422,20 @@ impl CodeInfo { } } - /// Next swappable index after `i` in `instrs`, skipping NOPs. + /// Next swappable index after `i` in `instructions`, skipping NOPs. /// Returns None if a non-NOP non-swappable instruction blocks, or /// if `lineno >= 0` and a different lineno is encountered. - fn next_swappable(instrs: &[InstructionInfo], mut i: usize, lineno: i32) -> Option { + fn next_swappable( + instructions: &[InstructionInfo], + mut i: usize, + lineno: i32, + ) -> Option { loop { i += 1; - if i >= instrs.len() { + if i >= instructions.len() { return None; } - let info = &instrs[i]; + let info = &instructions[i]; let info_lineno = info.location.line.get() as i32; if lineno >= 0 && info_lineno > 0 && info_lineno != lineno { return None; @@ -1444,13 +1451,15 @@ impl CodeInfo { } for block in &mut self.blocks { - let instrs = &mut block.instructions; - let len = instrs.len(); + let instructions = &mut block.instructions; + let len = instructions.len(); // Walk forward; for each SWAP attempt elimination. let mut i = 0; while i < len { - let swap_arg = match instrs[i].instr { - AnyInstruction::Real(Instruction::Swap { .. }) => u32::from(instrs[i].arg), + let swap_arg = match instructions[i].instr { + AnyInstruction::Real(Instruction::Swap { .. }) => { + u32::from(instructions[i].arg) + } _ => { i += 1; continue; @@ -1463,17 +1472,17 @@ impl CodeInfo { continue; } // Find first swappable after SWAP (lineno = -1 initially). - let Some(j) = next_swappable(instrs, i, -1) else { + let Some(j) = next_swappable(instructions, i, -1) else { i += 1; continue; }; - let lineno = instrs[j].location.line.get() as i32; + let lineno = instructions[j].location.line.get() as i32; // Walk (swap_arg - 1) more swappable instructions, with // lineno constraint. let mut k = j; let mut ok = true; for _ in 1..swap_arg { - match next_swappable(instrs, k, lineno) { + match next_swappable(instructions, k, lineno) { Some(next) => k = next, None => { ok = false; @@ -1488,14 +1497,14 @@ impl CodeInfo { // Conflict check: if either j or k is a STORE_FAST, no // intervening store may target the same variable, and // they must not target the same variable themselves. - let store_j = stores_to(&instrs[j]); - let store_k = stores_to(&instrs[k]); + let store_j = stores_to(&instructions[j]); + let store_k = stores_to(&instructions[k]); if store_j.is_some() || store_k.is_some() { if store_j == store_k { i += 1; continue; } - let conflict = instrs[(j + 1)..k].iter().any(|info| { + let conflict = instructions[(j + 1)..k].iter().any(|info| { if let Some(store_idx) = stores_to(info) { Some(store_idx) == store_j || Some(store_idx) == store_k } else { @@ -1508,9 +1517,9 @@ impl CodeInfo { } } // Safe to reorder. SWAP -> NOP, swap j and k. - instrs[i].instr = AnyInstruction::Real(Instruction::Nop); - instrs[i].arg = OpArg::new(0); - instrs.swap(j, k); + instructions[i].instr = AnyInstruction::Real(Instruction::Nop); + instructions[i].arg = OpArg::new(0); + instructions.swap(j, k); i += 1; } } @@ -1529,13 +1538,13 @@ impl CodeInfo { /// store to a variable carries the final value, earlier ones are dead. fn eliminate_dead_stores(&mut self) { for block in &mut self.blocks { - let instrs = &mut block.instructions; - let len = instrs.len(); + let instructions = &mut block.instructions; + let len = instructions.len(); let mut i = 0; while i < len { // Look for UNPACK_SEQUENCE or UNPACK_EX let is_unpack = matches!( - instrs[i].instr, + instructions[i].instr, AnyInstruction::Real( Instruction::UnpackSequence { .. } | Instruction::UnpackEx { .. } ) @@ -1549,7 +1558,7 @@ impl CodeInfo { let mut run_end = run_start; while run_end < len && matches!( - instrs[run_end].instr, + instructions[run_end].instr, AnyInstruction::Real(Instruction::StoreFast { .. }) ) { @@ -1558,11 +1567,11 @@ impl CodeInfo { if run_end - run_start >= 2 { // Pass 1: find the LAST occurrence of each variable let mut last_occurrence = std::collections::HashMap::new(); - for (j, instr) in instrs[run_start..run_end].iter().enumerate() { + for (j, instr) in instructions[run_start..run_end].iter().enumerate() { last_occurrence.insert(u32::from(instr.arg), j); } // Pass 2: non-last stores to the same variable are dead - for (j, instr) in instrs[run_start..run_end].iter_mut().enumerate() { + for (j, instr) in instructions[run_start..run_end].iter_mut().enumerate() { let idx = u32::from(instr.arg); if last_occurrence[&idx] != j { instr.instr = AnyInstruction::Real(Instruction::PopTop); @@ -1580,18 +1589,29 @@ impl CodeInfo { 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]; + + // Only combine if both are real instructions (not pseudo) + let (Some(curr_instr), Some(next_instr)) = (curr.instr.real(), next.instr.real()) + else { + i += 1; + continue; + }; + + if matches!( + next_instr, + Instruction::PopJumpIfFalse { .. } | Instruction::PopJumpIfTrue { .. } + ) && matches!(curr_instr, Instruction::CompareOp { .. }) + { + block.instructions[i].arg = OpArg::new( + u32::from(block.instructions[i].arg) | oparg::COMPARE_OP_BOOL_MASK, + ); + i += 1; + continue; + } + let combined = { - let curr = &block.instructions[i]; - let next = &block.instructions[i + 1]; - - // Only combine if both are real instructions (not pseudo) - let (Some(curr_instr), Some(next_instr)) = - (curr.instr.real(), next.instr.real()) - else { - i += 1; - continue; - }; - match (curr_instr, next_instr) { // LoadFast + LoadFast -> LoadFastLoadFast (if both indices < 16) (Instruction::LoadFast { .. }, Instruction::LoadFast { .. }) => { @@ -1652,6 +1672,14 @@ impl CodeInfo { None } } + (Instruction::CompareOp { .. }, Instruction::ToBool) => Some(( + curr_instr, + OpArg::new(u32::from(curr.arg) | oparg::COMPARE_OP_BOOL_MASK), + )), + (Instruction::ContainsOp { .. }, Instruction::ToBool) + | (Instruction::IsOp { .. }, Instruction::ToBool) => { + Some((curr_instr, curr.arg)) + } (Instruction::LoadConst { consti }, Instruction::UnaryNot) => { let constant = &self.metadata.consts[consti.get(curr.arg).as_usize()]; match constant { @@ -2724,6 +2752,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 +2812,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 +2832,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 +2847,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 +3106,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 +3309,21 @@ fn is_scope_exit_block(block: &Block) -> bool { .is_some_and(|instr| instr.instr.is_scope_exit()) } +fn is_exception_cleanup_block(block: &Block) -> bool { + block + .instructions + .iter() + .any(|instr| matches!(instr.instr.real(), Some(Instruction::PopExcept))) + && block + .instructions + .last() + .is_some_and(|instr| matches!(instr.instr.real(), Some(Instruction::Reraise { .. }))) +} + +fn block_is_exceptional(block: &Block) -> bool { + block.except_handler || block.preserve_lasti || is_exception_cleanup_block(block) +} + 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 +3374,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 +3402,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 +3463,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 +3496,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; @@ -3404,6 +3530,92 @@ fn reorder_conditional_jump_and_exit_blocks(blocks: &mut [Block]) { } } +#[allow(dead_code)] +fn reorder_jump_over_exception_cleanup_blocks(blocks: &mut [Block]) { + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + let idx = current.idx(); + let next = blocks[idx].next; + let Some(last) = blocks[idx].instructions.last().copied() else { + current = next; + continue; + }; + if !matches!(last.instr.real(), Some(Instruction::JumpForward { .. })) + || last.target == BlockIdx::NULL + { + current = next; + continue; + } + + let cleanup_start = next; + let target_start = last.target; + let target = next_nonempty_block(blocks, target_start); + if cleanup_start == BlockIdx::NULL || target == BlockIdx::NULL || cleanup_start == target { + current = next; + continue; + } + // Keep the target anchored to the first target block. If we have to + // skip leading empty blocks here, reordering can leave the jump shape + // inconsistent in nested cleanup chains such as poplib.POP3.close(). + if target_start != target { + current = next; + continue; + } + + let mut cleanup_end = BlockIdx::NULL; + let mut saw_exceptional = false; + let mut cursor = cleanup_start; + while cursor != BlockIdx::NULL && cursor != target { + if blocks[cursor.idx()].instructions.is_empty() { + cleanup_end = cursor; + cursor = blocks[cursor.idx()].next; + continue; + } + if !block_is_exceptional(&blocks[cursor.idx()]) + && !is_exception_cleanup_block(&blocks[cursor.idx()]) + { + cleanup_end = BlockIdx::NULL; + break; + } + saw_exceptional = true; + cleanup_end = cursor; + cursor = blocks[cursor.idx()].next; + } + if !saw_exceptional || cleanup_end == BlockIdx::NULL || cursor != target { + current = next; + continue; + } + + let mut target_end = BlockIdx::NULL; + let mut target_exit = BlockIdx::NULL; + cursor = target; + while cursor != BlockIdx::NULL { + if block_is_exceptional(&blocks[cursor.idx()]) { + break; + } + target_end = cursor; + if !blocks[cursor.idx()].instructions.is_empty() { + target_exit = cursor; + } + cursor = blocks[cursor.idx()].next; + } + + if target_end == BlockIdx::NULL + || target_exit == BlockIdx::NULL + || !is_scope_exit_block(&blocks[target_exit.idx()]) + { + current = next; + continue; + } + + let after_target = blocks[target_end.idx()].next; + blocks[idx].next = target_start; + blocks[target_end.idx()].next = cleanup_start; + blocks[cleanup_end.idx()].next = after_target; + current = after_target; + } +} + fn maybe_propagate_location( instr: &mut InstructionInfo, location: SourceLocation, @@ -3619,9 +3831,23 @@ fn resolve_line_numbers(blocks: &mut Vec) { propagate_line_numbers(blocks, &predecessors); } +fn find_layout_predecessor(blocks: &[Block], target: BlockIdx) -> BlockIdx { + if target == BlockIdx::NULL { + return BlockIdx::NULL; + } + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + if blocks[current.idx()].next == target { + return current; + } + current = blocks[current.idx()].next; + } + BlockIdx::NULL +} + /// 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 +3878,18 @@ 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 an unconditional 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 +3905,64 @@ 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 + && trailing_conditional_jump_index(block).is_none() + && !already_has_return + { + fallthrough_blocks_to_fix.push(current); + } + if predecessors[last_block.idx()] > 1 + && 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 is_conditional = is_conditional_jump(&jump.instr); + let new_block = Block { + cold: blocks[last_block.idx()].cold, + except_handler: blocks[last_block.idx()].except_handler, + instructions: cloned_return, + next: if is_conditional { + last_block + } else { + blocks[block_idx.idx()].next + }, + ..Block::default() + }; + blocks.push(new_block); + if is_conditional { + let layout_pred = find_layout_predecessor(blocks, last_block); + if layout_pred != BlockIdx::NULL { + blocks[layout_pred.idx()].next = new_idx; + } + } else { + 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 4783c0f2d..0cdb7f0a3 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,6 +1,6 @@ --- 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) 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 043bf380a..8e9118991 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,6 +1,6 @@ --- 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) 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 bf4960582..944563545 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,6 +1,6 @@ --- 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) 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 4573b7a94..64ae0cfd5 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: 10936 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) @@ -34,7 +34,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 26 CACHE 27 STORE_FAST (0, stop_exc) - 3 28 LOAD_GLOBAL (4, self) + 3 >> 28 LOAD_GLOBAL (4, self) 29 CACHE 30 CACHE 31 CACHE @@ -151,7 +151,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 136 CACHE 137 CACHE 138 CHECK_EXC_MATCH - 139 POP_JUMP_IF_FALSE (32) + 139 POP_JUMP_IF_FALSE (28) 140 CACHE 141 NOT_TAKEN 142 STORE_FAST (1, ex) @@ -182,11 +182,11 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 166 STORE_FAST (1, ex) 167 DELETE_FAST (1, ex) 168 JUMP_FORWARD (32) - 169 LOAD_CONST (None) - 170 STORE_FAST (1, ex) - 171 DELETE_FAST (1, ex) - 172 RERAISE (1) - 173 RERAISE (0) + 169 RERAISE (0) + 170 LOAD_CONST (None) + 171 STORE_FAST (1, ex) + 172 DELETE_FAST (1, ex) + 173 RERAISE (1) 174 COPY (3) 175 POP_EXCEPT 176 RERAISE (1) @@ -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/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index 76dc6a793..b0edc922a 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -1169,7 +1169,14 @@ impl InstructionMetadata for Instruction { Self::CheckEgMatch => w!(CHECK_EG_MATCH), Self::CheckExcMatch => w!(CHECK_EXC_MATCH), Self::CleanupThrow => w!(CLEANUP_THROW), - Self::CompareOp { opname } => w!(COMPARE_OP, ?opname), + Self::CompareOp { opname } => { + let op = opname.get(arg); + if u32::from(arg) & oparg::COMPARE_OP_BOOL_MASK != 0 { + write!(f, "{:pad$}(bool({}))", "COMPARE_OP", op) + } else { + write!(f, "{:pad$}({})", "COMPARE_OP", op) + } + } Self::ContainsOp { invert } => w!(CONTAINS_OP, ?invert), Self::ConvertValue { oparg } => write!(f, "{:pad$}{}", "CONVERT_VALUE", oparg.get(arg)), Self::Copy { i } => w!(COPY, i), diff --git a/crates/compiler-core/src/bytecode/oparg.rs b/crates/compiler-core/src/bytecode/oparg.rs index a0ab45a3a..0a2a660a4 100644 --- a/crates/compiler-core/src/bytecode/oparg.rs +++ b/crates/compiler-core/src/bytecode/oparg.rs @@ -388,9 +388,12 @@ impl From for u32 { impl OpArgType for MakeFunctionFlag {} -/// `COMPARE_OP` arg is `(cmp_index << 5) | mask`. Only the upper -/// 3 bits identify the comparison; the lower 5 bits are an inline -/// cache mask for adaptive specialization. +/// `COMPARE_OP` arg is `(cmp_index << 5) | mask`. +/// +/// The low four bits are the CPython comparison mask used by specialized +/// compare opcodes, and bit 4 requests bool-conversion of the compare result. +pub const COMPARE_OP_BOOL_MASK: u32 = 1 << 4; + #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum ComparisonOperator { Less, @@ -425,15 +428,15 @@ impl TryFrom for ComparisonOperator { } impl From for u8 { - /// Encode as `cmp_index << 5` (mask bits zero). + /// Encode using CPython's comparison mask layout. fn from(value: ComparisonOperator) -> Self { match value { - ComparisonOperator::Less => 0, - ComparisonOperator::LessOrEqual => 1 << 5, - ComparisonOperator::Equal => 2 << 5, - ComparisonOperator::NotEqual => 3 << 5, - ComparisonOperator::Greater => 4 << 5, - ComparisonOperator::GreaterOrEqual => 5 << 5, + ComparisonOperator::Less => 2, + ComparisonOperator::LessOrEqual => (1 << 5) | 2 | 8, + ComparisonOperator::Equal => (2 << 5) | 8, + ComparisonOperator::NotEqual => (3 << 5) | 1 | 2 | 4, + ComparisonOperator::Greater => (4 << 5) | 4, + ComparisonOperator::GreaterOrEqual => (5 << 5) | 4 | 8, } } } @@ -446,6 +449,20 @@ impl From for u32 { impl OpArgType for ComparisonOperator {} +impl fmt::Display for ComparisonOperator { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let op = match self { + Self::Less => "<", + Self::LessOrEqual => "<=", + Self::Equal => "==", + Self::NotEqual => "!=", + Self::Greater => ">", + Self::GreaterOrEqual => ">=", + }; + f.write_str(op) + } +} + oparg_enum!( /// The possible Binary operators /// diff --git a/crates/vm/src/coroutine.rs b/crates/vm/src/coroutine.rs index 07158c488..5c24b2d6d 100644 --- a/crates/vm/src/coroutine.rs +++ b/crates/vm/src/coroutine.rs @@ -73,7 +73,10 @@ impl Coro { } } - fn maybe_close(&self, res: &PyResult) { + fn maybe_close(&self, res: &PyResult, entered_frame: bool) { + if !entered_frame { + return; + } match res { Ok(ExecutionResult::Return(_)) | Err(_) => { self.closed.store(true); @@ -82,6 +85,9 @@ impl Coro { FrameOwner::FrameObject as i8, core::sync::atomic::Ordering::Release, ); + // Completed generators/coroutines should not keep their locals + // alive while the wrapper object itself remains referenced. + self.frame.clear_locals_and_stack(); } Ok(ExecutionResult::Yield(_)) => {} } @@ -92,12 +98,15 @@ impl Coro { jen: &PyObject, vm: &VirtualMachine, func: F, - ) -> PyResult + ) -> (PyResult, bool) where F: FnOnce(&Py) -> PyResult, { if self.running.compare_exchange(false, true).is_err() { - return Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm)))); + return ( + Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm)))), + false, + ); } // SAFETY: running.compare_exchange guarantees exclusive access @@ -112,16 +121,17 @@ impl Coro { }); self.running.store(false); - result + (result, true) } fn finalize_send_result( &self, - jen: &PyObject, result: PyResult, + entered_frame: bool, + jen: &PyObject, vm: &VirtualMachine, ) -> PyResult { - self.maybe_close(&result); + self.maybe_close(&result, entered_frame); match result { Ok(exec_res) => Ok(exec_res.into_iter_return(vm)), Err(e) => { @@ -147,14 +157,19 @@ impl Coro { if self.closed.load() { return Ok(PyIterReturn::StopIteration(None)); } - self.frame.locals_to_fast(vm)?; + if self.running.load() { + return Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm)))); + } let value = if self.frame.lasti() > 0 { Some(vm.ctx.none()) } else { None }; - let result = self.run_with_context(jen, vm, |f| f.resume(value, vm)); - self.finalize_send_result(jen, result, vm) + let (result, entered_frame) = self.run_with_context(jen, vm, |f| { + self.frame.locals_to_fast(vm)?; + f.resume(value, vm) + }); + self.finalize_send_result(result, entered_frame, jen, vm) } pub fn send( @@ -166,7 +181,9 @@ impl Coro { if self.closed.load() { return Ok(PyIterReturn::StopIteration(None)); } - self.frame.locals_to_fast(vm)?; + if self.running.load() { + return Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm)))); + } let value = if self.frame.lasti() > 0 { Some(value) } else if !vm.is_none(&value) { @@ -177,8 +194,11 @@ impl Coro { } else { None }; - let result = self.run_with_context(jen, vm, |f| f.resume(value, vm)); - self.finalize_send_result(jen, result, vm) + let (result, entered_frame) = self.run_with_context(jen, vm, |f| { + self.frame.locals_to_fast(vm)?; + f.resume(value, vm) + }); + self.finalize_send_result(result, entered_frame, jen, vm) } pub fn throw( @@ -203,8 +223,9 @@ impl Coro { // Validate exception type before entering generator context. // Invalid types propagate to caller without closing the generator. crate::exceptions::ExceptionCtor::try_from_object(vm, exc_type.clone())?; - let result = self.run_with_context(jen, vm, |f| f.gen_throw(vm, exc_type, exc_val, exc_tb)); - self.maybe_close(&result); + let (result, entered_frame) = + self.run_with_context(jen, vm, |f| f.gen_throw(vm, exc_type, exc_val, exc_tb)); + self.maybe_close(&result, entered_frame); Ok(result?.into_iter_return(vm)) } @@ -217,7 +238,7 @@ impl Coro { self.closed.store(true); return Ok(vm.ctx.none()); } - let result = self.run_with_context(jen, vm, |f| { + let (result, entered_frame) = self.run_with_context(jen, vm, |f| { f.gen_throw( vm, vm.ctx.exceptions.generator_exit.to_owned().into(), @@ -225,6 +246,12 @@ impl Coro { vm.ctx.none(), ) }); + if !entered_frame { + return match result { + Err(err) => Err(err), + Ok(_) => unreachable!("run_with_context preflight returned without an error"), + }; + } self.closed.store(true); // Release frame locals and stack to free references held by the // closed generator, matching gen_send_ex2 with close_on_completion. diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 66b8a87f2..9bdaf4775 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -1522,11 +1522,16 @@ impl ExecutingFrame<'_> { idx: usize, vm: &VirtualMachine, ) -> FrameResult { - let (loc, _end_loc) = frame.code.locations[idx]; - let next = exception.__traceback__(); - let new_traceback = - PyTraceback::new(next, frame.object.to_owned(), idx as u32 * 2, loc.line); - exception.set_traceback_typed(Some(new_traceback.into_ref(&vm.ctx))); + if let Some((loc, _end_loc)) = frame.code.locations.get(idx) { + let next = exception.__traceback__(); + let new_traceback = PyTraceback::new( + next, + frame.object.to_owned(), + idx as u32 * 2, + loc.line, + ); + exception.set_traceback_typed(Some(new_traceback.into_ref(&vm.ctx))); + } vm.contextualize_exception(&exception); frame.unwind_blocks(vm, UnwindReason::Raising { exception }) } @@ -1579,17 +1584,23 @@ impl ExecutingFrame<'_> { // checking for duplicates. Each time an exception passes through // a frame (e.g., in a loop with repeated raise statements), // a new traceback entry is added. - let (loc, _end_loc) = frame.code.locations[idx]; - let next = exception.__traceback__(); + if let Some((loc, _end_loc)) = frame.code.locations.get(idx) { + let next = exception.__traceback__(); - let new_traceback = PyTraceback::new( - next, - frame.object.to_owned(), - idx as u32 * 2, - loc.line, - ); - vm_trace!("Adding to traceback: {:?} {:?}", new_traceback, loc.line); - exception.set_traceback_typed(Some(new_traceback.into_ref(&vm.ctx))); + let new_traceback = PyTraceback::new( + next, + frame.object.to_owned(), + idx as u32 * 2, + loc.line, + ); + vm_trace!( + "Adding to traceback: {:?} {:?}", + new_traceback, + loc.line + ); + exception + .set_traceback_typed(Some(new_traceback.into_ref(&vm.ctx))); + } // _PyErr_SetObject sets __context__ only when the exception // is first raised. When an exception propagates through frames, @@ -2240,7 +2251,7 @@ impl ExecutingFrame<'_> { Instruction::CompareOp { opname: op } => { let op_val = op.get(arg); self.adaptive(|s, ii, cb| s.specialize_compare_op(vm, op_val, ii, cb)); - self.execute_compare(vm, op_val) + self.execute_compare(vm, arg) } Instruction::ContainsOp { invert } => { self.adaptive(|s, ii, cb| s.specialize_contains_op(vm, ii, cb)); @@ -5337,9 +5348,7 @@ impl ExecutingFrame<'_> { self.push_value(vm.ctx.new_bool(result).into()); Ok(None) } else { - let op = bytecode::ComparisonOperator::try_from(u32::from(arg)) - .unwrap_or(bytecode::ComparisonOperator::Equal); - self.execute_compare(vm, op) + self.execute_compare(vm, arg) } } Instruction::CompareOpFloat => { @@ -5361,9 +5370,7 @@ impl ExecutingFrame<'_> { self.push_value(vm.ctx.new_bool(result).into()); Ok(None) } else { - let op = bytecode::ComparisonOperator::try_from(u32::from(arg)) - .unwrap_or(bytecode::ComparisonOperator::Equal); - self.execute_compare(vm, op) + self.execute_compare(vm, arg) } } Instruction::CompareOpStr => { @@ -5375,9 +5382,7 @@ impl ExecutingFrame<'_> { ) { let op = self.compare_op_from_arg(arg); if op != PyComparisonOp::Eq && op != PyComparisonOp::Ne { - let op = bytecode::ComparisonOperator::try_from(u32::from(arg)) - .unwrap_or(bytecode::ComparisonOperator::Equal); - return self.execute_compare(vm, op); + return self.execute_compare(vm, arg); } let result = op.eval_ord(a_str.as_wtf8().cmp(b_str.as_wtf8())); self.pop_value(); @@ -5385,9 +5390,7 @@ impl ExecutingFrame<'_> { self.push_value(vm.ctx.new_bool(result).into()); Ok(None) } else { - let op = bytecode::ComparisonOperator::try_from(u32::from(arg)) - .unwrap_or(bytecode::ComparisonOperator::Equal); - self.execute_compare(vm, op) + self.execute_compare(vm, arg) } } Instruction::ToBoolBool => { @@ -7256,14 +7259,13 @@ impl ExecutingFrame<'_> { } #[cfg_attr(feature = "flame-it", flame("Frame"))] - fn execute_compare( - &mut self, - vm: &VirtualMachine, - op: bytecode::ComparisonOperator, - ) -> FrameResult { + fn execute_compare(&mut self, vm: &VirtualMachine, arg: bytecode::OpArg) -> FrameResult { + let op = bytecode::ComparisonOperator::try_from(u32::from(arg)) + .unwrap_or(bytecode::ComparisonOperator::Equal); let b = self.pop_value(); let a = self.pop_value(); let cmp_op: PyComparisonOp = op.into(); + let force_bool = u32::from(arg) & bytecode::oparg::COMPARE_OP_BOOL_MASK != 0; // COMPARE_OP_INT: leaf type, cannot recurse — skip rich_compare dispatch if let (Some(a_int), Some(b_int)) = ( @@ -7287,6 +7289,12 @@ impl ExecutingFrame<'_> { } let value = a.rich_compare(b, cmp_op, vm)?; + let value = if force_bool { + let bool_val = value.try_to_bool(vm)?; + vm.ctx.new_bool(bool_val).into() + } else { + value + }; self.push_value(value); Ok(None) } diff --git a/crates/vm/src/stdlib/_ast/python.rs b/crates/vm/src/stdlib/_ast/python.rs index 5c97d7599..9a70a4567 100644 --- a/crates/vm/src/stdlib/_ast/python.rs +++ b/crates/vm/src/stdlib/_ast/python.rs @@ -8,7 +8,9 @@ use super::{ pub(crate) mod _ast { use crate::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, - builtins::{PyStr, PyStrRef, PyTupleRef, PyType, PyTypeRef, PyUtf8Str, PyUtf8StrRef}, + builtins::{ + PyDictRef, PyStr, PyStrRef, PyTupleRef, PyType, PyTypeRef, PyUtf8Str, PyUtf8StrRef, + }, class::{PyClassImpl, StaticType}, common::wtf8::Wtf8, function::{FuncArgs, KwArgs, PyMethodDef, PyMethodFlags}, @@ -54,9 +56,22 @@ pub(crate) mod _ast { PyMethodFlags::METHOD, None, ); + const AST_DEEPCOPY: PyMethodDef = PyMethodDef::new_const( + "__deepcopy__", + |zelf: PyObjectRef, memo: PyObjectRef, vm: &VirtualMachine| -> PyResult { + ast_deepcopy(zelf, memo, vm) + }, + PyMethodFlags::METHOD, + None, + ); class.set_str_attr("__reduce__", AST_REDUCE.to_proper_method(class, ctx), ctx); class.set_str_attr("__replace__", AST_REPLACE.to_proper_method(class, ctx), ctx); + class.set_str_attr( + "__deepcopy__", + AST_DEEPCOPY.to_proper_method(class, ctx), + ctx, + ); class.slots.repr.store(Some(ast_repr)); } @@ -84,6 +99,11 @@ pub(crate) mod _ast { fn __replace__(zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { ast_replace(zelf, args, vm) } + + #[pymethod] + fn __deepcopy__(zelf: PyObjectRef, memo: PyObjectRef, vm: &VirtualMachine) -> PyResult { + ast_deepcopy(zelf, memo, vm) + } } pub(crate) fn ast_reduce(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult { @@ -233,6 +253,46 @@ pub(crate) mod _ast { Ok(result) } + pub(crate) fn ast_deepcopy( + zelf: PyObjectRef, + memo: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + let memo_dict: PyDictRef = memo + .clone() + .downcast() + .map_err(|_| vm.new_type_error("__deepcopy__() memo must be a dict"))?; + let memo_key: PyObjectRef = vm.ctx.new_int(zelf.get_id() as i64).into(); + + if let Some(existing) = memo_dict.get_item_opt(&*memo_key, vm)? { + return Ok(existing); + } + + let cls = zelf.class(); + let copied_dict = if cls + .slots + .flags + .contains(crate::types::PyTypeFlags::HAS_DICT) + { + Some(vm.ctx.new_dict()) + } else { + None + }; + let copied = vm.ctx.new_base_object(cls.to_owned(), copied_dict.clone()); + + memo_dict.set_item(&*memo_key, copied.clone(), vm)?; + + if let (Some(src_dict), Some(dst_dict)) = (zelf.as_object().dict(), copied_dict) { + let deepcopy = vm.import("copy", 0)?.get_attr("deepcopy", vm)?; + for (key, value) in src_dict.items_vec() { + let copied_value = deepcopy.call((value, memo.clone()), vm)?; + dst_dict.set_item(&*key, copied_value, vm)?; + } + } + + Ok(copied) + } + pub(crate) fn ast_repr(zelf: &crate::PyObject, vm: &VirtualMachine) -> PyResult> { let repr = repr::repr_ast_node(vm, &zelf.to_owned(), 3)?; Ok(vm.ctx.new_str(repr)) @@ -242,15 +302,12 @@ pub(crate) mod _ast { type Args = FuncArgs; fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { - if args.args.is_empty() - && args.kwargs.is_empty() - && let Some(instance) = cls.get_attr(vm.ctx.intern_str("_instance")) - { - return Ok(instance); - } - - // AST nodes accept extra arguments (unlike object.__new__) - // This matches CPython's behavior where AST has its own tp_new + // Keep _instance for parser-internal shared operator/context nodes, + // but match CPython's public constructor behavior by allocating a + // fresh object for Python-level ast.Load()/ast.Add()/... calls. + // Returning the cached singleton here makes user-added attributes + // like `parent` leak across unrelated trees and breaks deepcopy. + // AST nodes accept extra arguments (unlike object.__new__). let dict = if cls .slots .flags diff --git a/crates/vm/src/stdlib/_thread.rs b/crates/vm/src/stdlib/_thread.rs index 765f25374..b3f29a7fa 100644 --- a/crates/vm/src/stdlib/_thread.rs +++ b/crates/vm/src/stdlib/_thread.rs @@ -55,7 +55,6 @@ pub(crate) mod _thread { // this is a value in seconds #[pyattr] const TIMEOUT_MAX: f64 = (TIMEOUT_MAX_IN_MICROSECONDS / 1_000_000) as f64; - #[pyattr] fn error(vm: &VirtualMachine) -> PyTypeRef { vm.ctx.exceptions.runtime_error.to_owned() @@ -551,11 +550,7 @@ pub(crate) mod _thread { .map(|(k, v)| (k.as_str().to_owned(), v)) .collect::(), ); - let mut thread_builder = thread::Builder::new(); - let stacksize = vm.state.stacksize.load(); - if stacksize != 0 { - thread_builder = thread_builder.stack_size(stacksize); - } + let thread_builder = apply_thread_stack_size(thread::Builder::new(), vm); thread_builder .spawn( vm.new_thread() @@ -593,6 +588,18 @@ pub(crate) mod _thread { vm.state.thread_count.fetch_sub(1); } + fn apply_thread_stack_size( + thread_builder: thread::Builder, + vm: &VirtualMachine, + ) -> thread::Builder { + let configured = vm.state.stacksize.load(); + if configured != 0 { + thread_builder.stack_size(configured) + } else { + thread_builder + } + } + /// Clean up thread-local data for the current thread. /// This triggers __del__ on objects stored in thread-local variables. fn cleanup_thread_local_data() { @@ -893,7 +900,7 @@ pub(crate) mod _thread { // Guard that removes thread-local data when dropped struct LocalGuard { local: Weak, - thread_id: std::thread::ThreadId, + thread_id: u64, } impl Drop for LocalGuard { @@ -909,7 +916,7 @@ pub(crate) mod _thread { // Shared data structure for Local struct LocalData { - data: parking_lot::Mutex>, + data: parking_lot::Mutex>, } impl fmt::Debug for LocalData { @@ -928,7 +935,7 @@ pub(crate) mod _thread { #[pyclass(with(GetAttr, SetAttr), flags(BASETYPE))] impl Local { fn l_dict(&self, vm: &VirtualMachine) -> PyDictRef { - let thread_id = std::thread::current().id(); + let thread_id = current_thread_id(); // Fast path: check if dict exists under lock if let Some(dict) = self.inner.data.lock().get(&thread_id).cloned() { @@ -964,6 +971,11 @@ pub(crate) mod _thread { dict } + #[pygetset(name = "__dict__")] + fn dict(zelf: PyRef, vm: &VirtualMachine) -> PyDictRef { + zelf.l_dict(vm) + } + #[pyslot] fn slot_new(cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult { Self { @@ -1626,11 +1638,7 @@ pub(crate) mod _thread { Arc::new((std::sync::Mutex::new(false), std::sync::Condvar::new())); let handle_ready_event_clone = Arc::clone(&handle_ready_event); - let mut thread_builder = thread::Builder::new(); - let stacksize = vm.state.stacksize.load(); - if stacksize != 0 { - thread_builder = thread_builder.stack_size(stacksize); - } + let thread_builder = apply_thread_stack_size(thread::Builder::new(), vm); let join_handle = thread_builder .spawn(vm.new_thread().make_spawn_func(move |vm| { diff --git a/crates/vm/src/stdlib/os.rs b/crates/vm/src/stdlib/os.rs index bc9d24aa4..949fda825 100644 --- a/crates/vm/src/stdlib/os.rs +++ b/crates/vm/src/stdlib/os.rs @@ -322,18 +322,17 @@ pub(super) mod _os { #[pyfunction] fn write(fd: crt_fd::Borrowed<'_>, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult { - data.with_ref(|b| { - loop { - match vm.allow_threads(|| crt_fd::write(fd, b)) { - Ok(n) => return Ok(n), - Err(e) if e.raw_os_error() == Some(libc::EINTR) => { - vm.check_signals()?; - continue; - } - Err(e) => return Err(e.into_pyexception(vm)), + let owned = data.with_ref(|b| b.to_vec()); + loop { + match vm.allow_threads(|| crt_fd::write(fd, &owned)) { + Ok(n) => return Ok(n), + Err(e) if e.raw_os_error() == Some(libc::EINTR) => { + vm.check_signals()?; + continue; } + Err(e) => return Err(e.into_pyexception(vm)), } - }) + } } #[cfg(not(windows))] diff --git a/crates/vm/src/stdlib/sys/monitoring.rs b/crates/vm/src/stdlib/sys/monitoring.rs index 367038658..f802cd94f 100644 --- a/crates/vm/src/stdlib/sys/monitoring.rs +++ b/crates/vm/src/stdlib/sys/monitoring.rs @@ -367,6 +367,11 @@ pub fn instrument_code(code: &PyCode, events: u32) { if events & EVENT_LINE != 0 { // is_line_start[i] = true if position i should have INSTRUMENTED_LINE let mut is_line_start = vec![false; len]; + let line_locations = rustpython_compiler_core::marshal::linetable_to_locations( + &code.code.linetable, + code.code.first_line_number.map_or(-1, |x| x.get() as i32), + len, + ); // Build NO_LOCATION mask from linetable let no_loc_mask = bytecode::build_no_location_mask(&code.code.linetable, len); @@ -402,7 +407,7 @@ pub fn instrument_code(code: &PyCode, events: u32) { if no_loc_mask.get(i).copied().unwrap_or(false) { continue; } - if let Some((loc, _)) = code.code.locations.get(i) { + if let Some((loc, _)) = line_locations.get(i) { let line = loc.line.get() as u32; let is_new = prev_line != Some(line); prev_line = Some(line); @@ -456,12 +461,12 @@ pub fn instrument_code(code: &PyCode, events: u32) { { let target_op = code.code.instructions[target_idx].op; let target_base = target_op.to_base().map_or(target_op, |b| b); - // Skip POP_ITER targets + // Skip synthetic cleanup targets. if matches!(target_base, Instruction::PopIter) { instr_idx += 1; continue; } - if let Some((loc, _)) = code.code.locations.get(target_idx) + if let Some((loc, _)) = line_locations.get(target_idx) && loc.line.get() > 0 { is_line_start[target_idx] = true; @@ -480,7 +485,7 @@ pub fn instrument_code(code: &PyCode, events: u32) { let target_op = code.code.instructions[target_idx].op; let target_base = target_op.to_base().map_or(target_op, |b| b); if !matches!(target_base, Instruction::PopIter) - && let Some((loc, _)) = code.code.locations.get(target_idx) + && let Some((loc, _)) = line_locations.get(target_idx) && loc.line.get() > 0 { is_line_start[target_idx] = true; diff --git a/crates/vm/src/vm/thread.rs b/crates/vm/src/vm/thread.rs index e7cc64f00..e5a961e1c 100644 --- a/crates/vm/src/vm/thread.rs +++ b/crates/vm/src/vm/thread.rs @@ -557,6 +557,10 @@ impl ThreadedVirtualMachine { F: FnOnce(&VirtualMachine) -> R, { let vm = &self.vm; + // Each spawned thread has its own native stack bounds. Recompute the + // soft limit here instead of inheriting the parent thread's value. + vm.c_stack_soft_limit + .set(VirtualMachine::calculate_c_stack_soft_limit()); enter_vm(vm, || f(vm)) } } diff --git a/extra_tests/snippets/stdlib_threading.py b/extra_tests/snippets/stdlib_threading.py new file mode 100644 index 000000000..f35d7e9d0 --- /dev/null +++ b/extra_tests/snippets/stdlib_threading.py @@ -0,0 +1,77 @@ +import multiprocessing +import os +import threading + + +def import_in_thread(module_name): + outcome = {} + error = {} + + def worker(): + try: + module = __import__(module_name, fromlist=["*"]) + outcome["name"] = module.__name__ + except Exception as exc: + error["exc"] = exc + + thread = threading.Thread(target=worker) + thread.start() + thread.join(timeout=5) + assert not thread.is_alive(), "thread did not finish in time" + if "exc" in error: + raise error["exc"] + + assert outcome["name"] == module_name + + +def run_exec(code): + result = {} + error = {} + + def worker(): + try: + scope = {"__builtins__": __builtins__} + exec(code, scope, scope) # noqa: S102 - intentional threaded exec regression test + result["scope"] = scope + except Exception as exc: + error["exc"] = exc + + thread = threading.Thread(target=worker) + thread.start() + thread.join(timeout=5) + assert not thread.is_alive(), "thread did not finish in time" + if "exc" in error: + raise error["exc"] + return result["scope"] + + +def child_process(): + return None + + +def start_fork_process_after_thread(): + if not hasattr(os, "fork"): + return + + import_in_thread("multiprocessing.connection") + + ctx = multiprocessing.get_context("fork") + process = ctx.Process(target=child_process) + process.start() + process.join(timeout=10) + assert process.exitcode == 0, process.exitcode + + +import_in_thread("functools") +import_in_thread("tempfile") +import_in_thread("multiprocessing.connection") +start_fork_process_after_thread() + +scope = run_exec("import functools") +assert scope["functools"].__name__ == "functools" + +scope = run_exec("from collections import namedtuple") +assert scope["namedtuple"].__name__ == "namedtuple" + +scope = run_exec("module = __import__('multiprocessing.connection', fromlist=['*'])") +assert scope["module"].__name__ == "multiprocessing.connection" diff --git a/scripts/dis_dump.py b/scripts/dis_dump.py index 4a8039726..e30a8955f 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/ @@ -43,6 +43,7 @@ _JUMP_OPNAMES = frozenset( ) _JUMP_OPCODES = None +_ABSOLUTE_JUMP_OPCODES = frozenset(getattr(dis, "hasjabs", ())) def _jump_opcodes(): @@ -74,11 +75,6 @@ def _normalize_argrepr(argrepr): if idx >= 0: name = name[:idx] return "" % name.rstrip(">").strip() - # Normalize COMPARE_OP: strip bool(...) wrapper from CPython 3.14 - # e.g. "bool(==)" -> "==", "bool(<)" -> "<" - m = re.match(r"^bool\((.+)\)$", argrepr) - if m: - return m.group(1) # Remove memory addresses from other reprs argrepr = re.sub(r" at 0x[0-9a-fA-F]+", "", argrepr) # Remove LOAD_ATTR/LOAD_SUPER_ATTR suffixes: " + NULL|self", " + NULL" @@ -167,42 +163,58 @@ 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 = [] - 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) + 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) - # 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 + explicit_offsets = {inst.offset for inst in raw} + cache_counts = {} + stream = [] + offset_to_idx = {} + 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,43 +244,29 @@ 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) + if inst.opcode in _ABSOLUTE_JUMP_OPCODES: + target_off = inst.arg * 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]) elif inst.opname == "COMPARE_OP": - # Normalize COMPARE_OP across interpreters (different encodings) if _IS_RUSTPYTHON: - cmp_str = _RP_CMP_OPS.get(inst.arg, inst.argrepr) + cmp_idx = inst.arg >> 5 if isinstance(inst.arg, int) else inst.arg + cmp_str = _RP_CMP_OPS.get(cmp_idx, inst.argrepr) + if isinstance(inst.arg, int) and inst.arg & 16: + cmp_str = f"bool({cmp_str})" else: - cmp_str = ( - _normalize_argrepr(inst.argrepr) if inst.argrepr else str(inst.arg) - ) + cmp_str = inst.argrepr if inst.argrepr else str(inst.arg) result.append([opname, cmp_str]) elif inst.arg is not None and inst.argrepr: # If argrepr is just a number, try to resolve it via fallback