diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index 27d1394e516..08830c25ae7 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -529,7 +529,6 @@ def f(x): self.assertEqual(len(returns), 1) self.check_lnotab(f) - @unittest.expectedFailure # TODO: RUSTPYTHON; KeyError: 20 def test_elim_jump_to_return(self): # JUMP_FORWARD to RETURN --> RETURN def f(cond, true_value, false_value): diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index e3f1d187370..e0debb48cd2 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -2481,15 +2481,10 @@ impl Compiler { idx: bytecode::CommonConstant::AssertionError } ); - emit!(self, Instruction::PushNull); - match msg { - Some(e) => { - self.compile_expression(e)?; - emit!(self, Instruction::Call { argc: 1 }); - } - None => { - emit!(self, Instruction::Call { argc: 0 }); - } + if let Some(e) = msg { + emit!(self, Instruction::PushNull); + self.compile_expression(e)?; + emit!(self, Instruction::Call { argc: 1 }); } emit!( self, @@ -3221,20 +3216,19 @@ impl Compiler { emit!(self, PseudoInstruction::PopBlock); } - // Create a block for normal path continuation (after handler body succeeds) - let handler_normal_exit = self.new_block(); - emit!( - self, - PseudoInstruction::JumpNoInterrupt { - delta: handler_normal_exit, - } - ); - // cleanup_end block for named handler // IMPORTANT: In CPython, cleanup_end is within outer SETUP_CLEANUP scope. // so when RERAISE is executed, it goes to the cleanup block which does POP_EXCEPT. // We MUST compile cleanup_end BEFORE popping ExceptionHandler so RERAISE routes to cleanup_block. if let Some(cleanup_end) = handler_cleanup_block { + let handler_normal_exit = self.new_block(); + emit!( + self, + PseudoInstruction::JumpNoInterrupt { + delta: handler_normal_exit, + } + ); + self.switch_to_block(cleanup_end); if let Some(alias) = name { // name = None; del name; before RERAISE @@ -3247,10 +3241,10 @@ impl Compiler { // This RERAISE is within ExceptionHandler scope, so it routes to cleanup_block // which does COPY 3; POP_EXCEPT; RERAISE emit!(self, Instruction::Reraise { depth: 1 }); - } - // Switch to normal exit block - this is where handler body success continues - self.switch_to_block(handler_normal_exit); + // Switch to normal exit block - this is where handler body success continues + self.switch_to_block(handler_normal_exit); + } // PopBlock for outer SETUP_CLEANUP (ExceptionHandler) emit!(self, PseudoInstruction::PopBlock); @@ -6789,6 +6783,89 @@ impl Compiler { Ok(()) } + fn compile_jump_if_compare( + &mut self, + left: &ast::Expr, + ops: &[ast::CmpOp], + comparators: &[ast::Expr], + condition: bool, + target_block: BlockIdx, + ) -> CompileResult<()> { + let compare_range = self.current_source_range; + let (last_op, mid_ops) = ops.split_last().unwrap(); + let (last_comparator, mid_comparators) = comparators.split_last().unwrap(); + + self.compile_expression(left)?; + + if mid_comparators.is_empty() { + self.compile_expression(last_comparator)?; + self.set_source_range(compare_range); + self.compile_addcompare(last_op); + if condition { + emit!( + self, + Instruction::PopJumpIfTrue { + delta: target_block + } + ); + } else { + emit!( + self, + Instruction::PopJumpIfFalse { + delta: target_block, + } + ); + } + return Ok(()); + } + + let cleanup = self.new_block(); + let end = self.new_block(); + + for (op, comparator) in mid_ops.iter().zip(mid_comparators) { + self.compile_expression(comparator)?; + self.set_source_range(compare_range); + emit!(self, Instruction::Swap { i: 2 }); + emit!(self, Instruction::Copy { i: 2 }); + self.compile_addcompare(op); + emit!(self, Instruction::PopJumpIfFalse { delta: cleanup }); + } + + self.compile_expression(last_comparator)?; + self.set_source_range(compare_range); + self.compile_addcompare(last_op); + if condition { + emit!( + self, + Instruction::PopJumpIfTrue { + delta: target_block + } + ); + } else { + emit!( + self, + Instruction::PopJumpIfFalse { + delta: target_block, + } + ); + } + emit!(self, PseudoInstruction::Jump { delta: end }); + + self.switch_to_block(cleanup); + emit!(self, Instruction::PopTop); + if !condition { + emit!( + self, + PseudoInstruction::Jump { + delta: target_block + } + ); + } + + self.switch_to_block(end); + Ok(()) + } + fn compile_annotation(&mut self, annotation: &ast::Expr) -> CompileResult<()> { if self.future_annotations { self.emit_load_const(ConstantData::Str { @@ -7066,8 +7143,11 @@ impl Compiler { condition: bool, target_block: BlockIdx, ) -> CompileResult<()> { + let prev_source_range = self.current_source_range; + self.set_source_range(expression.range()); + // Compile expression for test, and jump to label if false - match &expression { + let result = match &expression { ast::Expr::BoolOp(ast::ExprBoolOp { op, values, .. }) => { match op { ast::BoolOp::And => { @@ -7113,13 +7193,20 @@ impl Compiler { } } } + Ok(()) } ast::Expr::UnaryOp(ast::ExprUnaryOp { op: ast::UnaryOp::Not, operand, .. - }) => { - self.compile_jump_if(operand, !condition, target_block)?; + }) => self.compile_jump_if(operand, !condition, target_block), + ast::Expr::Compare(ast::ExprCompare { + left, + ops, + comparators, + .. + }) if ops.len() > 1 => { + self.compile_jump_if_compare(left, ops, comparators, condition, target_block) } // `x is None` / `x is not None` → POP_JUMP_IF_NONE / POP_JUMP_IF_NOT_NONE ast::Expr::Compare(ast::ExprCompare { @@ -7154,6 +7241,7 @@ impl Compiler { } ); } + Ok(()) } _ => { // Fall back case which always will work! @@ -7177,9 +7265,12 @@ impl Compiler { } ); } + Ok(()) } - } - Ok(()) + }; + + self.set_source_range(prev_source_range); + result } /// Compile a boolean operation as an expression. @@ -7575,8 +7666,11 @@ impl Compiler { lower, upper, step, .. }) => { // Try constant slice folding first - if self.try_fold_constant_slice(lower.as_deref(), upper.as_deref(), step.as_deref()) - { + if self.try_fold_constant_slice( + lower.as_deref(), + upper.as_deref(), + step.as_deref(), + )? { return Ok(()); } let mut compile_bound = |bound: Option<&ast::Expr>| match bound { @@ -7991,12 +8085,15 @@ impl Compiler { // Save the call expression's source range so CALL instructions use the // call start line, not the last argument's line. let call_range = self.current_source_range; + let uses_ex_call = self.call_uses_ex_call(args); // Method call: obj → LOAD_ATTR_METHOD → [method, self_or_null] → args → CALL // Regular call: func → PUSH_NULL → args → CALL if let ast::Expr::Attribute(ast::ExprAttribute { value, attr, .. }) = &func { // Check for super() method call optimization - if let Some(super_type) = self.can_optimize_super_call(value, attr.as_str()) { + if !uses_ex_call + && let Some(super_type) = self.can_optimize_super_call(value, attr.as_str()) + { // super().method() or super(cls, self).method() optimization // Stack: [global_super, class, self] → LOAD_SUPER_METHOD → [method, self] // Set source range to the super() call for LOAD_GLOBAL/LOAD_DEREF/etc. @@ -8021,12 +8118,12 @@ impl Compiler { } else { self.compile_expression(value)?; let idx = self.name(attr.as_str()); - // Imported names use plain LOAD_ATTR + PUSH_NULL; - // other names use method call mode LOAD_ATTR. + // Imported names and CALL_FUNCTION_EX-style calls use plain + // LOAD_ATTR + PUSH_NULL; other names use method-call mode. // Check current scope and enclosing scopes for IMPORTED flag. let is_import = matches!(value.as_ref(), ast::Expr::Name(ast::ExprName { id, .. }) if self.is_name_imported(id.as_str())); - if is_import { + if is_import || uses_ex_call { self.emit_load_attr(idx); emit!(self, Instruction::PushNull); } else { @@ -8044,6 +8141,16 @@ impl Compiler { Ok(()) } + fn call_uses_ex_call(&self, arguments: &ast::Arguments) -> bool { + let has_starred = arguments + .args + .iter() + .any(|arg| matches!(arg, ast::Expr::Starred(_))); + let has_double_star = arguments.keywords.iter().any(|k| k.arg.is_none()); + let too_big = arguments.args.len() + arguments.keywords.len() > 15; + has_starred || has_double_star || too_big + } + /// Compile subkwargs: emit key-value pairs for BUILD_MAP fn codegen_subkwargs( &mut self, @@ -8929,6 +9036,28 @@ impl Compiler { } } + fn compile_fstring_literal_value( + &self, + string: &ast::InterpolatedStringLiteralElement, + flags: ast::FStringFlags, + ) -> Wtf8Buf { + if string.value.contains(char::REPLACEMENT_CHARACTER) { + let source = self.source_file.slice(string.range); + crate::string_parser::parse_fstring_literal_element(source.into(), flags.into()).into() + } else { + string.value.to_string().into() + } + } + + fn compile_fstring_part_literal_value(&self, string: &ast::StringLiteral) -> Wtf8Buf { + if string.value.contains(char::REPLACEMENT_CHARACTER) { + let source = self.source_file.slice(string.range); + crate::string_parser::parse_string_literal(source, string.flags.into()).into() + } else { + string.value.to_string().into() + } + } + fn arg_constant(&mut self, constant: ConstantData) -> oparg::ConstIdx { let info = self.current_code_info(); info.metadata.consts.insert_full(constant).0.to_u32().into() @@ -8942,48 +9071,40 @@ impl Compiler { ) -> CompileResult> { let mut constants = Vec::with_capacity(elts.len()); for elt in elts { - match elt { - ast::Expr::NumberLiteral(num) => match &num.value { - ast::Number::Int(int) => { - let value = ruff_int_to_bigint(int).map_err(|e| self.error(e))?; - constants.push(ConstantData::Integer { value }); - } - ast::Number::Float(f) => { - constants.push(ConstantData::Float { value: *f }); - } - ast::Number::Complex { real, imag } => { - constants.push(ConstantData::Complex { - value: Complex::new(*real, *imag), - }); - } - }, - ast::Expr::StringLiteral(s) => { - constants.push(ConstantData::Str { - value: self.compile_string_value(s), - }); - } - ast::Expr::BytesLiteral(b) => { - constants.push(ConstantData::Bytes { - value: b.value.bytes().collect(), - }); - } - ast::Expr::BooleanLiteral(b) => { - constants.push(ConstantData::Boolean { value: b.value }); - } - ast::Expr::NoneLiteral(_) => { - constants.push(ConstantData::None); - } - ast::Expr::EllipsisLiteral(_) => { - constants.push(ConstantData::Ellipsis); - } - _ => return Ok(None), - } + let Some(constant) = self.try_fold_constant_expr(elt)? else { + return Ok(None); + }; + constants.push(constant); } Ok(Some(ConstantData::Tuple { elements: constants, })) } + fn try_fold_constant_expr(&mut self, expr: &ast::Expr) -> CompileResult> { + Ok(Some(match expr { + ast::Expr::NumberLiteral(num) => match &num.value { + ast::Number::Int(int) => ConstantData::Integer { + value: ruff_int_to_bigint(int).map_err(|e| self.error(e))?, + }, + ast::Number::Float(f) => ConstantData::Float { value: *f }, + ast::Number::Complex { real, imag } => ConstantData::Complex { + value: Complex::new(*real, *imag), + }, + }, + ast::Expr::StringLiteral(s) => ConstantData::Str { + value: self.compile_string_value(s), + }, + ast::Expr::BytesLiteral(b) => ConstantData::Bytes { + value: b.value.bytes().collect(), + }, + ast::Expr::BooleanLiteral(b) => ConstantData::Boolean { value: b.value }, + ast::Expr::NoneLiteral(_) => ConstantData::None, + ast::Expr::EllipsisLiteral(_) => ConstantData::Ellipsis, + _ => return Ok(None), + })) + } + fn emit_load_const(&mut self, constant: ConstantData) { let idx = self.arg_constant(constant); self.emit_arg(idx, |consti| Instruction::LoadConst { consti }) @@ -8995,38 +9116,26 @@ impl Compiler { lower: Option<&ast::Expr>, upper: Option<&ast::Expr>, step: Option<&ast::Expr>, - ) -> bool { - fn to_const(expr: Option<&ast::Expr>) -> Option { + ) -> CompileResult { + let to_const = |expr: Option<&ast::Expr>, this: &mut Self| -> CompileResult<_> { match expr { - None | Some(ast::Expr::NoneLiteral(_)) => Some(ConstantData::None), - Some(ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. })) => match value - { - ast::Number::Int(i) => Some(ConstantData::Integer { - value: i.as_i64()?.into(), - }), - ast::Number::Float(f) => Some(ConstantData::Float { value: *f }), - ast::Number::Complex { real, imag } => Some(ConstantData::Complex { - value: num_complex::Complex64::new(*real, *imag), - }), - }, - Some(ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. })) => { - Some(ConstantData::Boolean { value: *value }) - } - // Only match Constant_kind nodes (no UnaryOp) - _ => None, + None => Ok(Some(ConstantData::None)), + Some(expr) => this.try_fold_constant_expr(expr), } - } + }; - let (Some(start), Some(stop), Some(step_val)) = - (to_const(lower), to_const(upper), to_const(step)) - else { - return false; + let (Some(start), Some(stop), Some(step_val)) = ( + to_const(lower, self)?, + to_const(upper, self)?, + to_const(step, self)?, + ) else { + return Ok(false); }; self.emit_load_const(ConstantData::Slice { elements: Box::new([start, stop, step_val]), }); - true + Ok(true) } fn emit_return_const(&mut self, constant: ConstantData) { @@ -9493,45 +9602,63 @@ impl Compiler { fn compile_expr_fstring(&mut self, fstring: &ast::ExprFString) -> CompileResult<()> { let fstring = &fstring.value; + let mut element_count = 0; + let mut pending_literal = None; for part in fstring { - self.compile_fstring_part(part)?; + self.compile_fstring_part_into(part, &mut pending_literal, &mut element_count)?; } - let part_count: u32 = fstring - .iter() - .len() - .try_into() - .expect("BuildString size overflowed"); - if part_count > 1 { - emit!(self, Instruction::BuildString { count: part_count }); - } - - Ok(()) + self.finish_fstring(pending_literal, element_count) } - fn compile_fstring_part(&mut self, part: &ast::FStringPart) -> CompileResult<()> { + fn compile_fstring_part_into( + &mut self, + part: &ast::FStringPart, + pending_literal: &mut Option, + element_count: &mut u32, + ) -> CompileResult<()> { match part { ast::FStringPart::Literal(string) => { - if string.value.contains(char::REPLACEMENT_CHARACTER) { - // might have a surrogate literal; should reparse to be sure - let source = self.source_file.slice(string.range); - let value = - crate::string_parser::parse_string_literal(source, string.flags.into()); - self.emit_load_const(ConstantData::Str { - value: value.into(), - }); + let value = self.compile_fstring_part_literal_value(string); + if let Some(pending) = pending_literal.as_mut() { + pending.push_wtf8(value.as_ref()); } else { - self.emit_load_const(ConstantData::Str { - value: string.value.to_string().into(), - }); + *pending_literal = Some(value); } Ok(()) } - ast::FStringPart::FString(fstring) => self.compile_fstring(fstring), + ast::FStringPart::FString(fstring) => self.compile_fstring_elements_into( + fstring.flags, + &fstring.elements, + pending_literal, + element_count, + ), } } - fn compile_fstring(&mut self, fstring: &ast::FString) -> CompileResult<()> { - self.compile_fstring_elements(fstring.flags, &fstring.elements) + fn finish_fstring( + &mut self, + mut pending_literal: Option, + mut element_count: u32, + ) -> CompileResult<()> { + if let Some(value) = pending_literal.take() { + self.emit_load_const(ConstantData::Str { value }); + element_count += 1; + } + + if element_count == 0 { + self.emit_load_const(ConstantData::Str { + value: Wtf8Buf::new(), + }); + } else if element_count > 1 { + emit!( + self, + Instruction::BuildString { + count: element_count + } + ); + } + + Ok(()) } /// Optimize `'format_str' % (args,)` into f-string bytecode. @@ -9659,24 +9786,31 @@ impl Compiler { fstring_elements: &ast::InterpolatedStringElements, ) -> CompileResult<()> { let mut element_count = 0; + let mut pending_literal: Option = None; + self.compile_fstring_elements_into( + flags, + fstring_elements, + &mut pending_literal, + &mut element_count, + )?; + self.finish_fstring(pending_literal, element_count) + } + + fn compile_fstring_elements_into( + &mut self, + flags: ast::FStringFlags, + fstring_elements: &ast::InterpolatedStringElements, + pending_literal: &mut Option, + element_count: &mut u32, + ) -> CompileResult<()> { for element in fstring_elements { - element_count += 1; match element { ast::InterpolatedStringElement::Literal(string) => { - if string.value.contains(char::REPLACEMENT_CHARACTER) { - // might have a surrogate literal; should reparse to be sure - let source = self.source_file.slice(string.range); - let value = crate::string_parser::parse_fstring_literal_element( - source.into(), - flags.into(), - ); - self.emit_load_const(ConstantData::Str { - value: value.into(), - }); + let value = self.compile_fstring_literal_value(string, flags); + if let Some(pending) = pending_literal.as_mut() { + pending.push_wtf8(value.as_ref()); } else { - self.emit_load_const(ConstantData::Str { - value: string.value.to_string().into(), - }); + *pending_literal = Some(value); } } ast::InterpolatedStringElement::Interpolation(fstring_expr) => { @@ -9697,8 +9831,10 @@ impl Compiler { ] .concat(); - self.emit_load_const(ConstantData::Str { value: text.into() }); - element_count += 1; + let text: Wtf8Buf = text.into(); + pending_literal + .get_or_insert_with(Wtf8Buf::new) + .push_wtf8(text.as_ref()); // If debug text is present, apply repr conversion when no `format_spec` specified. // See action_helpers.c: fstring_find_expr_replacement @@ -9710,6 +9846,11 @@ impl Compiler { } } + if let Some(value) = pending_literal.take() { + self.emit_load_const(ConstantData::Str { value }); + *element_count += 1; + } + self.compile_expression(&fstring_expr.expression)?; match conversion { @@ -9731,22 +9872,10 @@ impl Compiler { emit!(self, Instruction::FormatSimple); } } - } - } - } - if element_count == 0 { - // ensure to put an empty string on the stack if there aren't any fstring elements - self.emit_load_const(ConstantData::Str { - value: Wtf8Buf::new(), - }); - } else if element_count > 1 { - emit!( - self, - Instruction::BuildString { - count: element_count + *element_count += 1; } - ); + } } Ok(()) @@ -10195,7 +10324,28 @@ mod ruff_tests { #[cfg(test)] mod tests { use super::*; - use rustpython_compiler_core::SourceFileBuilder; + use rustpython_compiler_core::{SourceFileBuilder, bytecode::OpArg}; + + fn assert_scope_exit_locations(code: &CodeObject) { + for (instr, (location, _)) in code.instructions.iter().zip(code.locations.iter()) { + if matches!( + instr.op, + Instruction::ReturnValue + | Instruction::RaiseVarargs { .. } + | Instruction::Reraise { .. } + ) { + assert!( + location.line.get() > 0, + "scope-exit instruction {instr:?} is missing a line number" + ); + } + } + for constant in code.constants.iter() { + if let ConstantData::Code { code } = constant { + assert_scope_exit_locations(code); + } + } + } fn compile_exec(source: &str) -> CodeObject { let opts = CompileOpts::default(); @@ -10230,6 +10380,19 @@ mod tests { compiler.exit_scope() } + fn find_code<'a>(code: &'a CodeObject, name: &str) -> Option<&'a CodeObject> { + if code.obj_name == name { + return Some(code); + } + code.constants.iter().find_map(|constant| { + if let ConstantData::Code { code } = constant { + find_code(code, name) + } else { + None + } + }) + } + macro_rules! assert_dis_snapshot { ($value:expr) => { insta::assert_snapshot!( @@ -10306,6 +10469,245 @@ async def test(): )); } + #[test] + fn test_scope_exit_instructions_keep_line_numbers() { + let code = compile_exec( + "\ +async def test(): + for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')): + with self.subTest(type=type(stop_exc)): + try: + async with egg(): + raise stop_exc + except Exception as ex: + self.assertIs(ex, stop_exc) + else: + self.fail(f'{stop_exc} was suppressed') +", + ); + assert_scope_exit_locations(&code); + } + + #[test] + fn test_attribute_ex_call_uses_plain_load_attr() { + let code = compile_exec( + "\ +def f(cls, args, kwargs): + cls.__new__(cls, *args) + cls.__new__(cls, *args, **kwargs) +", + ); + let f = find_code(&code, "f").expect("missing function code"); + + let ex_call_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::CallFunctionEx)) + .count(); + let load_attr_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::LoadAttr { .. })) + .count(); + + assert_eq!(ex_call_count, 2); + assert_eq!(load_attr_count, 2); + + for unit in f.instructions.iter() { + if let Instruction::LoadAttr { namei } = unit.op { + let load_attr = namei.get(OpArg::new(u32::from(u8::from(unit.arg)))); + assert!( + !load_attr.is_method(), + "CALL_FUNCTION_EX should use plain LOAD_ATTR" + ); + } + } + } + + #[test] + fn test_simple_attribute_call_keeps_method_load() { + let code = compile_exec( + "\ +def f(obj, arg): + return obj.method(arg) +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let load_attr = f + .instructions + .iter() + .find_map(|unit| match unit.op { + Instruction::LoadAttr { namei } => { + Some(namei.get(OpArg::new(u32::from(u8::from(unit.arg))))) + } + _ => None, + }) + .expect("missing LOAD_ATTR"); + + assert!( + load_attr.is_method(), + "simple method calls should stay optimized" + ); + } + + #[test] + fn test_conditional_return_epilogue_is_duplicated() { + let code = compile_exec( + "\ +def f(base, cls, state): + if base is object: + obj = object.__new__(cls) + else: + obj = base.__new__(cls, state) + return obj +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let return_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::ReturnValue)) + .count(); + + assert_eq!(return_count, 2); + } + + #[test] + fn test_assert_without_message_raises_class_directly() { + let code = compile_exec( + "\ +def f(x): + assert x +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let call_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::Call { .. })) + .count(); + let push_null_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::PushNull)) + .count(); + + assert_eq!(call_count, 0); + assert_eq!(push_null_count, 0); + } + + #[test] + fn test_chained_compare_jump_uses_single_cleanup_copy() { + let code = compile_exec( + "\ +def f(code): + if not 1 <= code <= 2147483647: + raise ValueError('x') +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let copy_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::Copy { .. })) + .count(); + let pop_top_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::PopTop)) + .count(); + + assert_eq!(copy_count, 1); + assert_eq!(pop_top_count, 1); + } + + #[test] + fn test_constant_slice_folding_handles_string_and_bigint_bounds() { + let code = compile_exec( + "\ +def f(obj): + return obj['a':123456789012345678901234567890] +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let slice = f + .constants + .iter() + .find_map(|constant| match constant { + ConstantData::Slice { elements } => Some(elements), + _ => None, + }) + .expect("missing folded slice constant"); + + assert!(matches!(slice[0], ConstantData::Str { .. })); + assert!(matches!(slice[1], ConstantData::Integer { .. })); + assert!(matches!(slice[2], ConstantData::None)); + } + + #[test] + fn test_exception_cleanup_jump_to_return_is_inlined() { + let code = compile_exec( + "\ +def f(names, cls): + try: + cls.attr = names + except: + pass + return names +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let return_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::ReturnValue)) + .count(); + + assert_eq!(return_count, 2); + } + + #[test] + fn test_fstring_adjacent_literals_are_merged() { + let code = compile_exec( + "\ +def f(cls, proto): + raise TypeError( + f\"cannot pickle {cls.__name__!r} object: \" + f\"a class that defines __slots__ without \" + f\"defining __getstate__ cannot be pickled \" + f\"with protocol {proto}\" + ) +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let string_consts = f + .instructions + .iter() + .filter_map(|unit| match unit.op { + Instruction::LoadConst { consti } => { + Some(&f.constants[consti.get(OpArg::new(u32::from(u8::from(unit.arg))))]) + } + _ => None, + }) + .filter_map(|constant| match constant { + ConstantData::Str { value } => Some(value.to_string()), + _ => None, + }) + .collect::>(); + + assert!( + string_consts.iter().any(|value| { + value + == " object: a class that defines __slots__ without defining __getstate__ cannot be pickled with protocol " + }), + "expected merged trailing f-string literal, got {string_consts:?}" + ); + assert!( + !string_consts.iter().any(|value| value == " object: "), + "did not expect split trailing literal, got {string_consts:?}" + ); + } + #[test] fn test_optimized_assert_preserves_nested_scope_order() { compile_exec_optimized( diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 63e31a5b1f9..12c1961e88f 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -223,10 +223,10 @@ impl CodeInfo { // Phase 2: _PyCfg_OptimizedCfgToInstructionSequence (flowgraph.c) normalize_jumps(&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(); - // TODO: duplicate_exits_without_lineno disabled pending stackdepth fix - // duplicate_exits_without_lineno(&mut self.blocks); + resolve_line_numbers(&mut self.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 @@ -1630,21 +1630,9 @@ impl CodeInfo { let mut maxdepth = 0u32; let mut stack = Vec::with_capacity(self.blocks.len()); let mut start_depths = vec![u32::MAX; self.blocks.len()]; - start_depths[0] = 0; - stack.push(BlockIdx(0)); + stackdepth_push(&mut stack, &mut start_depths, BlockIdx(0), 0); const DEBUG: bool = false; - // Global iteration limit as safety guard - // The algorithm is monotonic (depths only increase), so it should converge quickly. - // Max iterations = blocks * max_possible_depth_increases per block - let max_iterations = self.blocks.len() * 100; - let mut iterations = 0usize; 'process_blocks: while let Some(block_idx) = stack.pop() { - iterations += 1; - if iterations > max_iterations { - // Safety guard: should never happen in valid code - // Return error instead of silently breaking to avoid underestimated stack depth - return Err(InternalError::StackOverflow); - } let idx = block_idx.idx(); let mut depth = start_depths[idx]; if DEBUG { @@ -1727,6 +1715,10 @@ impl CodeInfo { eprintln!("DONE: {maxdepth}"); } + for (block, &start_depth) in self.blocks.iter_mut().zip(&start_depths) { + block.start_depth = (start_depth != u32::MAX).then_some(start_depth); + } + // Fix up handler stack_depth in ExceptHandlerInfo using start_depths // computed above: depth = start_depth - 1 - preserve_lasti for block in self.blocks.iter_mut() { @@ -1789,7 +1781,6 @@ fn stackdepth_push( let idx = target.idx(); let block_depth = &mut start_depths[idx]; if depth > *block_depth || *block_depth == u32::MAX { - // Found a path with higher depth (or first visit): update max and queue *block_depth = depth; stack.push(target); } @@ -2410,7 +2401,7 @@ fn normalize_jumps(blocks: &mut Vec) { current = blocks[current.idx()].next; } - // Replace JUMP → LOAD_CONST + RETURN_VALUE with inline return. + // Replace JUMP → value-producing-instr + 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(); @@ -2427,11 +2418,13 @@ fn normalize_jumps(blocks: &mut Vec) { target_idx = blocks[target_idx].next.idx(); } let target_block = &blocks[target_idx]; - // Target must start with LOAD_CONST + RETURN_VALUE - if target_block.instructions.len() >= 2 { + // Target must be exactly `value; RETURN_VALUE`. + 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 { .. })) + if matches!(t0.instr, AnyInstruction::Real(_)) + && !t0.instr.is_scope_exit() + && !t0.instr.is_unconditional_jump() && matches!(t1.instr.real(), Some(Instruction::ReturnValue)) { let mut load = *t0; @@ -2500,13 +2493,61 @@ fn normalize_jumps(blocks: &mut Vec) { } } -#[allow(dead_code)] -fn is_simple_return_block(block: &Block) -> bool { - block.instructions.len() == 1 - && matches!( - block.instructions[0].instr.real(), - Some(Instruction::ReturnValue) - ) +/// flowgraph.c inline_small_or_no_lineno_blocks +fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { + const MAX_COPY_SIZE: usize = 4; + + let block_exits_scope = |block: &Block| { + block + .instructions + .last() + .is_some_and(|ins| ins.instr.is_scope_exit()) + }; + let block_has_no_lineno = |block: &Block| { + block + .instructions + .iter() + .all(|ins| !instruction_has_lineno(ins)) + }; + + loop { + let mut changes = false; + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + let next = blocks[current.idx()].next; + let Some(last) = blocks[current.idx()].instructions.last().copied() else { + current = next; + continue; + }; + if !last.instr.is_unconditional_jump() || last.target == BlockIdx::NULL { + current = next; + continue; + } + + let target = last.target; + 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()]) + && !block_has_fallthrough(&blocks[target.idx()]); + + if small_exit_block || no_lineno_no_fallthrough { + if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() { + last_instr.instr = Instruction::Nop.into(); + last_instr.arg = OpArg::new(0); + last_instr.target = BlockIdx::NULL; + } + let appended = blocks[target.idx()].instructions.clone(); + blocks[current.idx()].instructions.extend(appended); + changes = true; + } + + current = next; + } + + if !changes { + break; + } + } } /// Follow chain of empty blocks to find first non-empty block. @@ -2520,25 +2561,79 @@ fn next_nonempty_block(blocks: &[Block], mut idx: BlockIdx) -> BlockIdx { idx } -#[allow(dead_code)] -fn duplicate_exits_without_lineno(blocks: &mut Vec) { - // Count predecessors for each block +fn instruction_lineno(instr: &InstructionInfo) -> i32 { + instr + .lineno_override + .unwrap_or_else(|| instr.location.line.get() as i32) +} + +fn instruction_has_lineno(instr: &InstructionInfo) -> bool { + instruction_lineno(instr) > 0 +} + +fn block_has_fallthrough(block: &Block) -> bool { + block + .instructions + .last() + .is_none_or(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump()) +} + +fn is_jump_instruction(instr: &InstructionInfo) -> bool { + instr.instr.is_unconditional_jump() || is_conditional_jump(&instr.instr) +} + +fn is_exit_without_lineno(block: &Block) -> bool { + let Some(first) = block.instructions.first() else { + return false; + }; + let Some(last) = block.instructions.last() else { + return false; + }; + !instruction_has_lineno(first) && last.instr.is_scope_exit() +} + +fn maybe_propagate_location( + instr: &mut InstructionInfo, + location: SourceLocation, + end_location: SourceLocation, +) { + if !instruction_has_lineno(instr) { + instr.location = location; + instr.end_location = end_location; + instr.lineno_override = None; + } +} + +fn propagate_locations_in_block( + block: &mut Block, + location: SourceLocation, + end_location: SourceLocation, +) { + let mut prev_location = location; + let mut prev_end_location = end_location; + for instr in &mut block.instructions { + maybe_propagate_location(instr, prev_location, prev_end_location); + prev_location = instr.location; + prev_end_location = instr.end_location; + } +} + +fn compute_predecessors(blocks: &[Block]) -> Vec { let mut predecessors = vec![0u32; blocks.len()]; + if blocks.is_empty() { + return predecessors; + } + + predecessors[0] = 1; let mut current = BlockIdx(0); while current != BlockIdx::NULL { let block = &blocks[current.idx()]; - // Fall-through predecessor - let next = next_nonempty_block(blocks, block.next); - if next != BlockIdx::NULL { - let has_fallthrough = block - .instructions - .last() - .is_none_or(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump()); - if has_fallthrough { + if block_has_fallthrough(block) { + let next = next_nonempty_block(blocks, block.next); + if next != BlockIdx::NULL { predecessors[next.idx()] += 1; } } - // Jump target predecessor for ins in &block.instructions { if ins.target != BlockIdx::NULL { let target = next_nonempty_block(blocks, ins.target); @@ -2549,26 +2644,23 @@ fn duplicate_exits_without_lineno(blocks: &mut Vec) { } current = block.next; } + predecessors +} - // For each block, if its last instruction jumps to an exit block without lineno - // that has >1 predecessor, copy that exit block. - current = BlockIdx(0); +fn duplicate_exits_without_lineno(blocks: &mut Vec, predecessors: &mut Vec) { + let mut current = BlockIdx(0); while current != BlockIdx::NULL { let block = &blocks[current.idx()]; let last = match block.instructions.last() { - Some(ins) if ins.target != BlockIdx::NULL => ins, + Some(ins) if ins.target != BlockIdx::NULL && is_jump_instruction(ins) => ins, _ => { current = blocks[current.idx()].next; continue; } }; - if !last.instr.is_unconditional_jump() && !is_conditional_jump(&last.instr) { - current = blocks[current.idx()].next; - continue; - } let target = next_nonempty_block(blocks, last.target); - if target == BlockIdx::NULL || !is_simple_return_block(&blocks[target.idx()]) { + if target == BlockIdx::NULL || !is_exit_without_lineno(&blocks[target.idx()]) { current = blocks[current.idx()].next; continue; } @@ -2580,17 +2672,9 @@ fn duplicate_exits_without_lineno(blocks: &mut Vec) { // Copy the exit block let new_idx = BlockIdx(blocks.len() as u32); let mut new_block = blocks[target.idx()].clone(); - // Set location from the jump instruction - let jump_loc = blocks[current.idx()].instructions.last().unwrap().location; - let jump_end_loc = blocks[current.idx()] - .instructions - .last() - .unwrap() - .end_location; - if let Some(first) = new_block.instructions.first_mut() { - first.location = jump_loc; - first.end_location = jump_end_loc; - } + let jump_loc = last.location; + let jump_end_loc = last.end_location; + propagate_locations_in_block(&mut new_block, jump_loc, jump_end_loc); new_block.next = blocks[target.idx()].next; blocks.push(new_block); @@ -2598,26 +2682,106 @@ fn duplicate_exits_without_lineno(blocks: &mut Vec) { let last_mut = blocks[current.idx()].instructions.last_mut().unwrap(); last_mut.target = new_idx; predecessors[target.idx()] -= 1; + predecessors.push(1); + + current = blocks[current.idx()].next; + } + current = BlockIdx(0); + while current != BlockIdx::NULL { + let block = &blocks[current.idx()]; + if let Some(last) = block.instructions.last() + && block_has_fallthrough(block) + { + let target = next_nonempty_block(blocks, block.next); + if target != BlockIdx::NULL + && predecessors[target.idx()] == 1 + && is_exit_without_lineno(&blocks[target.idx()]) + { + let last_location = last.location; + let last_end_location = last.end_location; + propagate_locations_in_block( + &mut blocks[target.idx()], + last_location, + last_end_location, + ); + } + } current = blocks[current.idx()].next; } } +fn propagate_line_numbers(blocks: &mut [Block], predecessors: &[u32]) { + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + let last = blocks[current.idx()].instructions.last().copied(); + if let Some(last) = last { + let (next_block, has_fallthrough) = { + let block = &blocks[current.idx()]; + (block.next, block_has_fallthrough(block)) + }; + + { + let block = &mut blocks[current.idx()]; + let mut prev_location = None; + for instr in &mut block.instructions { + if let Some((location, end_location)) = prev_location { + maybe_propagate_location(instr, location, end_location); + } + prev_location = Some((instr.location, instr.end_location)); + } + } + + if has_fallthrough { + let target = next_nonempty_block(blocks, next_block); + if target != BlockIdx::NULL && predecessors[target.idx()] == 1 { + propagate_locations_in_block( + &mut blocks[target.idx()], + last.location, + last.end_location, + ); + } + } + + if is_jump_instruction(&last) { + let target = next_nonempty_block(blocks, last.target); + if target != BlockIdx::NULL && predecessors[target.idx()] == 1 { + propagate_locations_in_block( + &mut blocks[target.idx()], + last.location, + last.end_location, + ); + } + } + } + current = blocks[current.idx()].next; + } +} + +fn resolve_line_numbers(blocks: &mut Vec) { + let mut predecessors = compute_predecessors(blocks); + duplicate_exits_without_lineno(blocks, &mut predecessors); + propagate_line_numbers(blocks, &predecessors); +} + /// Duplicate `LOAD_CONST None + RETURN_VALUE` for blocks that fall through /// to the final return block. fn duplicate_end_returns(blocks: &mut [Block]) { - // Walk the block chain to find the last block - let mut last_block = BlockIdx(0); + // Walk the block chain and keep the last non-empty block. + let mut last_block = BlockIdx::NULL; let mut current = BlockIdx(0); while current != BlockIdx::NULL { - last_block = current; + if !blocks[current.idx()].instructions.is_empty() { + last_block = current; + } current = blocks[current.idx()].next; } + if last_block == BlockIdx::NULL { + return; + } - // Check if the last block ends with LOAD_CONST + RETURN_VALUE (the implicit return) let last_insts = &blocks[last_block.idx()].instructions; // Only apply when the last block is EXACTLY a return-None epilogue - // AND the return instructions have no explicit line number (lineno <= 0) let is_return_block = last_insts.len() == 2 && matches!( last_insts[0].instr, @@ -2639,8 +2803,8 @@ fn duplicate_end_returns(blocks: &mut [Block]) { current = BlockIdx(0); while current != BlockIdx::NULL { let block = &blocks[current.idx()]; - if current != last_block && block.next == last_block && !block.cold && !block.except_handler - { + let next = next_nonempty_block(blocks, block.next); + if current != last_block && next == 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()) diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index 629cdccc7fd..349b3b6f39c 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -42,6 +42,7 @@ impl Function { } } +#[allow(dead_code)] #[derive(Debug, Clone)] enum StackValue { String(String), @@ -49,6 +50,8 @@ enum StackValue { Map(HashMap), Code(Box), Function(Function), + Slice(Box<[StackValue; 3]>), + Frozenset(Vec), } impl From for StackValue { @@ -59,6 +62,13 @@ impl From for StackValue { } ConstantData::None => StackValue::None, ConstantData::Code { code } => StackValue::Code(code), + ConstantData::Slice { elements } => { + let [start, stop, step] = *elements; + StackValue::Slice(Box::new([start.into(), stop.into(), step.into()])) + } + ConstantData::Frozenset { elements } => { + StackValue::Frozenset(elements.into_iter().map(Into::into).collect()) + } c => unimplemented!("constant {:?} isn't yet supported in py_function!", c), } }