diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs index d3d412e9c52..98b70dfbb20 100644 --- a/compiler/codegen/src/compile.rs +++ b/compiler/codegen/src/compile.rs @@ -784,7 +784,12 @@ impl Compiler<'_> { if import_star { // from .... import * - emit!(self, Instruction::ImportStar); + emit!( + self, + Instruction::CallIntrinsic1 { + func: bytecode::IntrinsicFunction1::ImportStar + } + ); } else { // from mod import a, b as c @@ -1556,6 +1561,14 @@ impl Compiler<'_> { .constants .insert_full(ConstantData::None); + // Emit RESUME instruction at function start + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AtFuncStart as u32 + } + ); + self.compile_statements(body)?; // Emit None at end: @@ -1971,6 +1984,12 @@ impl Compiler<'_> { emit!(self, Instruction::GetAwaitable); self.emit_load_const(ConstantData::None); emit!(self, Instruction::YieldFrom); + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AfterAwait as u32 + } + ); emit!(self, Instruction::SetupAsyncWith { end: final_block }); } else { emit!(self, Instruction::SetupWith { end: final_block }); @@ -2012,6 +2031,12 @@ impl Compiler<'_> { emit!(self, Instruction::GetAwaitable); self.emit_load_const(ConstantData::None); emit!(self, Instruction::YieldFrom); + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AfterAwait as u32 + } + ); } emit!(self, Instruction::WithCleanupFinish); @@ -2050,6 +2075,12 @@ impl Compiler<'_> { emit!(self, Instruction::GetANext); self.emit_load_const(ConstantData::None); emit!(self, Instruction::YieldFrom); + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AfterAwait as u32 + } + ); self.compile_store(target)?; emit!(self, Instruction::PopBlock); } else { @@ -3521,6 +3552,12 @@ impl Compiler<'_> { Option::None => self.emit_load_const(ConstantData::None), }; emit!(self, Instruction::YieldValue); + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AfterYield as u32 + } + ); } Expr::Await(ExprAwait { value, .. }) => { if self.ctx.func != FunctionContext::AsyncFunction { @@ -3530,6 +3567,12 @@ impl Compiler<'_> { emit!(self, Instruction::GetAwaitable); self.emit_load_const(ConstantData::None); emit!(self, Instruction::YieldFrom); + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AfterAwait as u32 + } + ); } Expr::YieldFrom(ExprYieldFrom { value, .. }) => { match self.ctx.func { @@ -3546,6 +3589,12 @@ impl Compiler<'_> { emit!(self, Instruction::GetIter); self.emit_load_const(ConstantData::None); emit!(self, Instruction::YieldFrom); + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AfterYieldFrom as u32 + } + ); } Expr::Name(ExprName { id, .. }) => self.load_name(id.as_str())?, Expr::Lambda(ExprLambda { @@ -3672,6 +3721,12 @@ impl Compiler<'_> { compiler.compile_comprehension_element(elt)?; compiler.mark_generator(); emit!(compiler, Instruction::YieldValue); + emit!( + compiler, + Instruction::Resume { + arg: bytecode::ResumeType::AfterYield as u32 + } + ); emit!(compiler, Instruction::Pop); Ok(()) @@ -4067,6 +4122,12 @@ impl Compiler<'_> { emit!(self, Instruction::GetANext); self.emit_load_const(ConstantData::None); emit!(self, Instruction::YieldFrom); + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AfterAwait as u32 + } + ); self.compile_store(&generator.target)?; emit!(self, Instruction::PopBlock); } else { @@ -4145,6 +4206,12 @@ impl Compiler<'_> { emit!(self, Instruction::GetAwaitable); self.emit_load_const(ConstantData::None); emit!(self, Instruction::YieldFrom); + emit!( + self, + Instruction::Resume { + arg: bytecode::ResumeType::AfterAwait as u32 + } + ); } Ok(()) diff --git a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap index 36b00c567d7..9165a6cfbf1 100644 --- a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap +++ b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap @@ -11,7 +11,7 @@ expression: "compile_exec(\"\\\nfor stop_exc in (StopIteration('spam'), StopAsyn 6 CallFunctionPositional(1) 7 BuildTuple (2) 8 GetIter - >> 9 ForIter (71) + >> 9 ForIter (73) 10 StoreLocal (2, stop_exc) 2 11 LoadNameAny (3, self) @@ -21,10 +21,10 @@ expression: "compile_exec(\"\\\nfor stop_exc in (StopIteration('spam'), StopAsyn 15 CallFunctionPositional(1) 16 LoadConst (("type")) 17 CallMethodKeyword (1) - 18 SetupWith (68) + 18 SetupWith (70) 19 Pop - 3 20 SetupExcept (40) + 3 20 SetupExcept (42) 4 21 LoadNameAny (6, egg) 22 CallFunctionPositional(0) @@ -32,55 +32,57 @@ expression: "compile_exec(\"\\\nfor stop_exc in (StopIteration('spam'), StopAsyn 24 GetAwaitable 25 LoadConst (None) 26 YieldFrom - 27 SetupAsyncWith (33) - 28 Pop + 27 Resume (3) + 28 SetupAsyncWith (34) + 29 Pop - 5 29 LoadNameAny (2, stop_exc) - 30 Raise (Raise) + 5 30 LoadNameAny (2, stop_exc) + 31 Raise (Raise) - 4 31 PopBlock - 32 EnterFinally - >> 33 WithCleanupStart - 34 GetAwaitable - 35 LoadConst (None) - 36 YieldFrom - 37 WithCleanupFinish - 38 PopBlock - 39 Jump (57) - >> 40 Duplicate + 4 32 PopBlock + 33 EnterFinally + >> 34 WithCleanupStart + 35 GetAwaitable + 36 LoadConst (None) + 37 YieldFrom + 38 Resume (3) + 39 WithCleanupFinish + 40 PopBlock + 41 Jump (59) + >> 42 Duplicate - 6 41 LoadNameAny (7, Exception) - 42 TestOperation (ExceptionMatch) - 43 JumpIfFalse (56) - 44 StoreLocal (8, ex) + 6 43 LoadNameAny (7, Exception) + 44 TestOperation (ExceptionMatch) + 45 JumpIfFalse (58) + 46 StoreLocal (8, ex) - 7 45 LoadNameAny (3, self) - 46 LoadMethod (9, assertIs) - 47 LoadNameAny (8, ex) - 48 LoadNameAny (2, stop_exc) - 49 CallMethodPositional (2) - 50 Pop - 51 PopException - 52 LoadConst (None) - 53 StoreLocal (8, ex) - 54 DeleteLocal (8, ex) - 55 Jump (66) - >> 56 Raise (Reraise) + 7 47 LoadNameAny (3, self) + 48 LoadMethod (9, assertIs) + 49 LoadNameAny (8, ex) + 50 LoadNameAny (2, stop_exc) + 51 CallMethodPositional (2) + 52 Pop + 53 PopException + 54 LoadConst (None) + 55 StoreLocal (8, ex) + 56 DeleteLocal (8, ex) + 57 Jump (68) + >> 58 Raise (Reraise) - 9 >> 57 LoadNameAny (3, self) - 58 LoadMethod (10, fail) - 59 LoadConst ("") - 60 LoadNameAny (2, stop_exc) - 61 FormatValue (None) - 62 LoadConst (" was suppressed") - 63 BuildString (2) - 64 CallMethodPositional (1) - 65 Pop + 9 >> 59 LoadNameAny (3, self) + 60 LoadMethod (10, fail) + 61 LoadConst ("") + 62 LoadNameAny (2, stop_exc) + 63 FormatValue (None) + 64 LoadConst (" was suppressed") + 65 BuildString (2) + 66 CallMethodPositional (1) + 67 Pop - 2 >> 66 PopBlock - 67 EnterFinally - >> 68 WithCleanupStart - 69 WithCleanupFinish - 70 Jump (9) - >> 71 PopBlock - 72 ReturnConst (None) + 2 >> 68 PopBlock + 69 EnterFinally + >> 70 WithCleanupStart + 71 WithCleanupFinish + 72 Jump (9) + >> 73 PopBlock + 74 ReturnConst (None) diff --git a/compiler/core/src/bytecode.rs b/compiler/core/src/bytecode.rs index cef332bfbc6..3e74fe62738 100644 --- a/compiler/core/src/bytecode.rs +++ b/compiler/core/src/bytecode.rs @@ -24,6 +24,16 @@ pub enum ConversionFlag { Repr = b'r' as i8, } +/// Resume type for the RESUME instruction +#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)] +#[repr(u32)] +pub enum ResumeType { + AtFuncStart = 0, + AfterYield = 1, + AfterYieldFrom = 2, + AfterAwait = 3, +} + pub trait Constant: Sized { type Name: AsRef; @@ -382,6 +392,8 @@ op_arg_enum!( #[derive(Copy, Clone, Debug, PartialEq, Eq)] #[repr(u8)] pub enum IntrinsicFunction1 { + /// Import * operation + ImportStar = 2, /// Type parameter related TypeVar = 7, ParamSpec = 8, @@ -419,8 +431,6 @@ pub enum Instruction { }, /// Importing without name ImportNameless, - /// Import * - ImportStar, /// from ... import ... ImportFrom { idx: Arg, @@ -549,6 +559,12 @@ pub enum Instruction { }, YieldValue, YieldFrom, + + /// Resume execution (e.g., at function start, after yield, etc.) + Resume { + arg: Arg, + }, + SetupAnnotation, SetupLoop, @@ -1240,7 +1256,6 @@ impl Instruction { match self { Nop => 0, ImportName { .. } | ImportNameless => -1, - ImportStar => -1, ImportFrom { .. } => 1, LoadFast(_) | LoadNameAny(_) | LoadGlobal(_) | LoadDeref(_) | LoadClassDeref(_) => 1, StoreFast(_) | StoreLocal(_) | StoreGlobal(_) | StoreDeref(_) => -1, @@ -1305,6 +1320,7 @@ impl Instruction { } ReturnValue => -1, ReturnConst { .. } => 0, + Resume { .. } => 0, YieldValue => 0, YieldFrom => -1, SetupAnnotation | SetupLoop | SetupFinally { .. } | EnterFinally | EndFinally => 0, @@ -1433,7 +1449,6 @@ impl Instruction { Nop => w!(Nop), ImportName { idx } => w!(ImportName, name = idx), ImportNameless => w!(ImportNameless), - ImportStar => w!(ImportStar), ImportFrom { idx } => w!(ImportFrom, name = idx), LoadFast(idx) => w!(LoadFast, varname = idx), LoadNameAny(idx) => w!(LoadNameAny, name = idx), @@ -1493,6 +1508,7 @@ impl Instruction { ForIter { target } => w!(ForIter, target), ReturnValue => w!(ReturnValue), ReturnConst { idx } => fmt_const("ReturnConst", arg, f, idx), + Resume { arg } => w!(Resume, arg), YieldValue => w!(YieldValue), YieldFrom => w!(YieldFrom), SetupAnnotation => w!(SetupAnnotation), diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs index 9ec0a4385ee..5f0123d22ba 100644 --- a/jit/src/instructions.rs +++ b/jit/src/instructions.rs @@ -612,6 +612,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { self.stack.pop(); Ok(()) } + Instruction::Resume { arg: _resume_arg } => { + // TODO: Implement the resume instruction + Ok(()) + } _ => Err(JitCompileError::NotSupported), } } diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 7c935b814b3..a3e31c5c2bb 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -541,10 +541,6 @@ impl ExecutingFrame<'_> { self.import(vm, None)?; Ok(None) } - bytecode::Instruction::ImportStar => { - self.import_star(vm)?; - Ok(None) - } bytecode::Instruction::ImportFrom { idx } => { let obj = self.import_from(vm, idx.get(arg))?; self.push_value(obj); @@ -893,6 +889,18 @@ impl ExecutingFrame<'_> { Ok(Some(ExecutionResult::Yield(value))) } bytecode::Instruction::YieldFrom => self.execute_yield_from(vm), + bytecode::Instruction::Resume { arg: resume_arg } => { + // Resume execution after yield, await, or at function start + // In CPython, this checks instrumentation and eval breaker + // For now, we just check for signals/interrupts + let _resume_type = resume_arg.get(arg); + + // Check for interrupts if not resuming from yield_from + // if resume_type < bytecode::ResumeType::AfterYieldFrom as u32 { + // vm.check_signals()?; + // } + Ok(None) + } bytecode::Instruction::SetupAnnotation => self.setup_annotations(vm), bytecode::Instruction::SetupLoop => { self.push_block(BlockType::Loop); @@ -2203,6 +2211,12 @@ impl ExecutingFrame<'_> { vm: &VirtualMachine, ) -> PyResult { match func { + bytecode::IntrinsicFunction1::ImportStar => { + // arg is the module object + self.push_value(arg); // Push module back on stack for import_star + self.import_star(vm)?; + Ok(vm.ctx.none()) + } bytecode::IntrinsicFunction1::SubscriptGeneric => { // Used for PEP 695: Generic[*type_params] crate::builtins::genericalias::subscript_generic(arg, vm)