diff --git a/.claude/settings.json b/.claude/settings.json deleted file mode 100644 index 22f0a9a8a01..00000000000 --- a/.claude/settings.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "hooks": { - "SessionStart": [ - { - "matcher": "", - "hooks": [ - { - "type": "command", - "command": "bash .claude/scripts/setup-env.sh" - } - ] - } - ] - } -} diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 39195202af7..0fc4119b805 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -64,9 +64,9 @@ fn format_missing_args( #[pyclass(module = false, name = "function", traverse = "manual")] #[derive(Debug)] pub struct PyFunction { - code: PyAtomicRef, - globals: PyDictRef, - builtins: PyObjectRef, + pub(crate) code: PyAtomicRef, + pub(crate) globals: PyDictRef, + pub(crate) builtins: PyObjectRef, pub(crate) closure: Option>>, defaults_and_kwdefaults: PyMutex<(Option, Option)>, name: PyMutex, @@ -617,11 +617,6 @@ impl Py { // Fast path: stack-allocated InterpreterFrame, no FrameObject. // No refcount inc for code — it's alive via self.code for the call duration. - let nlocalsplus = code.localspluskinds.len(); - let max_stackdepth = code.max_stackdepth as usize; - let localsplus = - crate::frame::LocalsPlus::new_on_datastack(nlocalsplus, max_stackdepth, vm); - let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { crate::frame::FrameLocals::lazy() } else if let Some(locals) = locals { @@ -634,22 +629,21 @@ impl Py { // Use self.as_object() as raw pointer — no refcount inc/dec. // The function is alive on the caller's stack for the call duration. - let mut iframe = crate::frame::InterpreterFrame::new( + let iframe = crate::frame::InterpreterFrame::new_on_datastack( &self.code, &self.globals, &self.builtins, Some(self.as_object()), - localsplus, locals, self.closure.as_ref().map_or(&[], |c| c.as_slice()), - crate::frame::FrameOwner::Thread, + vm, ); let result = self - .fill_locals_from_args_iframe(&mut iframe, func_args, vm) - .and_then(|()| vm.run_frame_fast(&mut iframe)); + .fill_locals_from_args_iframe(iframe, func_args, vm) + .and_then(|()| vm.run_frame_fast(iframe)); // Release data stack memory — must happen on both success and error. unsafe { - if let Some(base) = iframe.localsplus.release_datastack() { + if let Some(base) = iframe.release_datastack_frame() { vm.datastack_pop(base); } } @@ -797,10 +791,6 @@ impl Py { vm: &VirtualMachine, ) -> PyResult { let code = &*self.code; - let nlocalsplus = code.localspluskinds.len(); - let max_stackdepth = code.max_stackdepth as usize; - let localsplus = - crate::frame::LocalsPlus::new_on_datastack(nlocalsplus, max_stackdepth, vm); let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { crate::frame::FrameLocals::lazy() @@ -810,15 +800,14 @@ impl Py { )) }; - let mut iframe = crate::frame::InterpreterFrame::new( + let iframe = crate::frame::InterpreterFrame::new_on_datastack( code, &self.globals, &self.builtins, Some(self.as_object()), - localsplus, locals, self.closure.as_ref().map_or(&[], |c| c.as_slice()), - crate::frame::FrameOwner::Thread, + vm, ); // Fill arguments directly into fastlocals @@ -829,9 +818,9 @@ impl Py { } } - let result = vm.run_frame_fast(&mut iframe); + let result = vm.run_frame_fast(iframe); unsafe { - if let Some(base) = iframe.localsplus.release_datastack() { + if let Some(base) = iframe.release_datastack_frame() { vm.datastack_pop(base); } } @@ -887,8 +876,10 @@ pub(crate) fn datastack_frame_size_bytes_for_code(code: &Py) -> Option()) + Some(crate::frame::datastack_iframe_total_bytes( + nlocalsplus, + code.max_stackdepth as usize, + )) } impl PyPayload for PyFunction { diff --git a/crates/vm/src/coroutine.rs b/crates/vm/src/coroutine.rs index c7252c66d12..43a28320e00 100644 --- a/crates/vm/src/coroutine.rs +++ b/crates/vm/src/coroutine.rs @@ -23,6 +23,7 @@ impl ExecutionResult { }; PyIterReturn::StopIteration(arg) } + Self::TailCall => unreachable!("TailCall in generator/coroutine"), } } } @@ -104,6 +105,7 @@ impl Coro { self.clear_frame_locals_on_close(); } Ok(ExecutionResult::Yield(_)) => {} + Ok(ExecutionResult::TailCall) => unreachable!("TailCall in generator/coroutine"), } } diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index bf276ec45e0..e16007ad0e1 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -544,6 +544,13 @@ impl LocalsPlus { Ok(()) } + /// Push a PyObjectRef onto the evaluation stack. + /// Panics on overflow. + pub(crate) fn push_stack(&mut self, value: PyObjectRef) { + self.stack_try_push(Some(PyStackRef::new_owned(value))) + .expect("stack overflow in push_stack"); + } + /// Pop a value from the evaluation stack. #[inline(always)] fn stack_pop(&mut self) -> Option { @@ -861,6 +868,9 @@ pub struct InterpreterFrame { /// Used by `frame.clear()` to reject clearing an executing frame, /// even when called from a different thread. pub(crate) owner: atomic::AtomicI8, + /// Base pointer of the datastack allocation when this frame and its + /// localsplus are bump-allocated together. Null for heap-backed frames. + pub(crate) datastack_base: *mut u8, /// Pointer to the owning `Py`, or null for stack-allocated /// frames that have not been materialized yet. /// Stored as `usize` for `PyAtomic` compatibility. @@ -944,11 +954,110 @@ impl InterpreterFrame { generator: PyAtomicBorrow::new(), previous: Radium::new(0), owner: atomic::AtomicI8::new(owner as i8), + datastack_base: core::ptr::null_mut(), materialized: Radium::new(0), cold: OnceCell::new(), } } + /// Allocate an InterpreterFrame and its LocalsPlus data together on the + /// thread data stack in a single bump allocation. + /// + /// Layout: `[InterpreterFrame | localsplus usize×capacity]` + /// + /// Returns a mutable reference whose lifetime is bounded by the data + /// stack's LIFO discipline. The caller must call + /// `release_datastack_frame()` (unsafe) when done, then + /// `vm.datastack_pop(base)`. The reference must not be used after + /// `release_datastack_frame` returns. + #[allow(clippy::too_many_arguments)] + pub(crate) fn new_on_datastack<'a>( + code: &Py, + globals: &Py, + builtins: &PyObject, + func_obj: Option<&PyObject>, + locals: FrameLocals, + closure: &[PyCellRef], + vm: &VirtualMachine, + ) -> &'a mut Self { + let nlocalsplus = code.localspluskinds.len(); + let stacksize = code.max_stackdepth as usize; + let capacity = nlocalsplus + .checked_add(stacksize) + .expect("LocalsPlus capacity overflow"); + + let total_bytes = datastack_iframe_total_bytes(nlocalsplus, stacksize); + let base = vm.datastack_push(total_bytes); + + // InterpreterFrame lives at the start of the allocation. + let iframe_ptr = base as *mut Self; + // LocalsPlus data follows the InterpreterFrame, aligned to usize. + let localsplus_data_ptr = + unsafe { base.add(datastack_iframe_localsplus_offset()) } as *mut usize; + + // Zero-initialize localsplus data. + unsafe { core::ptr::write_bytes(localsplus_data_ptr, 0, capacity) }; + + let nlocalsplus_u32 = u32::try_from(nlocalsplus).expect("nlocalsplus exceeds u32"); + let localsplus = LocalsPlus { + data: LocalsPlusData::DataStack { + ptr: localsplus_data_ptr, + capacity, + }, + nlocalsplus: nlocalsplus_u32, + stack_top: 0, + }; + + let mut iframe = Self::new( + code, + globals, + builtins, + func_obj, + localsplus, + locals, + closure, + FrameOwner::Thread, + ); + iframe.datastack_base = base; + + // Write the fully initialized InterpreterFrame into the datastack. + unsafe { + core::ptr::write(iframe_ptr, iframe); + &mut *iframe_ptr + } + } + + /// Release this datastack-allocated frame's resources and return the + /// base pointer for `vm.datastack_pop()`. + /// + /// Drops all localsplus values, runs destructors for all frame fields + /// (trace, temporary_refs, retained_back, etc.), and detaches the + /// backing store. + /// Returns `None` if this frame is not datastack-allocated. + /// + /// After this call, the InterpreterFrame at `self` is logically dead — + /// the caller must not use `self` again except to pass the returned + /// base to `vm.datastack_pop()`. + pub(crate) unsafe fn release_datastack_frame(&mut self) -> Option<*mut u8> { + let base = self.datastack_base; + if base.is_null() { + return None; + } + self.datastack_base = core::ptr::null_mut(); + // Drop all localsplus values while the backing store is still valid. + self.localsplus.drop_values(); + // Detach from the data stack so further accesses see an empty frame. + self.localsplus.data = LocalsPlusData::Heap(Box::default()); + self.localsplus.nlocalsplus = 0; + // Drop remaining frame fields (trace, temporary_refs, retained_back, + // etc.) by running destructors in place. The localsplus is already + // empty/heap-backed, so this only drops non-localsplus fields. + // SAFETY: `self` points to valid, initialized memory on the data + // stack. After this call the memory is logically dead. + unsafe { core::ptr::drop_in_place(self) }; + Some(base) + } + /// Get the last instruction index. #[inline(always)] pub fn get_lasti(&self) -> u32 { @@ -1038,6 +1147,7 @@ impl InterpreterFrame { // If we copied Thread from the source iframe, frame.clear() would // reject the frame with "cannot clear an executing frame". owner: atomic::AtomicI8::new(FrameOwner::FrameObject as i8), + datastack_base: core::ptr::null_mut(), materialized: Radium::new(0), cold: OnceCell::from(Box::new(FrameColdData { escaped: atomic::AtomicBool::new(true), @@ -1117,6 +1227,7 @@ impl InterpreterFrame { generator: PyAtomicBorrow::new(), previous: Radium::new(0), owner: atomic::AtomicI8::new(FrameOwner::FrameObject as i8), + datastack_base: core::ptr::null_mut(), materialized: Radium::new(0), cold: OnceCell::from(Box::new(FrameColdData { escaped: atomic::AtomicBool::new(true), @@ -1350,6 +1461,10 @@ unsafe impl Traverse for FrameObject { pub enum ExecutionResult { Return(PyObjectRef), Yield(PyObjectRef), + /// The bytecode loop wants to tail-call into a new frame that has + /// already been prepared on the datastack. The trampoline reads the + /// pending frame pointer from `vm.pending_tailcall_frame`. + TailCall, } /// A valid execution result, or an exception @@ -2094,6 +2209,7 @@ impl Py { func_obj, prev_line: &iframe.prev_line, monitoring_mask: 0, + tailcall_enabled: false, }; f(exec) } @@ -2156,6 +2272,7 @@ impl Py { func_obj, prev_line: &iframe.prev_line, monitoring_mask: 0, + tailcall_enabled: false, }; exec.yield_from_target().map(PyObject::to_owned) } @@ -2179,6 +2296,92 @@ impl Py { } } +/// Byte offset from the start of a datastack allocation to the LocalsPlus data, +/// accounting for alignment padding after the InterpreterFrame header. +#[inline] +fn datastack_iframe_localsplus_offset() -> usize { + let iframe_size = core::mem::size_of::(); + (iframe_size + core::mem::align_of::() - 1) & !(core::mem::align_of::() - 1) +} + +/// Total bytes needed to co-allocate an InterpreterFrame and its LocalsPlus +/// data on the thread data stack. +pub(crate) fn datastack_iframe_total_bytes(nlocalsplus: usize, stacksize: usize) -> usize { + let iframe_padded = datastack_iframe_localsplus_offset(); + let capacity = nlocalsplus + .checked_add(stacksize) + .expect("LocalsPlus capacity overflow"); + let data_bytes = capacity + .checked_mul(core::mem::size_of::()) + .expect("LocalsPlus byte size overflow"); + iframe_padded + .checked_add(data_bytes) + .expect("datastack iframe total size overflow") +} + +/// Handle an exception propagating into a suspended caller frame in the +/// trampoline. Adds a traceback entry at the caller's call site, then +/// tries the caller's exception table via `unwind_blocks`. +/// +/// Returns: +/// - `Ok(None)` — handler found, the caller's `run_iframe` can be re-entered +/// - `Ok(Some(result))` — handler returned a result (break from the run loop) +/// - `Err(exc)` — no handler, exception propagates to the next caller +pub(crate) fn trampoline_handle_exception( + iframe: &mut InterpreterFrame, + exception: &PyBaseExceptionRef, + vm: &VirtualMachine, +) -> FrameResult { + let code: &Py = unsafe { &*iframe.code }; + let globals: &Py = unsafe { &*iframe.globals }; + let builtins: &PyObject = unsafe { &*iframe.builtins }; + let func_obj: Option<&PyObject> = if iframe.func_obj.is_null() { + None + } else { + Some(unsafe { &*iframe.func_obj }) + }; + let builtins_dict = if globals.class().is(vm.ctx.types.dict_type) { + builtins + .downcast_ref_if_exact::(vm) + .map(|d| unsafe { PyExact::ref_unchecked(d) }) + } else { + None + }; + let iframe_ptr = iframe as *const InterpreterFrame; + let mut exec = ExecutingFrame { + code, + localsplus: &mut iframe.localsplus, + locals: &iframe.locals, + globals, + builtins, + builtins_dict, + lasti: &iframe.lasti, + iframe: iframe_ptr, + func_obj, + prev_line: &mut iframe.prev_line, + monitoring_mask: 0, + tailcall_enabled: false, + }; + + // lasti points past the CallPyExactArgs instruction (+ cache entries). + // The exception occurred at the previous instruction (the call site). + let idx = (exec.lasti() as usize).saturating_sub(1); + + // Add traceback entry at the call site. + if let Some((loc, _end_loc)) = exec.code.locations.get(idx) { + let next = exception.__traceback__(); + let new_traceback = PyTraceback::new(next, exec.frame_object(vm), idx as u32 * 2, loc.line); + exception.set_traceback_typed(Some(new_traceback.into_ref(&vm.ctx))); + } + + exec.unwind_blocks( + vm, + UnwindReason::Raising { + exception: exception.clone(), + }, + ) +} + /// Execute an InterpreterFrame's bytecode directly, without a FrameObject. /// /// # Safety @@ -2217,6 +2420,7 @@ pub(crate) fn run_iframe( func_obj, prev_line: &mut iframe.prev_line, monitoring_mask: 0, + tailcall_enabled: true, }; exec.run(vm) } @@ -2246,6 +2450,9 @@ pub(crate) struct ExecutingFrame<'a> { prev_line: &'a core::cell::Cell, /// Cached monitoring events mask. Reloaded at Resume instruction only, monitoring_mask: u32, + /// Whether TailCall is allowed. True when running under the trampoline + /// (`run_frame_fast`), false for FrameObject-based execution. + tailcall_enabled: bool, } #[inline] @@ -5647,7 +5854,11 @@ impl ExecutingFrame<'_> { if self.specialization_call_recursion_guard(vm) { return self.execute_call_vectorcall(nargs, vm); } - // Stage args without a per-call Vec: [self?, arg1, ..., argN] + if self.tailcall_enabled && !func.is_generator_like() { + self.tailcall_prepare_frame(nargs, self_or_null_is_some, vm); + return Ok(Some(ExecutionResult::TailCall)); + } + // Recursive path: pop args and call. let base = usize::from(self_or_null_is_some); let mut arg_buf = CallArgBuffer::new(nargs as usize + base); let args = arg_buf.slots(); @@ -5704,7 +5915,16 @@ impl ExecutingFrame<'_> { if self.specialization_call_recursion_guard(vm) { return self.execute_call_vectorcall(nargs, vm); } - // Stage args without a per-call Vec: + if self.tailcall_enabled && !func.is_generator_like() { + self.tailcall_prepare_bound_method_frame( + nargs, + bound_function, + bound_self, + vm, + ); + return Ok(Some(ExecutionResult::TailCall)); + } + // Recursive path: stage args without a per-call Vec. // [bound_self, arg1, ..., argN] let mut arg_buf = CallArgBuffer::new(nargs as usize + 1); let args = arg_buf.slots(); @@ -10403,6 +10623,129 @@ impl ExecutingFrame<'_> { >= vm.recursion_limit.get() } + /// Prepare a callee frame on the datastack for a TailCall. + /// Pops args, self_or_null, and callable from the caller's stack, + /// builds the callee InterpreterFrame, and stores its pointer in + /// `vm.pending_tailcall_frame`. + /// + /// The callable must be at stack position `nargs + 1` (already validated). + fn tailcall_prepare_frame( + &mut self, + nargs: u32, + self_or_null_is_some: bool, + vm: &VirtualMachine, + ) { + let base = usize::from(self_or_null_is_some); + let effective_nargs = nargs as usize + base; + + // Peek at the callable (still on the stack) to build the callee + // frame. The callable stays on the caller's stack until we're done + // constructing the callee. + let callable = self.nth_value(nargs + 1); + let func = callable.downcast_ref_if_exact::(vm).unwrap(); + + let code: &Py = &func.code; + + let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { + FrameLocals::lazy() + } else { + FrameLocals::with_locals(crate::function::ArgMapping::from_dict_exact( + func.globals.clone(), + )) + }; + + let callee_iframe = InterpreterFrame::new_on_datastack( + code, + &func.globals, + &func.builtins, + Some(func.as_object()), + locals, + func.closure.as_ref().map_or(&[], |c| c.as_slice()), + vm, + ); + + // Move args directly from the caller's stack into callee fastlocals, + // avoiding an intermediate buffer. + { + let fastlocals = callee_iframe.localsplus.fastlocals_mut(); + for (dst, arg) in fastlocals[base..effective_nargs] + .iter_mut() + .zip(self.pop_multiple(nargs as usize)) + { + *dst = Some(arg); + } + let self_or_null = self.pop_value_opt(); + debug_assert_eq!(self_or_null.is_some(), self_or_null_is_some); + if self_or_null.is_some() { + fastlocals[0] = self_or_null; + } + } + + // Pop the callable and transfer ownership to the trampoline via + // the VM side channel, avoiding a per-frame mutex lock on + // temporary_refs. + let callable = self.pop_value(); + unsafe { &mut *vm.pending_tailcall_refs.get() }.push(callable); + + vm.set_pending_tailcall(callee_iframe); + } + + /// Prepare a callee frame for a bound method TailCall. + /// Pops args, self_or_null (null), and callable from the caller's stack, + /// builds the callee InterpreterFrame with bound_self prepended, and + /// stores its pointer in `vm.pending_tailcall_frame`. + fn tailcall_prepare_bound_method_frame( + &mut self, + nargs: u32, + bound_function: PyObjectRef, + bound_self: PyObjectRef, + vm: &VirtualMachine, + ) { + let effective_nargs = nargs as usize + 1; // +1 for bound_self + + let func = bound_function + .downcast_ref_if_exact::(vm) + .unwrap(); + let code: &Py = &func.code; + + let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { + FrameLocals::lazy() + } else { + FrameLocals::with_locals(crate::function::ArgMapping::from_dict_exact( + func.globals.clone(), + )) + }; + + let callee_iframe = InterpreterFrame::new_on_datastack( + code, + &func.globals, + &func.builtins, + Some(func.as_object()), + locals, + func.closure.as_ref().map_or(&[], |c| c.as_slice()), + vm, + ); + + // Move args directly from the caller's stack into callee fastlocals. + let fastlocals = callee_iframe.localsplus.fastlocals_mut(); + for (dst, arg) in fastlocals[1..effective_nargs] + .iter_mut() + .zip(self.pop_multiple(nargs as usize)) + { + *dst = Some(arg); + } + self.pop_value_opt(); // null (self_or_null) + let callable = self.pop_value(); // callable (bound method) + fastlocals[0] = Some(bound_self); + + // Transfer ownership to the trampoline via the VM side channel. + let refs = unsafe { &mut *vm.pending_tailcall_refs.get() }; + refs.push(bound_function); + refs.push(callable); + + vm.set_pending_tailcall(callee_iframe); + } + #[inline] fn for_iter_has_end_for_shape(&self, instr_idx: usize, jump_delta: u32) -> bool { let target_idx = instr_idx diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index a48ce56c14a..e0a086c10db 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -106,6 +106,15 @@ pub struct VirtualMachine { pub asyncio_running_task: RefCell>, pub(crate) callable_cache: CallableCache, pub(crate) audit_hooks: RefCell>, + /// Side channel for TailCall: the bytecode loop stores the new frame + /// pointer here before returning `ExecutionResult::TailCall`. + /// Access only via `set_pending_tailcall` / `take_pending_tailcall`. + pending_tailcall_frame: Cell>, + /// Owned references that keep callee raw pointers valid during TailCall. + /// Set by `tailcall_prepare_frame`, drained by the trampoline into + /// its local `owned_refs` Vec. Uses UnsafeCell because the VM is + /// per-thread and this field is only accessed on the owning thread. + pub(crate) pending_tailcall_refs: core::cell::UnsafeCell>, } /// Non-owning frame pointer for the non-unix threading frames stack. @@ -777,6 +786,60 @@ pub fn process_hash_secret_seed() -> u32 { *SEED.get_or_init(|| u32::from_ne_bytes(rustpython_common::rand::os_random())) } +/// A `NonNull` wrapper that implements `Send + Sync`. +/// +/// # Safety contract +/// +/// This type bypasses Rust's `Send`/`Sync` bounds on `NonNull`. It is +/// sound **only** when the pointer is exclusively accessed by one thread +/// at a time. In this codebase, that invariant is upheld because +/// `VirtualMachine` is per-thread. +/// +/// **Do not use this type outside `pending_tailcall_frame`.** It exists +/// solely to let a `Cell>` field on the per-thread +/// VM satisfy `Send + Sync`. If you need a `Send`-able pointer +/// elsewhere, justify and document the safety invariant at that site. +#[repr(transparent)] +struct PendingFrame(core::ptr::NonNull); + +impl Copy for PendingFrame {} +impl Clone for PendingFrame { + fn clone(&self) -> Self { + *self + } +} + +// SAFETY: VirtualMachine is per-thread; the pointer is only ever +// accessed on the thread that wrote it. The pointed-to InterpreterFrame +// lives on that thread's datastack and is valid from set to take. +unsafe impl Send for PendingFrame {} +unsafe impl Sync for PendingFrame {} + +/// Saved state from `enter_iframe`, needed by `exit_iframe` to restore +/// the previous frame chain and exception state. +pub(crate) struct IframeEntryState { + pub(crate) iframe_ptr: *const crate::frame::InterpreterFrame, + pub(crate) old_chain: *const crate::frame::InterpreterFrame, + pub(crate) saved_exc: Option, + pub(crate) save_exc: bool, +} + +/// Caller frame suspended by a TailCall in the trampoline. +struct SuspendedFrame { + iframe: *mut crate::frame::InterpreterFrame, + entry_state: IframeEntryState, + /// Owned references that keep callee's raw pointers (code, globals, + /// builtins borrowed from PyFunction) valid. Drained from + /// `vm.pending_tailcall_refs` when the callee's TailCall is consumed. + /// Dropped when this SuspendedFrame is popped (after callee returns/errors). + owned_refs: Vec, + /// True for the initial frame passed into the trampoline by the caller. + /// The caller owns the datastack allocation for the entry frame, so the + /// trampoline must NOT release it — only callee-allocated frames are + /// released here. + is_entry: bool, +} + impl VirtualMachine { fn init_callable_cache(&mut self) -> PyResult<()> { self.callable_cache.len = Some(self.builtins.get_attr("len", self)?); @@ -891,6 +954,8 @@ impl VirtualMachine { asyncio_running_task: RefCell::new(None), callable_cache: CallableCache::default(), audit_hooks: RefCell::new(vec![]), + pending_tailcall_frame: Cell::new(None), + pending_tailcall_refs: core::cell::UnsafeCell::new(Vec::with_capacity(2)), }; if vm.state.hash_secret.hash_str("") @@ -1331,13 +1396,294 @@ impl VirtualMachine { } } + /// Store a callee frame pointer for the trampoline to pick up after + /// `TailCall` is returned. The pointed-to InterpreterFrame must live + /// on the current thread's datastack and remain valid until the + /// trampoline calls `take_pending_tailcall`. + #[inline(always)] + pub(crate) fn set_pending_tailcall(&self, iframe: &mut crate::frame::InterpreterFrame) { + self.pending_tailcall_frame + .set(Some(PendingFrame(core::ptr::NonNull::from(iframe)))); + } + + /// Take the pending tailcall frame pointer, resetting the side channel. + #[inline(always)] + fn take_pending_tailcall(&self) -> *mut crate::frame::InterpreterFrame { + self.pending_tailcall_frame + .take() + .expect("TailCall without pending frame") + .0 + .as_ptr() + } + #[inline(always)] /// Run a stack-allocated InterpreterFrame without heap allocation. - /// This is the fast path for regular (non-generator) function calls. + /// Uses a trampoline loop to flatten Python-to-Python calls: when the + /// bytecode loop returns `TailCall`, the trampoline swaps to the new + /// frame without adding a Rust stack frame. pub fn run_frame_fast(&self, iframe: &mut crate::frame::InterpreterFrame) -> PyResult { - match self.with_iframe(iframe, |iframe| crate::frame::run_iframe(iframe, self))? { - ExecutionResult::Return(value) => Ok(value), - _ => panic!("Got unexpected result from function"), + use crate::frame::ExecutionResult; + + let entry_state = self.enter_iframe(iframe)?; + let result = crate::frame::run_iframe(iframe, self); + + match result { + Ok(ExecutionResult::Return(value)) => { + self.exit_iframe(entry_state); + Ok(value) + } + Ok(ExecutionResult::TailCall) => self.run_frame_fast_trampoline(iframe, entry_state), + Ok(ExecutionResult::Yield(_)) => panic!("Yield in non-generator frame"), + Err(exc) => { + self.exit_iframe(entry_state); + Err(exc) + } + } + } + + /// Cold path: at least one TailCall was issued. Run the trampoline. + /// All frame dispatch happens in this single loop — no mutual recursion + /// between helper functions, so C stack depth is bounded. + #[cold] + #[inline(never)] + fn run_frame_fast_trampoline( + &self, + iframe: &mut crate::frame::InterpreterFrame, + entry_state: IframeEntryState, + ) -> PyResult { + use crate::frame::ExecutionResult; + + let mut frame_stack: Vec = Vec::with_capacity(8); + + // What we need to do next. + enum Action { + /// Enter and run a new callee frame (pointer from pending_tailcall_frame). + EnterCallee(*mut crate::frame::InterpreterFrame), + /// Push a return value onto the next caller and re-enter it. + ReturnValue(PyObjectRef), + /// Propagate an exception through suspended callers. + Unwind(PyBaseExceptionRef), + } + + let initial_ptr = self.take_pending_tailcall(); + // Drain the refs that keep the initial callee's raw pointers alive. + let initial_refs = unsafe { &mut *self.pending_tailcall_refs.get() } + .drain(..) + .collect(); + frame_stack.push(SuspendedFrame { + iframe: iframe as *mut crate::frame::InterpreterFrame, + entry_state, + owned_refs: initial_refs, + is_entry: true, + }); + let mut action = Action::EnterCallee(initial_ptr); + + loop { + match action { + Action::EnterCallee(callee_ptr) => { + let callee = unsafe { &mut *callee_ptr }; + let callee_entry = match self.enter_iframe_unchecked(callee) { + Ok(state) => state, + Err(exc) => { + unsafe { + if let Some(base) = callee.release_datastack_frame() { + self.datastack_pop(base); + } + } + action = Action::Unwind(exc); + continue; + } + }; + + let result = crate::frame::run_iframe(callee, self); + match result { + Ok(ExecutionResult::TailCall) => { + let refs = unsafe { &mut *self.pending_tailcall_refs.get() } + .drain(..) + .collect(); + frame_stack.push(SuspendedFrame { + iframe: callee_ptr, + entry_state: callee_entry, + owned_refs: refs, + is_entry: false, + }); + action = Action::EnterCallee(self.take_pending_tailcall()); + } + Ok(ExecutionResult::Return(value)) => { + self.exit_iframe(callee_entry); + unsafe { + if let Some(base) = callee.release_datastack_frame() { + self.datastack_pop(base); + } + } + action = Action::ReturnValue(value); + } + Ok(ExecutionResult::Yield(_)) => panic!("Yield in non-generator frame"), + Err(exc) => { + self.exit_iframe(callee_entry); + unsafe { + if let Some(base) = callee.release_datastack_frame() { + self.datastack_pop(base); + } + } + action = Action::Unwind(exc); + } + } + } + + Action::ReturnValue(value) => { + let Some(caller) = frame_stack.pop() else { + // All frames consumed — this is the final return. + return Ok(value); + }; + let SuspendedFrame { + iframe: caller_iframe_ptr, + entry_state: caller_entry, + owned_refs: _caller_refs, + is_entry: caller_is_entry, + } = caller; + let caller_iframe = unsafe { &mut *caller_iframe_ptr }; + caller_iframe.localsplus.push_stack(value); + + let result = crate::frame::run_iframe(caller_iframe, self); + match result { + Ok(ExecutionResult::TailCall) => { + let refs = unsafe { &mut *self.pending_tailcall_refs.get() } + .drain(..) + .collect(); + drop(_caller_refs); + frame_stack.push(SuspendedFrame { + iframe: caller_iframe_ptr, + entry_state: caller_entry, + owned_refs: refs, + is_entry: caller_is_entry, + }); + action = Action::EnterCallee(self.take_pending_tailcall()); + } + Ok(ExecutionResult::Return(value)) => { + drop(_caller_refs); + self.exit_iframe(caller_entry); + if !caller_is_entry { + unsafe { + if let Some(base) = caller_iframe.release_datastack_frame() { + self.datastack_pop(base); + } + } + } + action = Action::ReturnValue(value); + } + Ok(ExecutionResult::Yield(_)) => panic!("Yield in non-generator frame"), + Err(exc) => { + drop(_caller_refs); + self.exit_iframe(caller_entry); + if !caller_is_entry { + unsafe { + if let Some(base) = caller_iframe.release_datastack_frame() { + self.datastack_pop(base); + } + } + } + action = Action::Unwind(exc); + } + } + } + + Action::Unwind(exc) => { + let Some(caller) = frame_stack.pop() else { + return Err(exc); + }; + let SuspendedFrame { + iframe: caller_iframe_ptr, + entry_state: caller_entry, + owned_refs: _caller_refs, + is_entry: caller_is_entry, + } = caller; + let caller_iframe = unsafe { &mut *caller_iframe_ptr }; + + let handled = + crate::frame::trampoline_handle_exception(caller_iframe, &exc, self); + + match handled { + Ok(None) => { + // Handler found — resume the caller's dispatch loop. + let result = crate::frame::run_iframe(caller_iframe, self); + match result { + Ok(ExecutionResult::TailCall) => { + let refs = unsafe { &mut *self.pending_tailcall_refs.get() } + .drain(..) + .collect(); + drop(_caller_refs); + frame_stack.push(SuspendedFrame { + iframe: caller_iframe_ptr, + entry_state: caller_entry, + owned_refs: refs, + is_entry: caller_is_entry, + }); + action = Action::EnterCallee(self.take_pending_tailcall()); + } + Ok(ExecutionResult::Return(value)) => { + drop(_caller_refs); + self.exit_iframe(caller_entry); + if !caller_is_entry { + unsafe { + if let Some(base) = + caller_iframe.release_datastack_frame() + { + self.datastack_pop(base); + } + } + } + action = Action::ReturnValue(value); + } + Ok(ExecutionResult::Yield(_)) => { + panic!("Yield in non-generator frame") + } + Err(new_exc) => { + drop(_caller_refs); + self.exit_iframe(caller_entry); + if !caller_is_entry { + unsafe { + if let Some(base) = + caller_iframe.release_datastack_frame() + { + self.datastack_pop(base); + } + } + } + action = Action::Unwind(new_exc); + } + } + } + Ok(Some(ExecutionResult::Return(value))) => { + drop(_caller_refs); + self.exit_iframe(caller_entry); + if !caller_is_entry { + unsafe { + if let Some(base) = caller_iframe.release_datastack_frame() { + self.datastack_pop(base); + } + } + } + action = Action::ReturnValue(value); + } + Ok(Some(_)) => { + panic!("Unexpected execution result in trampoline unwind") + } + Err(new_exc) => { + drop(_caller_refs); + self.exit_iframe(caller_entry); + if !caller_is_entry { + unsafe { + if let Some(base) = caller_iframe.release_datastack_frame() { + self.datastack_pop(base); + } + } + } + action = Action::Unwind(new_exc); + } + } + } + } } } @@ -1805,15 +2151,14 @@ impl VirtualMachine { result } - /// Execute a stack-allocated InterpreterFrame without heap-allocating - /// a FrameObject. This is the fast path for regular function calls. - /// The frame is pushed onto the chain as a `*const InterpreterFrame`. - #[inline(always)] - pub fn with_iframe( + /// Push `iframe` onto the frame chain: recursion/C-stack check, TLS + /// link, exception save. Returns the saved state needed by + /// `exit_iframe`. + #[inline] + pub(crate) fn enter_iframe( &self, iframe: &mut crate::frame::InterpreterFrame, - f: impl FnOnce(&mut crate::frame::InterpreterFrame) -> PyResult, - ) -> PyResult { + ) -> PyResult { self.check_recursive_call("")?; let depth = self.recursion_depth.get(); @@ -1821,6 +2166,23 @@ impl VirtualMachine { return Err(self.new_recursion_error(String::new())); } + self.enter_iframe_unchecked(iframe) + } + + /// Like `enter_iframe` but skips the Python recursion depth check + /// (already verified by `specialization_call_recursion_guard`). + /// Still checks C-stack overflow since each `run_iframe` call + /// consumes Rust stack space. + #[inline(always)] + pub(crate) fn enter_iframe_unchecked( + &self, + iframe: &mut crate::frame::InterpreterFrame, + ) -> PyResult { + let depth = self.recursion_depth.get(); + if depth & 7 == 0 && self.check_c_stack_overflow() { + return Err(self.new_recursion_error(String::new())); + } + self.recursion_depth.update(|d| d + 1); let iframe_ptr = iframe as *const crate::frame::InterpreterFrame; @@ -1839,29 +2201,35 @@ impl VirtualMachine { None }; - let result = f(iframe); + Ok(IframeEntryState { + iframe_ptr, + old_chain, + saved_exc, + save_exc, + }) + } + + /// Pop `iframe` from the frame chain: sync materialized state, restore + /// exception, TLS unlink, GC tracking. + pub(crate) fn exit_iframe(&self, state: IframeEntryState) { + let IframeEntryState { + iframe_ptr, + old_chain, + saved_exc, + save_exc, + } = state; // If this iframe was materialized, capture f_back so that code - // holding a reference to the FrameObject (e.g. sys._getframe() - // return value, traceback frames) can walk the chain after return. - // - // Read materialized through the raw TLS pointer instead of the - // &mut iframe reference. During f(iframe), bytecode can - // materialize the frame via the TLS chain (a raw pointer alias); - // the &mut borrow lets LLVM assume no aliased writes, which can - // cause the store to be invisible through `iframe.materialized`. + // holding a reference to the FrameObject can walk the chain after + // return. Read materialized through read_volatile to bypass + // LLVM's noalias on the &mut iframe borrow. { - // Use read_volatile through the original raw pointer to bypass - // LLVM's noalias assumptions on the &mut iframe borrow. let mat_ptr = unsafe { let field_ptr = core::ptr::addr_of!((*iframe_ptr).materialized); core::ptr::read_volatile(field_ptr as *const usize) }; if mat_ptr != 0 { let fo = unsafe { &*(mat_ptr as *const crate::Py) }; - // Sync localsplus, prev_line, lasti from the live iframe to - // the materialized FrameObject so f_locals, f_lineno, f_lasti - // reflect the final state after execution. unsafe { let live_iframe = &*iframe_ptr; fo.iframe_mut() @@ -1879,15 +2247,9 @@ impl VirtualMachine { } if !old_chain.is_null() { let prev_iframe = unsafe { &*old_chain }; - // Use materialize_chain to avoid cloning localsplus, which - // would create extra refcounts on local variables. The - // lightweight frame has empty localsplus; live values are - // read through find_live_source_iframe when needed. let back_fo = prev_iframe.materialize_chain(self); *fo.iframe().cold().retained_back.lock() = Some(back_fo); } - // Set owner to FrameObject since this frame is no longer - // executing on a thread. fo.iframe().owner.store( crate::frame::FrameOwner::FrameObject as i8, core::sync::atomic::Ordering::Release, @@ -1898,15 +2260,22 @@ impl VirtualMachine { if save_exc { self.restore_exception(saved_exc); } - // Restore the frame chain BEFORE clearing temporary_refs, so - // top_frame no longer points at the materialized FrameObject - // when its last strong reference is released. + // Clear previous before popping — it may point to a stack-allocated + // iframe that will be freed when the caller's with_iframe exits. + { + #[allow(unused_imports)] + use rustpython_common::atomic::Radium; + unsafe { + (*iframe_ptr) + .previous + .store(0, core::sync::atomic::Ordering::Relaxed); + } + } let _ = crate::vm::thread::set_current_frame(old_chain); self.recursion_depth.update(|d| d - 1); - // Now that the frame is off the chain, track the materialized - // FrameObject in the GC and release temporary_refs so cycle - // collection can detect and reclaim reference cycles. + // Track the materialized FrameObject in GC and release + // temporary_refs after the frame is off the chain. { let mat_ptr = unsafe { let field_ptr = core::ptr::addr_of!((*iframe_ptr).materialized); @@ -1922,7 +2291,19 @@ impl VirtualMachine { } } } + } + pub fn with_iframe( + &self, + iframe: &mut crate::frame::InterpreterFrame, + f: impl FnOnce(&mut crate::frame::InterpreterFrame) -> PyResult, + ) -> PyResult { + let state = self.enter_iframe(iframe)?; + // Ensure exit_iframe runs even if f(iframe) panics. + let guard = scopeguard::guard(state, |s| self.exit_iframe(s)); + let result = f(iframe); + let state = scopeguard::ScopeGuard::into_inner(guard); + self.exit_iframe(state); result } diff --git a/crates/vm/src/vm/thread.rs b/crates/vm/src/vm/thread.rs index 1520b2cd883..1bab539a0a6 100644 --- a/crates/vm/src/vm/thread.rs +++ b/crates/vm/src/vm/thread.rs @@ -1022,6 +1022,8 @@ impl VirtualMachine { asyncio_running_task: RefCell::new(None), callable_cache: self.callable_cache.clone(), audit_hooks: RefCell::new(vec![]), + pending_tailcall_frame: Cell::new(None), + pending_tailcall_refs: core::cell::UnsafeCell::new(Vec::with_capacity(2)), }; ThreadedVirtualMachine { vm } }