Add datastack-backed FastLocals for non-generator frames

youknowone · youknowone · commit 968372a13022 · 2026-03-04T23:56:43.000+09:00
Introduce FastLocalsData enum with Heap and DataStack variants
so non-generator/coroutine frames allocate localsplus on the VM
datastack instead of the heap. Includes materialize_to_heap for
migration when needed (e.g. generator suspension).
diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs
@@ -419,8 +419,10 @@ impl PyFunction {
                 .zip(&*code.varnames)
                 .skip(code.arg_count as usize)
                 .take(code.kwonlyarg_count as usize)
-                .filter(|(slot, _)| slot.is_none())
             {
+                if slot.is_some() {
+                    continue;
+                }
                 if let Some(defaults) = &get_defaults!().1
                     && let Some(default) = defaults.get_item_opt(&**kwarg, vm)?
                 {
@@ -563,7 +565,6 @@ impl Py<PyFunction> {
         let is_gen = code.flags.contains(bytecode::CodeFlags::GENERATOR);
         let is_coro = code.flags.contains(bytecode::CodeFlags::COROUTINE);
         let use_datastack = !(is_gen || is_coro);
-
         // Construct frame:
         let frame = Frame::new(
             code,
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -112,8 +112,16 @@ impl FrameOwner {
 /// if acquired, the frame is not executing and access is exclusive; if not,
 /// the caller is on the same thread as `with_exec()` (trace callback) and
 /// access is safe because frame execution is single-threaded.
+enum FastLocalsData {
+    Heap(Box<[Option<PyObjectRef>]>),
+    DataStack {
+        ptr: *mut Option<PyObjectRef>,
+        len: usize,
+    },
+}
+
 pub struct FastLocals {
-    inner: UnsafeCell<Box<[Option<PyObjectRef>]>>,
+    inner: UnsafeCell<FastLocalsData>,
 }
 
 // SAFETY: Frame execution is serialized by the state mutex.
@@ -125,31 +133,96 @@ unsafe impl Sync for FastLocals {}
 impl FastLocals {
     fn new(data: Box<[Option<PyObjectRef>]>) -> Self {
         Self {
-            inner: UnsafeCell::new(data),
+            inner: UnsafeCell::new(FastLocalsData::Heap(data)),
+        }
+    }
+
+    fn new_on_datastack(len: usize, vm: &VirtualMachine) -> Self {
+        let byte_size = len
+            .checked_mul(core::mem::size_of::<Option<PyObjectRef>>())
+            .expect("FastLocals byte size overflow");
+        let ptr = vm.datastack_push(byte_size) as *mut Option<PyObjectRef>;
+        for i in 0..len {
+            // SAFETY: `ptr` points to `len` contiguous slots from this thread's datastack.
+            unsafe { ptr.add(i).write(None) };
+        }
+        Self {
+            inner: UnsafeCell::new(FastLocalsData::DataStack { ptr, len }),
+        }
+    }
+
+    /// Migrate data-stack-backed fastlocals to the heap and return the
+    /// datastack base pointer for `VirtualMachine::datastack_pop`.
+    ///
+    /// # Safety
+    /// Caller must ensure exclusive access and pop the returned base pointer
+    /// in LIFO order on the owning VM.
+    pub unsafe fn materialize_to_heap(&self) -> Option<*mut u8> {
+        let inner = unsafe { &mut *self.inner.get() };
+        let (ptr, len) = match inner {
+            FastLocalsData::Heap(_) => return None,
+            FastLocalsData::DataStack { ptr, len } => (*ptr, *len),
+        };
+        let base = ptr as *mut u8;
+        let mut data = Vec::with_capacity(len);
+        for i in 0..len {
+            // SAFETY: `ptr` is valid for `len` initialized slots.
+            data.push(unsafe { ptr.add(i).read() });
         }
+        *inner = FastLocalsData::Heap(data.into_boxed_slice());
+        Some(base)
     }
 
     /// # Safety
     /// Caller must ensure exclusive access (frame state locked or frame
     /// not executing).
     #[inline(always)]
     pub unsafe fn borrow(&self) -> &[Option<PyObjectRef>] {
-        unsafe { &*self.inner.get() }
+        match unsafe { &*self.inner.get() } {
+            FastLocalsData::Heap(data) => data,
+            FastLocalsData::DataStack { ptr, len } => {
+                // SAFETY: `ptr` points to `len` initialized slots.
+                unsafe { core::slice::from_raw_parts(*ptr, *len) }
+            }
+        }
     }
 
     /// # Safety
     /// Caller must ensure exclusive mutable access.
     #[inline(always)]
     #[allow(clippy::mut_from_ref)]
     pub unsafe fn borrow_mut(&self) -> &mut [Option<PyObjectRef>] {
-        unsafe { &mut *self.inner.get() }
+        match unsafe { &mut *self.inner.get() } {
+            FastLocalsData::Heap(data) => data,
+            FastLocalsData::DataStack { ptr, len } => {
+                // SAFETY: `ptr` points to `len` initialized slots.
+                unsafe { core::slice::from_raw_parts_mut(*ptr, *len) }
+            }
+        }
+    }
+}
+
+impl Drop for FastLocals {
+    fn drop(&mut self) {
+        let inner = unsafe { &mut *self.inner.get() };
+        if let FastLocalsData::DataStack { ptr, len } = inner {
+            for i in 0..*len {
+                // SAFETY: `ptr` points to `len` initialized slots.
+                unsafe { core::ptr::drop_in_place(ptr.add(i)) };
+            }
+        }
     }
 }
 
 unsafe impl Traverse for FastLocals {
     fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
         // SAFETY: GC runs on the same thread; no concurrent mutation.
-        let data = unsafe { &*self.inner.get() };
+        let data: &[Option<PyObjectRef>] = match unsafe { &*self.inner.get() } {
+            FastLocalsData::Heap(data) => data,
+            FastLocalsData::DataStack { ptr, len } => unsafe {
+                core::slice::from_raw_parts(*ptr, *len)
+            },
+        };
         data.traverse(traverse_fn);
     }
 }
@@ -322,6 +395,7 @@ impl Frame {
         builtins: PyObjectRef,
         closure: &[PyCellRef],
         func_obj: Option<PyObjectRef>,
+        use_datastack: bool,
         vm: &VirtualMachine,
     ) -> Self {
         let nlocals = code.varnames.len();
@@ -336,11 +410,19 @@ impl Frame {
 
         // Extend fastlocals to include varnames + cellvars + freevars (localsplus)
         let total_locals = nlocals + num_cells + nfrees;
-        let mut fastlocals_vec: Vec<Option<PyObjectRef>> = vec![None; total_locals];
+        let fastlocals = if use_datastack {
+            FastLocals::new_on_datastack(total_locals, vm)
+        } else {
+            FastLocals::new(vec![None; total_locals].into_boxed_slice())
+        };
 
         // Store cell objects at cellvars and freevars positions
-        for (i, cell) in cells_frees.iter().enumerate() {
-            fastlocals_vec[nlocals + i] = Some(cell.clone().into());
+        {
+            // SAFETY: Frame is under construction and not visible to others yet.
+            let slots = unsafe { fastlocals.borrow_mut() };
+            for (i, cell) in cells_frees.iter().enumerate() {
+                slots[nlocals + i] = Some(cell.clone().into());
+            }
         }
 
         let state = FrameState {
@@ -350,7 +432,7 @@ impl Frame {
         };
 
         Self {
-            fastlocals: FastLocals::new(fastlocals_vec.into_boxed_slice()),
+            fastlocals,
             locals: match scope.locals {
                 Some(locals) => FrameLocals::with_locals(locals),
                 None if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) => FrameLocals::lazy(),
@@ -377,6 +459,34 @@ impl Frame {
         }
     }
 
+    /// Access fastlocals immutably.
+    ///
+    /// # Safety
+    /// Caller must ensure no concurrent mutable access.
+    #[inline(always)]
+    pub unsafe fn fastlocals(&self) -> &[Option<PyObjectRef>] {
+        unsafe { self.fastlocals.borrow() }
+    }
+
+    /// Access fastlocals mutably.
+    ///
+    /// # Safety
+    /// Caller must ensure exclusive access.
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub unsafe fn fastlocals_mut(&self) -> &mut [Option<PyObjectRef>] {
+        unsafe { self.fastlocals.borrow_mut() }
+    }
+
+    /// Migrate data-stack-backed fastlocals to heap storage and return the
+    /// datastack base pointer for `VirtualMachine::datastack_pop`.
+    ///
+    /// # Safety
+    /// Caller must ensure frame is not executing.
+    pub(crate) unsafe fn materialize_localsplus(&self) -> Option<*mut u8> {
+        unsafe { self.fastlocals.materialize_to_heap() }
+    }
+
     /// Clear evaluation stack and state-owned cell/free references.
     /// For full local/cell cleanup, call `clear_locals_and_stack()`.
     pub(crate) fn clear_stack_and_cells(&self) {
