diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py index fabe0c971c1..d7145cd6fb2 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -99,8 +99,6 @@ def woohoo(): raise ZeroDivisionError() self.assertEqual(state, [1, 42, 999]) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_contextmanager_traceback(self): @contextmanager def f(): diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py index c3c303ad9a7..95d0bd7f92b 100644 --- a/Lib/test/test_contextlib_async.py +++ b/Lib/test/test_contextlib_async.py @@ -252,7 +252,6 @@ async def woohoo(): raise ZeroDivisionError(999) self.assertEqual(state, [1, 42, 999]) - @unittest.expectedFailure # TODO: RUSTPYTHON async def test_contextmanager_except_stopiter(self): @asynccontextmanager async def woohoo(): diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index 4d9713a9027..27d1394e516 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -132,7 +132,6 @@ def f(): self.assertInBytecode(f, 'LOAD_CONST', None) self.check_lnotab(f) - @unittest.expectedFailure # TODO: RUSTPYTHON; RETURN_VALUE def test_while_one(self): # Skip over: LOAD_CONST trueconst POP_JUMP_IF_FALSE xx def f(): @@ -545,7 +544,6 @@ def f(cond, true_value, false_value): self.assertEqual(len(returns), 2) self.check_lnotab(f) - @unittest.expectedFailure # TODO: RUSTPYTHON; absolute jump encoding def test_elim_jump_to_uncond_jump(self): # POP_JUMP_IF_FALSE to JUMP_FORWARD --> POP_JUMP_IF_FALSE to non-jump def f(): diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 3b94a36b3ee..bdc7fb7f1e2 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -154,6 +154,10 @@ struct Compiler { /// True when compiling in "single" (interactive) mode. /// Expression statements at module scope emit CALL_INTRINSIC_1(Print). interactive: bool, + /// Counter for dead-code elimination during constant folding. + /// When > 0, the compiler walks AST (consuming sub_tables) but emits no bytecode. + /// Mirrors CPython's `c_do_not_emit_bytecode`. + do_not_emit_bytecode: u32, } #[derive(Clone, Copy)] @@ -203,6 +207,12 @@ impl CompileContext { } } +/// Segment of a parsed %-format string for optimize_format_str. +struct FormatSegment { + literal: String, + conversion: Option, +} + #[derive(Debug, Clone, Copy, PartialEq)] enum ComprehensionType { Generator, @@ -465,6 +475,7 @@ impl Compiler { opts, in_annotation: false, interactive: false, + do_not_emit_bytecode: 0, } } @@ -1021,8 +1032,8 @@ impl Compiler { } /// Check if this is an inlined comprehension context (PEP 709). - /// Only inline in function-like scopes (fastlocals) — module/class - /// level uses STORE_NAME which is incompatible with LOAD_FAST_AND_CLEAR. + /// PEP 709: Inline comprehensions in function-like scopes. + /// TODO: Module/class scope inlining needs more work (Cell name resolution edge cases). /// Generator expressions are never inlined. fn is_inlined_comprehension_context(&self, comprehension_type: ComprehensionType) -> bool { if comprehension_type == ComprehensionType::Generator { @@ -2360,43 +2371,7 @@ impl Compiler { .. }) => { self.enter_conditional_block(); - match elif_else_clauses.as_slice() { - // Only if - [] => { - let after_block = self.new_block(); - self.compile_jump_if(test, false, after_block)?; - self.compile_statements(body)?; - self.switch_to_block(after_block); - } - // If, elif*, elif/else - [rest @ .., tail] => { - let after_block = self.new_block(); - let mut next_block = self.new_block(); - - self.compile_jump_if(test, false, next_block)?; - self.compile_statements(body)?; - emit!(self, PseudoInstruction::Jump { delta: after_block }); - - for clause in rest { - self.switch_to_block(next_block); - next_block = self.new_block(); - if let Some(test) = &clause.test { - self.compile_jump_if(test, false, next_block)?; - } else { - unreachable!() // must be elif - } - self.compile_statements(&clause.body)?; - emit!(self, PseudoInstruction::Jump { delta: after_block }); - } - - self.switch_to_block(next_block); - if let Some(test) = &tail.test { - self.compile_jump_if(test, false, after_block)?; - } - self.compile_statements(&tail.body)?; - self.switch_to_block(after_block); - } - } + self.compile_if(test, body, elif_else_clauses)?; self.leave_conditional_block(); } ast::Stmt::While(ast::StmtWhile { @@ -3899,6 +3874,23 @@ impl Compiler { // Set qualname self.set_qualname(); + // PEP 479: Wrap generator/coroutine body with StopIteration handler + let is_gen = is_async || self.current_symbol_table().is_generator; + let stop_iteration_block = if is_gen { + let handler_block = self.new_block(); + emit!( + self, + PseudoInstruction::SetupCleanup { + delta: handler_block + } + ); + self.set_no_location(); + self.push_fblock(FBlockType::StopIteration, handler_block, handler_block)?; + Some(handler_block) + } else { + None + }; + // Handle docstring - store in co_consts[0] if present let (doc_str, body) = split_doc(body, &self.opts); if let Some(doc) = &doc_str { @@ -3929,6 +3921,23 @@ impl Compiler { self.arg_constant(ConstantData::None); } + // Close StopIteration handler and emit handler code + if let Some(handler_block) = stop_iteration_block { + emit!(self, PseudoInstruction::PopBlock); + self.set_no_location(); + self.pop_fblock(FBlockType::StopIteration); + self.switch_to_block(handler_block); + emit!( + self, + Instruction::CallIntrinsic1 { + func: oparg::IntrinsicFunction1::StopIterationError + } + ); + self.set_no_location(); + emit!(self, Instruction::Reraise { depth: 1u32 }); + self.set_no_location(); + } + // Exit scope and create function object let code = self.exit_scope(); self.ctx = prev_ctx; @@ -5190,6 +5199,94 @@ impl Compiler { self.store_name(name) } + /// Compile an if statement with constant condition elimination. + /// = compiler_if in CPython codegen.c + fn compile_if( + &mut self, + test: &ast::Expr, + body: &[ast::Stmt], + elif_else_clauses: &[ast::ElifElseClause], + ) -> CompileResult<()> { + let constant = Self::expr_constant(test); + + // If the test is constant false, walk the body (consuming sub_tables) + // but don't emit bytecode + if constant == Some(false) { + self.emit_nop(); + self.do_not_emit_bytecode += 1; + self.compile_statements(body)?; + self.do_not_emit_bytecode -= 1; + // Compile the elif/else chain (if any) + match elif_else_clauses { + [] => {} + [first, rest @ ..] => { + if let Some(elif_test) = &first.test { + self.compile_if(elif_test, &first.body, rest)?; + } else { + self.compile_statements(&first.body)?; + } + } + } + return Ok(()); + } + + // If the test is constant true, compile body directly, + // but walk elif/else without emitting (including elif tests to consume sub_tables) + if constant == Some(true) { + self.emit_nop(); + self.compile_statements(body)?; + self.do_not_emit_bytecode += 1; + for clause in elif_else_clauses { + if let Some(elif_test) = &clause.test { + self.compile_expression(elif_test)?; + } + self.compile_statements(&clause.body)?; + } + self.do_not_emit_bytecode -= 1; + return Ok(()); + } + + // Non-constant test: normal compilation + match elif_else_clauses { + // Only if + [] => { + let after_block = self.new_block(); + self.compile_jump_if(test, false, after_block)?; + self.compile_statements(body)?; + self.switch_to_block(after_block); + } + // If, elif*, elif/else + [rest @ .., tail] => { + let after_block = self.new_block(); + let mut next_block = self.new_block(); + + self.compile_jump_if(test, false, next_block)?; + self.compile_statements(body)?; + emit!(self, PseudoInstruction::Jump { delta: after_block }); + + for clause in rest { + self.switch_to_block(next_block); + next_block = self.new_block(); + if let Some(test) = &clause.test { + self.compile_jump_if(test, false, next_block)?; + } else { + unreachable!() // must be elif + } + self.compile_statements(&clause.body)?; + emit!(self, PseudoInstruction::Jump { delta: after_block }); + } + + self.switch_to_block(next_block); + if let Some(test) = &tail.test { + self.compile_jump_if(test, false, after_block)?; + } + self.compile_statements(&tail.body)?; + self.switch_to_block(after_block); + } + } + Ok(()) + } + fn compile_while( &mut self, test: &ast::Expr, @@ -5198,17 +5295,37 @@ impl Compiler { ) -> CompileResult<()> { self.enter_conditional_block(); + let constant = Self::expr_constant(test); + + // while False: body → walk body (consuming sub_tables) but don't emit, + // then compile orelse + if constant == Some(false) { + self.emit_nop(); + let while_block = self.new_block(); + let after_block = self.new_block(); + self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?; + self.do_not_emit_bytecode += 1; + self.compile_statements(body)?; + self.do_not_emit_bytecode -= 1; + self.pop_fblock(FBlockType::WhileLoop); + self.compile_statements(orelse)?; + self.leave_conditional_block(); + return Ok(()); + } + let while_block = self.new_block(); let else_block = self.new_block(); let after_block = self.new_block(); - // Note: SetupLoop is no longer emitted (break/continue use direct jumps) self.switch_to_block(while_block); - - // Push fblock for while loop self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?; - self.compile_jump_if(test, false, else_block)?; + // while True: → no condition test, just NOP + if constant == Some(true) { + self.emit_nop(); + } else { + self.compile_jump_if(test, false, else_block)?; + } let was_in_loop = self.ctx.loop_data.replace((while_block, after_block)); self.compile_statements(body)?; @@ -5216,9 +5333,7 @@ impl Compiler { emit!(self, PseudoInstruction::Jump { delta: while_block }); self.switch_to_block(else_block); - // Pop fblock self.pop_fblock(FBlockType::WhileLoop); - // Note: PopBlock is no longer emitted for loops self.compile_statements(orelse)?; self.switch_to_block(after_block); @@ -7006,6 +7121,40 @@ impl Compiler { }) => { self.compile_jump_if(operand, !condition, target_block)?; } + // `x is None` / `x is not None` → POP_JUMP_IF_NONE / POP_JUMP_IF_NOT_NONE + ast::Expr::Compare(ast::ExprCompare { + left, + ops, + comparators, + .. + }) if ops.len() == 1 + && matches!(ops[0], ast::CmpOp::Is | ast::CmpOp::IsNot) + && comparators.len() == 1 + && matches!(&comparators[0], ast::Expr::NoneLiteral(_)) => + { + self.compile_expression(left)?; + let is_not = matches!(ops[0], ast::CmpOp::IsNot); + // is None + jump_if_false → POP_JUMP_IF_NOT_NONE + // is None + jump_if_true → POP_JUMP_IF_NONE + // is not None + jump_if_false → POP_JUMP_IF_NONE + // is not None + jump_if_true → POP_JUMP_IF_NOT_NONE + let jump_if_none = condition != is_not; + if jump_if_none { + emit!( + self, + Instruction::PopJumpIfNone { + delta: target_block, + } + ); + } else { + emit!( + self, + Instruction::PopJumpIfNotNone { + delta: target_block, + } + ); + } + } _ => { // Fall back case which always will work! self.compile_expression(expression)?; @@ -7103,14 +7252,14 @@ impl Compiler { let has_unpacking = items.iter().any(|item| item.key.is_none()); if !has_unpacking { - if items.len() >= 16 { - emit!(self, Instruction::BuildMap { count: 0 }); - for item in items { - self.compile_expression(item.key.as_ref().unwrap())?; - self.compile_expression(&item.value)?; - emit!(self, Instruction::MapAdd { i: 1 }); - } - } else { + // Match CPython's compiler_subdict chunking strategy: + // - n≤15: BUILD_MAP n (all pairs on stack) + // - n>15: BUILD_MAP 0 + MAP_ADD chunks of 17, last chunk uses + // BUILD_MAP n (if ≤15) or BUILD_MAP 0 + MAP_ADD + const STACK_LIMIT: usize = 15; + const BIG_MAP_CHUNK: usize = 17; + + if items.len() <= STACK_LIMIT { for item in items { self.compile_expression(item.key.as_ref().unwrap())?; self.compile_expression(&item.value)?; @@ -7121,6 +7270,63 @@ impl Compiler { count: u32::try_from(items.len()).expect("too many dict items"), } ); + } else { + // Split: leading full chunks of BIG_MAP_CHUNK via MAP_ADD, + // remainder via BUILD_MAP n or MAP_ADD depending on size + let n = items.len(); + let remainder = n % BIG_MAP_CHUNK; + let n_big_chunks = n / BIG_MAP_CHUNK; + // If remainder fits on stack (≤15), use BUILD_MAP n for it. + // Otherwise it becomes another MAP_ADD chunk. + let (big_count, tail_count) = if remainder > 0 && remainder <= STACK_LIMIT { + (n_big_chunks, remainder) + } else { + // remainder is 0 or >15: all chunks are MAP_ADD chunks + let total_map_add = if remainder == 0 { + n_big_chunks + } else { + n_big_chunks + 1 + }; + (total_map_add, 0usize) + }; + + emit!(self, Instruction::BuildMap { count: 0 }); + + let mut idx = 0; + for chunk_i in 0..big_count { + if chunk_i > 0 { + emit!(self, Instruction::BuildMap { count: 0 }); + } + let chunk_size = if idx + BIG_MAP_CHUNK <= n - tail_count { + BIG_MAP_CHUNK + } else { + n - tail_count - idx + }; + for item in &items[idx..idx + chunk_size] { + self.compile_expression(item.key.as_ref().unwrap())?; + self.compile_expression(&item.value)?; + emit!(self, Instruction::MapAdd { i: 1 }); + } + if chunk_i > 0 { + emit!(self, Instruction::DictUpdate { i: 1 }); + } + idx += chunk_size; + } + + // Tail: remaining pairs via BUILD_MAP n + DICT_UPDATE + if tail_count > 0 { + for item in &items[idx..idx + tail_count] { + self.compile_expression(item.key.as_ref().unwrap())?; + self.compile_expression(&item.value)?; + } + emit!( + self, + Instruction::BuildMap { + count: tail_count.to_u32(), + } + ); + emit!(self, Instruction::DictUpdate { i: 1 }); + } } return Ok(()); } @@ -7275,6 +7481,15 @@ impl Compiler { ast::Expr::BinOp(ast::ExprBinOp { left, op, right, .. }) => { + // optimize_format_str: 'format' % (args,) → f-string bytecode + if matches!(op, ast::Operator::Mod) + && let ast::Expr::StringLiteral(s) = left.as_ref() + && let ast::Expr::Tuple(ast::ExprTuple { elts, .. }) = right.as_ref() + && !elts.iter().any(|e| matches!(e, ast::Expr::Starred(_))) + && self.try_optimize_format_str(s.value.to_str(), elts, range)? + { + return Ok(()); + } self.compile_expression(left)?; self.compile_expression(right)?; @@ -7879,8 +8094,10 @@ impl Compiler { .any(|arg| matches!(arg, ast::Expr::Starred(_))); let has_double_star = arguments.keywords.iter().any(|k| k.arg.is_none()); - // Check if exceeds stack guideline - let too_big = nelts + nkwelts * 2 > 8; + // Check if exceeds stack guideline (STACK_USE_GUIDELINE / 2 = 15) + // With CALL_KW, kwargs values go on stack but keys go in a const tuple, + // so stack usage is: func + null + positional_args + kwarg_values + kwnames_tuple + let too_big = nelts + nkwelts > 15; if !has_starred && !has_double_star && !too_big { // Simple call path: no * or ** args @@ -8611,6 +8828,9 @@ impl Compiler { // Low level helper functions: fn _emit>(&mut self, instr: I, arg: OpArg, target: BlockIdx) { + if self.do_not_emit_bytecode > 0 { + return; + } let range = self.current_source_range; let source = self.source_file.to_source_code(); let location = source.source_location(range.start(), PositionEncoding::Utf8); @@ -8628,6 +8848,14 @@ impl Compiler { }); } + /// Mark the last emitted instruction as having no source location. + /// Prevents it from triggering LINE events in sys.monitoring. + fn set_no_location(&mut self) { + if let Some(last) = self.current_block().instructions.last_mut() { + last.lineno_override = Some(-1); + } + } + fn emit_no_arg>(&mut self, ins: I) { self._emit(ins, OpArg::NULL, BlockIdx::NULL) } @@ -8809,6 +9037,44 @@ impl Compiler { self.code_stack.last_mut().expect("no code on stack") } + /// Evaluate whether an expression is a compile-time constant boolean. + /// Returns Some(true) for truthy constants, Some(false) for falsy constants, + /// None for non-constant expressions. + /// = expr_constant in CPython compile.c + fn expr_constant(expr: &ast::Expr) -> Option { + match expr { + ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => Some(*value), + ast::Expr::NoneLiteral(_) => Some(false), + ast::Expr::EllipsisLiteral(_) => Some(true), + ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => match value { + ast::Number::Int(i) => { + let n: i64 = i.as_i64().unwrap_or(1); + Some(n != 0) + } + ast::Number::Float(f) => Some(*f != 0.0), + ast::Number::Complex { real, imag, .. } => Some(*real != 0.0 || *imag != 0.0), + }, + ast::Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => { + Some(!value.to_str().is_empty()) + } + ast::Expr::BytesLiteral(ast::ExprBytesLiteral { value, .. }) => { + Some(value.bytes().next().is_some()) + } + ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => { + if elts.is_empty() { + Some(false) + } else { + None // non-empty tuples may have side effects in elements + } + } + _ => None, + } + } + + fn emit_nop(&mut self) { + emit!(self, Instruction::Nop); + } + /// Enter a conditional block (if/for/while/match/try/with) /// PEP 649: Track conditional annotation context fn enter_conditional_block(&mut self) { @@ -8830,6 +9096,35 @@ impl Compiler { range: ruff_text_size::TextRange, is_break: bool, ) -> CompileResult<()> { + if self.do_not_emit_bytecode > 0 { + // Still validate that we're inside a loop even in dead code + let code = self.current_code_info(); + let mut found_loop = false; + for i in (0..code.fblock.len()).rev() { + match code.fblock[i].fb_type { + FBlockType::WhileLoop | FBlockType::ForLoop => { + found_loop = true; + break; + } + FBlockType::ExceptionGroupHandler => { + return Err(self.error_ranged( + CodegenErrorType::BreakContinueReturnInExceptStar, + range, + )); + } + _ => {} + } + } + if !found_loop { + if is_break { + return Err(self.error_ranged(CodegenErrorType::InvalidBreak, range)); + } else { + return Err(self.error_ranged(CodegenErrorType::InvalidContinue, range)); + } + } + return Ok(()); + } + // unwind_fblock_stack // We need to unwind fblocks and compile cleanup code. For FinallyTry blocks, // we need to compile the finally body inline, but we must temporarily pop @@ -9159,6 +9454,125 @@ impl Compiler { self.compile_fstring_elements(fstring.flags, &fstring.elements) } + /// Optimize `'format_str' % (args,)` into f-string bytecode. + /// Returns true if optimization was applied, false to fall back to normal BINARY_OP %. + /// Matches CPython's codegen.c `compiler_formatted_value` optimization. + fn try_optimize_format_str( + &mut self, + format_str: &str, + args: &[ast::Expr], + range: ruff_text_size::TextRange, + ) -> CompileResult { + // Parse format string into segments + let Some(segments) = Self::parse_percent_format(format_str) else { + return Ok(false); + }; + + // Verify arg count matches specifier count + let spec_count = segments.iter().filter(|s| s.conversion.is_some()).count(); + if spec_count != args.len() { + return Ok(false); + } + + self.set_source_range(range); + + // Special case: no specifiers, just %% escaping → constant fold + if spec_count == 0 { + let folded: String = segments.iter().map(|s| s.literal.as_str()).collect(); + self.emit_load_const(ConstantData::Str { + value: folded.into(), + }); + return Ok(true); + } + + // Emit f-string style bytecode + let mut part_count: u32 = 0; + let mut arg_idx = 0; + + for seg in &segments { + if !seg.literal.is_empty() { + self.emit_load_const(ConstantData::Str { + value: seg.literal.clone().into(), + }); + part_count += 1; + } + if let Some(conv) = seg.conversion { + self.compile_expression(&args[arg_idx])?; + self.set_source_range(range); + emit!(self, Instruction::ConvertValue { oparg: conv }); + emit!(self, Instruction::FormatSimple); + part_count += 1; + arg_idx += 1; + } + } + + if part_count == 0 { + self.emit_load_const(ConstantData::Str { + value: String::new().into(), + }); + } else if part_count > 1 { + emit!(self, Instruction::BuildString { count: part_count }); + } + + Ok(true) + } + + /// Parse a %-format string into segments of (literal_prefix, optional conversion). + /// Returns None if the format string contains unsupported specifiers. + fn parse_percent_format(format_str: &str) -> Option> { + let mut segments = Vec::new(); + let mut chars = format_str.chars().peekable(); + let mut current_literal = String::new(); + + while let Some(ch) = chars.next() { + if ch == '%' { + match chars.peek() { + Some('%') => { + chars.next(); + current_literal.push('%'); + } + Some('s') => { + chars.next(); + segments.push(FormatSegment { + literal: core::mem::take(&mut current_literal), + conversion: Some(oparg::ConvertValueOparg::Str), + }); + } + Some('r') => { + chars.next(); + segments.push(FormatSegment { + literal: core::mem::take(&mut current_literal), + conversion: Some(oparg::ConvertValueOparg::Repr), + }); + } + Some('a') => { + chars.next(); + segments.push(FormatSegment { + literal: core::mem::take(&mut current_literal), + conversion: Some(oparg::ConvertValueOparg::Ascii), + }); + } + _ => { + // Unsupported: %d, %f, %(name)s, %10s, etc. + return None; + } + } + } else { + current_literal.push(ch); + } + } + + // Trailing literal + if !current_literal.is_empty() { + segments.push(FormatSegment { + literal: current_literal, + conversion: None, + }); + } + + Some(segments) + } + fn compile_fstring_elements( &mut self, flags: ast::FStringFlags, @@ -9419,32 +9833,34 @@ impl EmitArg for BlockIdx { // = _PyCompile_CleanDoc fn clean_doc(doc: &str) -> String { let doc = expandtabs(doc, 8); - // First pass: find minimum indentation of any non-blank lines - // after first line. + // First pass: find minimum indentation of non-blank lines AFTER the first line. + // A "blank line" is one containing only spaces (or empty). let margin = doc - .lines() - // Find the non-blank lines - .filter(|line| !line.trim().is_empty()) - // get the one with the least indentation - .map(|line| line.chars().take_while(|c| c == &' ').count()) - .min(); - if let Some(margin) = margin { - let mut cleaned = String::with_capacity(doc.len()); - // copy first line without leading whitespace - if let Some(first_line) = doc.lines().next() { - cleaned.push_str(first_line.trim_start()); - } - // copy subsequent lines without margin. - for line in doc.split('\n').skip(1) { - cleaned.push('\n'); - let cleaned_line = line.chars().skip(margin).collect::(); - cleaned.push_str(&cleaned_line); - } - - cleaned - } else { - doc.to_owned() - } + .split('\n') + .skip(1) // skip first line + .filter(|line| line.chars().any(|c| c != ' ')) // non-blank lines only + .map(|line| line.chars().take_while(|c| *c == ' ').count()) + .min() + .unwrap_or(0); + + let mut cleaned = String::with_capacity(doc.len()); + // Strip all leading spaces from the first line + if let Some(first_line) = doc.split('\n').next() { + let trimmed = first_line.trim_start(); + // Early exit: no leading spaces on first line AND margin == 0 + if trimmed.len() == first_line.len() && margin == 0 { + return doc.to_owned(); + } + cleaned.push_str(trimmed); + } + // Subsequent lines: skip up to `margin` leading spaces + for line in doc.split('\n').skip(1) { + cleaned.push('\n'); + let skip = line.chars().take(margin).take_while(|c| *c == ' ').count(); + cleaned.push_str(&line[skip..]); + } + + cleaned } // copied from rustpython_common::str, so we don't have to depend on it just for this function diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 73da736ef44..52aa9b0c3e0 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -3,7 +3,7 @@ use core::ops; use crate::{IndexMap, IndexSet, error::InternalError}; use malachite_bigint::BigInt; -use num_traits::ToPrimitive; +use num_traits::{ToPrimitive, Zero}; use rustpython_compiler_core::{ OneIndexed, SourceLocation, @@ -190,8 +190,11 @@ impl CodeInfo { opts: &crate::compile::CompileOpts, ) -> crate::InternalResult { // Constant folding passes + self.fold_binop_constants(); + self.remove_nops(); self.fold_unary_negative(); - self.remove_nops(); // remove NOPs from unary folding so tuple/list/set see contiguous LOADs + self.fold_binop_constants(); // re-run after unary folding: -1 + 2 → 1 + self.remove_nops(); // remove NOPs so tuple/list/set see contiguous LOADs self.fold_tuple_constants(); self.fold_list_constants(); self.fold_set_constants(); @@ -206,19 +209,21 @@ impl CodeInfo { // Peephole optimizer creates superinstructions matching CPython self.peephole_optimize(); - // Always apply LOAD_FAST_BORROW optimization - self.optimize_load_fast_borrow(); - - // Post-codegen CFG analysis passes (flowgraph.c pipeline) + // Phase 1: _PyCfg_OptimizeCodeUnit (flowgraph.c) mark_except_handlers(&mut self.blocks); label_exception_targets(&mut self.blocks); + // TODO: insert_superinstructions disabled pending StoreFastLoadFast VM fix push_cold_blocks_to_end(&mut self.blocks); + + // Phase 2: _PyCfg_OptimizedCfgToInstructionSequence (flowgraph.c) normalize_jumps(&mut self.blocks); self.dce(); // re-run within-block DCE after normalize_jumps creates new instructions self.eliminate_unreachable_blocks(); duplicate_end_returns(&mut self.blocks); self.dce(); // truncate after terminal in blocks that got return duplicated self.eliminate_unreachable_blocks(); // remove now-unreachable last block + // optimize_load_fast: after normalize_jumps + self.optimize_load_fast_borrow(); self.optimize_load_global_push_null(); let max_stackdepth = self.max_stackdepth()?; @@ -455,8 +460,9 @@ impl CodeInfo { let cache_count = info.cache_entries as usize; let (extras, lo_arg) = info.arg.split(); + let loc_pair = (info.location, info.end_location); locations.extend(core::iter::repeat_n( - (info.location, info.end_location), + loc_pair, info.arg.instr_size() + cache_count, )); // Collect linetable locations with lineno_override support @@ -687,6 +693,228 @@ impl CodeInfo { } } + /// Constant folding: fold LOAD_CONST/LOAD_SMALL_INT + LOAD_CONST/LOAD_SMALL_INT + BINARY_OP + /// into a single LOAD_CONST when the result is computable at compile time. + /// = fold_binops_on_constants in CPython flowgraph.c + fn fold_binop_constants(&mut self) { + use oparg::BinaryOperator as BinOp; + + for block in &mut self.blocks { + let mut i = 0; + while i + 2 < block.instructions.len() { + // Check pattern: LOAD_CONST/LOAD_SMALL_INT, LOAD_CONST/LOAD_SMALL_INT, BINARY_OP + let Some(Instruction::BinaryOp { .. }) = block.instructions[i + 2].instr.real() + else { + i += 1; + continue; + }; + + let op_raw = u32::from(block.instructions[i + 2].arg); + let Ok(op) = BinOp::try_from(op_raw) else { + i += 1; + continue; + }; + + let left = Self::get_const_value_from(&self.metadata, &block.instructions[i]); + let right = Self::get_const_value_from(&self.metadata, &block.instructions[i + 1]); + + let (Some(left_val), Some(right_val)) = (left, right) else { + i += 1; + continue; + }; + + let result = Self::eval_binop(&left_val, &right_val, op); + + if let Some(result_const) = result { + // Check result size limit (CPython limits to 4096 bytes) + if Self::const_too_big(&result_const) { + i += 1; + continue; + } + let (const_idx, _) = self.metadata.consts.insert_full(result_const); + // Replace first instruction with LOAD_CONST result + block.instructions[i].instr = Instruction::LoadConst { + consti: Arg::marker(), + } + .into(); + block.instructions[i].arg = OpArg::new(const_idx as u32); + // NOP out the second and third instructions + let loc = block.instructions[i].location; + let end_loc = block.instructions[i].end_location; + block.instructions[i + 1].instr = Instruction::Nop.into(); + block.instructions[i + 1].location = loc; + block.instructions[i + 1].end_location = end_loc; + block.instructions[i + 2].instr = Instruction::Nop.into(); + block.instructions[i + 2].location = loc; + block.instructions[i + 2].end_location = end_loc; + // Don't advance - check if the result can be folded again + // (e.g., 2 ** 31 - 1) + i = i.saturating_sub(1); // re-check with previous instruction + } else { + i += 1; + } + } + } + } + + fn get_const_value_from( + metadata: &CodeUnitMetadata, + info: &InstructionInfo, + ) -> Option { + match info.instr.real() { + Some(Instruction::LoadConst { .. }) => { + let idx = u32::from(info.arg) as usize; + metadata.consts.get_index(idx).cloned() + } + Some(Instruction::LoadSmallInt { .. }) => { + let v = u32::from(info.arg) as i32; + Some(ConstantData::Integer { + value: BigInt::from(v), + }) + } + _ => None, + } + } + + fn eval_binop( + left: &ConstantData, + right: &ConstantData, + op: oparg::BinaryOperator, + ) -> Option { + use oparg::BinaryOperator as BinOp; + match (left, right) { + (ConstantData::Integer { value: l }, ConstantData::Integer { value: r }) => { + let result = match op { + BinOp::Add => l + r, + BinOp::Subtract => l - r, + BinOp::Multiply => l * r, + BinOp::FloorDivide => { + if r.is_zero() { + return None; + } + // Python floor division: round towards negative infinity + let (q, rem) = (l.clone() / r.clone(), l.clone() % r.clone()); + if !rem.is_zero() && (rem < BigInt::from(0)) != (*r < BigInt::from(0)) { + q - 1 + } else { + q + } + } + BinOp::Remainder => { + if r.is_zero() { + return None; + } + // Python modulo: result has same sign as divisor + let rem = l.clone() % r.clone(); + if !rem.is_zero() && (rem < BigInt::from(0)) != (*r < BigInt::from(0)) { + rem + r + } else { + rem + } + } + BinOp::Power => { + let exp: u32 = r.try_into().ok()?; + if exp > 128 { + return None; + } // prevent huge results + num_traits::pow::pow(l.clone(), exp as usize) + } + BinOp::Lshift => { + let shift: u32 = r.try_into().ok()?; + if shift > 128 { + return None; + } + l << (shift as usize) + } + BinOp::Rshift => { + let shift: u32 = r.try_into().ok()?; + l >> (shift as usize) + } + BinOp::And => l & r, + BinOp::Or => l | r, + BinOp::Xor => l ^ r, + _ => return None, + }; + Some(ConstantData::Integer { value: result }) + } + (ConstantData::Float { value: l }, ConstantData::Float { value: r }) => { + let result = match op { + BinOp::Add => l + r, + BinOp::Subtract => l - r, + BinOp::Multiply => l * r, + BinOp::TrueDivide => { + if *r == 0.0 { + return None; + } + l / r + } + BinOp::FloorDivide => { + // Float floor division uses runtime semantics; skip folding + return None; + } + BinOp::Remainder => { + // Float modulo uses fmod() at runtime; Rust arithmetic differs + return None; + } + BinOp::Power => l.powf(*r), + _ => return None, + }; + if !result.is_finite() { + return None; + } + Some(ConstantData::Float { value: result }) + } + // Int op Float or Float op Int → Float + (ConstantData::Integer { value: l }, ConstantData::Float { value: r }) => { + let l_f = l.to_f64()?; + Self::eval_binop( + &ConstantData::Float { value: l_f }, + &ConstantData::Float { value: *r }, + op, + ) + } + (ConstantData::Float { value: l }, ConstantData::Integer { value: r }) => { + let r_f = r.to_f64()?; + Self::eval_binop( + &ConstantData::Float { value: *l }, + &ConstantData::Float { value: r_f }, + op, + ) + } + // String concatenation and repetition + (ConstantData::Str { value: l }, ConstantData::Str { value: r }) + if matches!(op, BinOp::Add) => + { + let mut result = l.to_string(); + result.push_str(&r.to_string()); + Some(ConstantData::Str { + value: result.into(), + }) + } + (ConstantData::Str { value: s }, ConstantData::Integer { value: n }) + if matches!(op, BinOp::Multiply) => + { + let n: usize = n.try_into().ok()?; + if n > 4096 { + return None; + } + let result = s.to_string().repeat(n); + Some(ConstantData::Str { + value: result.into(), + }) + } + _ => None, + } + } + + fn const_too_big(c: &ConstantData) -> bool { + match c { + ConstantData::Integer { value } => value.bits() > 4096 * 8, + ConstantData::Str { value } => value.len() > 4096, + _ => false, + } + } + /// Constant folding: fold LOAD_CONST/LOAD_SMALL_INT + BUILD_TUPLE into LOAD_CONST tuple /// fold_tuple_of_constants fn fold_tuple_constants(&mut self) { @@ -1067,6 +1295,9 @@ impl CodeInfo { None } } + // Note: StoreFast + LoadFast → StoreFastLoadFast is done in a + // separate pass AFTER optimize_load_fast_borrow, because CPython + // only combines STORE_FAST + LOAD_FAST (not LOAD_FAST_BORROW). (Instruction::LoadConst { consti }, Instruction::ToBool) => { let consti = consti.get(curr.arg); let constant = &self.metadata.consts[consti.as_usize()]; @@ -1253,10 +1484,48 @@ impl CodeInfo { /// Optimize LOAD_FAST to LOAD_FAST_BORROW where safe. /// - /// A LOAD_FAST can be converted to LOAD_FAST_BORROW if its value is - /// consumed within the same basic block (not passed to another block). - /// This is a reference counting optimization in CPython; in RustPython - /// we implement it for bytecode compatibility. + /// insert_superinstructions (flowgraph.c): Combine STORE_FAST + LOAD_FAST → + /// STORE_FAST_LOAD_FAST. Currently disabled pending VM stack null investigation. + #[allow(dead_code)] + fn combine_store_fast_load_fast(&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::StoreFast { .. }), Some(Instruction::LoadFast { .. })) = + (curr.instr.real(), next.instr.real()) + else { + i += 1; + continue; + }; + // Skip if instructions are on different lines (matching make_super_instruction) + let line1 = curr.location.line; + let line2 = next.location.line; + if line1 != line2 { + i += 1; + continue; + } + let idx1 = u32::from(curr.arg); + let idx2 = u32::from(next.arg); + if idx1 < 16 && idx2 < 16 { + let packed = (idx1 << 4) | idx2; + block.instructions[i].instr = Instruction::StoreFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + block.instructions[i].arg = OpArg::new(packed); + // Replace second instruction with NOP (CPython: INSTR_SET_OP0(inst2, NOP)) + block.instructions[i + 1].instr = Instruction::Nop.into(); + block.instructions[i + 1].arg = OpArg::new(0); + i += 2; // skip the NOP + } else { + i += 1; + } + } + } + } + fn optimize_load_fast_borrow(&mut self) { // NOT_LOCAL marker: instruction didn't come from a LOAD_FAST const NOT_LOCAL: usize = usize::MAX; @@ -1972,6 +2241,42 @@ fn normalize_jumps(blocks: &mut [Block]) { } } + // Replace JUMP → LOAD_CONST + RETURN_VALUE with inline return. + // This matches CPython's optimize_basic_block: "Replace JUMP to a RETURN". + for &block_idx in &visit_order { + let idx = block_idx.idx(); + let mut replacements: Vec<(usize, Vec)> = Vec::new(); + for (i, ins) in blocks[idx].instructions.iter().enumerate() { + if !ins.instr.is_unconditional_jump() || ins.target == BlockIdx::NULL { + continue; + } + let target_block = &blocks[ins.target.idx()]; + // Target must be exactly LOAD_CONST + RETURN_VALUE (2 instructions) + if target_block.instructions.len() >= 2 { + let t0 = &target_block.instructions[0]; + let t1 = &target_block.instructions[1]; + if matches!(t0.instr.real(), Some(Instruction::LoadConst { .. })) + && matches!(t1.instr.real(), Some(Instruction::ReturnValue)) + { + let mut load = *t0; + let mut ret = *t1; + // Use the jump's location for the inlined return + load.location = ins.location; + load.end_location = ins.end_location; + load.except_handler = ins.except_handler; + ret.location = ins.location; + ret.end_location = ins.end_location; + ret.except_handler = ins.except_handler; + replacements.push((i, vec![load, ret])); + } + } + } + // Apply replacements in reverse order + for (i, new_insts) in replacements.into_iter().rev() { + blocks[idx].instructions.splice(i..i + 1, new_insts); + } + } + // 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() { @@ -2172,19 +2477,15 @@ pub(crate) fn label_exception_targets(blocks: &mut [Block]) { blocks[bi].instructions[i].except_handler = handler_info; // Track YIELD_VALUE except stack depth - // Only count for direct yield (arg=0), not yield-from/await (arg=1) - // The yield-from's internal SETUP_FINALLY is not an external except depth + // Record the except stack depth at the point of yield. + // With the StopIteration wrapper, depth is naturally correct: + // - plain yield outside try: depth=1 → DEPTH1 set + // - yield inside try: depth=2+ → no DEPTH1 + // - yield-from/await: has internal SETUP_FINALLY → depth=2+ → no DEPTH1 if let Some(Instruction::YieldValue { .. }) = blocks[bi].instructions[i].instr.real() { - let yield_arg = u32::from(blocks[bi].instructions[i].arg); - if yield_arg == 0 { - // Direct yield: count actual except depth - last_yield_except_depth = stack.len() as i32; - } else { - // yield-from/await: subtract 1 for the internal SETUP_FINALLY - last_yield_except_depth = (stack.len() as i32) - 1; - } + last_yield_except_depth = stack.len() as i32; } // Set RESUME DEPTH1 flag based on last yield's except depth 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 bfdddf4a722..88f81321cda 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: 9780 +assertion_line: 9962 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) @@ -250,6 +250,9 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 232 JUMP_BACKWARD (209) 233 CACHE + 2 234 CALL_INTRINSIC_1 (StopIterationError) + 235 RERAISE (1) + 2 MAKE_FUNCTION 3 STORE_NAME (0, test) 4 LOAD_CONST (None) diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index be7d6e77390..f8d7a2b26ea 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -2169,17 +2169,15 @@ impl SymbolTableBuilder { // Generator expressions need the is_generator flag self.tables.last_mut().unwrap().is_generator = is_generator; - // PEP 709: Mark non-generator comprehensions for inlining, - // but only inside function-like scopes (fastlocals). - // Module/class scope uses STORE_NAME which is incompatible - // with LOAD_FAST_AND_CLEAR / STORE_FAST save/restore. - // Async comprehensions cannot be inlined because they need - // their own coroutine scope. - // Note: tables.last() is the comprehension scope we just pushed, - // so we check the second-to-last for the parent scope. + // PEP 709: Mark non-generator comprehensions for inlining. + // Only in function-like scopes for now. Module/class scope inlining + // needs more work (Cell name resolution, __class__ handling). + // Also excluded: generator expressions, async comprehensions, + // and annotation scopes nested in classes (can_see_class_scope). let element_has_await = expr_contains_await(elt1) || elt2.is_some_and(expr_contains_await); if !is_generator && !has_async_gen && !element_has_await { let parent = self.tables.iter().rev().nth(1); + let parent_can_see_class = parent.is_some_and(|t| t.can_see_class_scope); let parent_is_func = parent.is_some_and(|t| { matches!( t.typ, @@ -2189,8 +2187,7 @@ impl SymbolTableBuilder { | CompilerScope::Comprehension ) }); - let parent_can_see_class = parent.is_some_and(|t| t.can_see_class_scope); - if parent_is_func && !parent_can_see_class { + if !parent_can_see_class && parent_is_func { self.tables.last_mut().unwrap().comp_inlined = true; } } diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs index d099ed251c1..4e354085712 100644 --- a/crates/compiler-core/src/bytecode.rs +++ b/crates/compiler-core/src/bytecode.rs @@ -135,6 +135,72 @@ pub fn decode_exception_table(table: &[u8]) -> Vec { entries } +/// Parse linetable to build a boolean mask indicating which code units +/// have NO_LOCATION (line == -1). Returns a Vec of length `num_units`. +pub fn build_no_location_mask(linetable: &[u8], num_units: usize) -> Vec { + let mut mask = Vec::new(); + mask.resize(num_units, false); + let mut pos = 0; + let mut unit_idx = 0; + + while pos < linetable.len() && unit_idx < num_units { + let header = linetable[pos]; + pos += 1; + let code = (header >> 3) & 0xf; + let length = ((header & 7) + 1) as usize; + + let is_no_location = code == PyCodeLocationInfoKind::None as u8; + + // Skip payload bytes based on location kind + match code { + 0..=9 => pos += 1, // Short forms: 1 byte payload + 10..=12 => pos += 2, // OneLine forms: 2 bytes payload + 13 => { + // NoColumns: signed varint (line delta) + while pos < linetable.len() { + let b = linetable[pos]; + pos += 1; + if b & 0x40 == 0 { + break; + } + } + } + 14 => { + // Long form: signed varint (line delta) + 3 unsigned varints + // line_delta + while pos < linetable.len() { + let b = linetable[pos]; + pos += 1; + if b & 0x40 == 0 { + break; + } + } + // end_line_delta, col+1, end_col+1 + for _ in 0..3 { + while pos < linetable.len() { + let b = linetable[pos]; + pos += 1; + if b & 0x40 == 0 { + break; + } + } + } + } + 15 => {} // None: no payload + _ => {} + } + + for _ in 0..length { + if unit_idx < num_units { + mask[unit_idx] = is_no_location; + unit_idx += 1; + } + } + } + + mask +} + /// CPython 3.11+ linetable location info codes #[derive(Copy, Clone, Debug, PartialEq, Eq)] #[repr(u8)] diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index db20b5fd7a6..88f6fe7411a 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -736,6 +736,20 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; self.store_variable(var_num.get(arg), val) } + Instruction::StoreFastLoadFast { var_nums } => { + let oparg = var_nums.get(arg); + let (store_idx, load_idx) = oparg.indexes(); + let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; + self.store_variable(store_idx, val)?; + let local = self.variables[load_idx] + .as_ref() + .ok_or(JitCompileError::BadBytecode)?; + self.stack.push(JitValue::from_type_and_value( + local.ty.clone(), + self.builder.use_var(local.var), + )); + Ok(()) + } Instruction::StoreFastStoreFast { var_nums } => { let oparg = var_nums.get(arg); let (idx1, idx2) = oparg.indexes(); diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 132fd8e26ec..5c2312015e9 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -1755,12 +1755,6 @@ impl ExecutingFrame<'_> { exc_tb: PyObjectRef, ) -> PyResult { self.monitoring_mask = vm.state.monitoring_events.load(); - // Reset prev_line so that LINE monitoring events fire even if - // the exception handler is on the same line as the yield point. - // In CPython, _Py_call_instrumentation_line has a special case - // for RESUME: it fires LINE even when prev_line == current_line. - // Since gen_throw bypasses RESUME, we reset prev_line instead. - *self.prev_line = 0; if let Some(jen) = self.yield_from_target() { // Check if the exception is GeneratorExit (type or instance). // For GeneratorExit, close the sub-iterator instead of throwing. @@ -1796,7 +1790,10 @@ impl ExecutingFrame<'_> { self.push_value(vm.ctx.none()); vm.contextualize_exception(&err); return match self.unwind_blocks(vm, UnwindReason::Raising { exception: err }) { - Ok(None) => self.run(vm), + Ok(None) => { + *self.prev_line = 0; + self.run(vm) + } Ok(Some(result)) => Ok(result), Err(exception) => Err(exception), }; @@ -1838,7 +1835,10 @@ impl ExecutingFrame<'_> { self.push_value(vm.ctx.none()); vm.contextualize_exception(&err); match self.unwind_blocks(vm, UnwindReason::Raising { exception: err }) { - Ok(None) => self.run(vm), + Ok(None) => { + *self.prev_line = 0; + self.run(vm) + } Ok(Some(result)) => Ok(result), Err(exception) => Err(exception), } @@ -1906,7 +1906,13 @@ impl ExecutingFrame<'_> { self.push_value(vm.ctx.none()); match self.unwind_blocks(vm, UnwindReason::Raising { exception }) { - Ok(None) => self.run(vm), + Ok(None) => { + // Reset prev_line so that the first instruction in the handler + // fires a LINE event. In CPython, gen_send_ex re-enters the + // eval loop which reinitializes its local prev_instr tracker. + *self.prev_line = 0; + self.run(vm) + } Ok(Some(result)) => Ok(result), Err(exception) => { // Fire PY_UNWIND: exception escapes the generator frame. @@ -9440,20 +9446,25 @@ impl ExecutingFrame<'_> { Ok(vm.ctx.new_tuple(list.borrow_vec().to_vec()).into()) } bytecode::IntrinsicFunction1::StopIterationError => { - // Convert StopIteration to RuntimeError - // Used to ensure async generators don't raise StopIteration directly - // _PyGen_FetchStopIterationValue - // Use fast_isinstance to handle subclasses of StopIteration + // Convert StopIteration to RuntimeError (PEP 479) + // Returns the exception object; RERAISE will re-raise it if arg.fast_isinstance(vm.ctx.exceptions.stop_iteration) { - Err(vm.new_runtime_error("coroutine raised StopIteration")) + let flags = &self.code.flags; + let msg = if flags + .contains(bytecode::CodeFlags::COROUTINE | bytecode::CodeFlags::GENERATOR) + { + "async generator raised StopIteration" + } else if flags.contains(bytecode::CodeFlags::COROUTINE) { + "coroutine raised StopIteration" + } else { + "generator raised StopIteration" + }; + let err = vm.new_runtime_error(msg); + err.set___cause__(arg.downcast().ok()); + Ok(err.into()) } else { - // If not StopIteration, just re-raise the original exception - Err(arg.downcast().unwrap_or_else(|obj| { - vm.new_runtime_error(format!( - "unexpected exception type: {:?}", - obj.class() - )) - })) + // Not StopIteration, pass through for RERAISE + Ok(arg) } } bytecode::IntrinsicFunction1::AsyncGenWrap => { diff --git a/crates/vm/src/stdlib/_symtable.rs b/crates/vm/src/stdlib/_symtable.rs index eede3d102c2..299e006bdce 100644 --- a/crates/vm/src/stdlib/_symtable.rs +++ b/crates/vm/src/stdlib/_symtable.rs @@ -174,6 +174,14 @@ mod _symtable { .symtable .sub_tables .iter() + .flat_map(|t| { + if t.comp_inlined { + // Flatten: replace inlined comprehension tables with their children + t.sub_tables.iter().collect::>() + } else { + vec![t] + } + }) .map(|t| to_py_symbol_table(t.clone()).into_pyobject(vm)) .collect(); Ok(children) diff --git a/crates/vm/src/stdlib/sys/monitoring.rs b/crates/vm/src/stdlib/sys/monitoring.rs index 739165073af..367038658c6 100644 --- a/crates/vm/src/stdlib/sys/monitoring.rs +++ b/crates/vm/src/stdlib/sys/monitoring.rs @@ -368,6 +368,9 @@ pub fn instrument_code(code: &PyCode, events: u32) { // is_line_start[i] = true if position i should have INSTRUMENTED_LINE let mut is_line_start = vec![false; len]; + // Build NO_LOCATION mask from linetable + let no_loc_mask = bytecode::build_no_location_mask(&code.code.linetable, len); + // First pass: mark positions where the source line changes let mut prev_line: Option = None; for (i, unit) in code @@ -395,6 +398,10 @@ pub fn instrument_code(code: &PyCode, events: u32) { ) { continue; } + // Skip NO_LOCATION instructions + if no_loc_mask.get(i).copied().unwrap_or(false) { + continue; + } if let Some((loc, _)) = code.code.locations.get(i) { let line = loc.line.get() as u32; let is_new = prev_line != Some(line); @@ -445,6 +452,7 @@ pub fn instrument_code(code: &PyCode, events: u32) { if let Some(target_idx) = target && target_idx < len && !is_line_start[target_idx] + && !no_loc_mask.get(target_idx).copied().unwrap_or(false) { let target_op = code.code.instructions[target_idx].op; let target_base = target_op.to_base().map_or(target_op, |b| b); @@ -465,7 +473,10 @@ pub fn instrument_code(code: &PyCode, events: u32) { // Third pass: mark exception handler targets as line starts. for entry in bytecode::decode_exception_table(&code.code.exceptiontable) { let target_idx = entry.target as usize; - if target_idx < len && !is_line_start[target_idx] { + if target_idx < len + && !is_line_start[target_idx] + && !no_loc_mask.get(target_idx).copied().unwrap_or(false) + { 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)