Use DataStack for LocalsPlus in non-generator function calls

Normal function calls now bump-allocate LocalsPlus data from the per-thread DataStack instead of a separate heap allocation. Generator/coroutine frames continue using heap allocation since they outlive the call. On frame exit, data is copied to the heap (materialize_to_heap) to preserve locals for tracebacks, then the DataStack is popped. VirtualMachine.datastack is wrapped in UnsafeCell for interior mutability (safe because frame allocation is single-threaded LIFO).
RustPython · youknowone · Mar 4, 2026 · Mar 3, 2026 · Mar 3, 2026 · Mar 3, 2026
commit 9dbe95d7ea80f8fa9e39299c9bdd3d4473ac1971
diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs
@@ -562,6 +562,7 @@ 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(
@@ -570,6 +571,7 @@ impl Py<PyFunction> {
             self.builtins.clone(),
             self.closure.as_ref().map_or(&[], |c| c.as_slice()),
             Some(self.to_owned().into()),
+            use_datastack,
             vm,
         )
         .into_ref(&vm.ctx);
@@ -594,7 +596,16 @@ impl Py<PyFunction> {
                 frame.set_generator(&obj);
                 Ok(obj)
             }
-            (false, false) => vm.run_frame(frame),
+            (false, false) => {
+                let result = vm.run_frame(frame.clone());
+                // Release data stack memory after frame execution completes.
+                unsafe {
+                    if let Some(base) = frame.materialize_localsplus() {
+                        vm.datastack_pop(base);
+                    }
+                }
+                result
+            }
         }
     }
 
@@ -665,6 +676,7 @@ impl Py<PyFunction> {
             self.builtins.clone(),
             self.closure.as_ref().map_or(&[], |c| c.as_slice()),
             Some(self.to_owned().into()),
+            true, // Always use datastack (invoke_exact_args is never gen/coro)
             vm,
         )
         .into_ref(&vm.ctx);
@@ -686,7 +698,13 @@ impl Py<PyFunction> {
             }
         }
 
-        vm.run_frame(frame)
+        let result = vm.run_frame(frame.clone());
+        unsafe {
+            if let Some(base) = frame.materialize_localsplus() {
+                vm.datastack_pop(base);
+            }
+        }
+        result
     }
 }
 

diff --git a/crates/vm/src/datastack.rs b/crates/vm/src/datastack.rs
@@ -101,7 +101,8 @@ impl DataStack {
     #[inline(never)]
     fn push_slow(&mut self, aligned_size: usize) -> *mut u8 {
         let mut chunk_size = MIN_CHUNK_SIZE;
-        let needed = aligned_size + MINIMUM_OVERHEAD + core::mem::size_of::<DataStackChunk>() + ALIGN;
+        let needed =
+            aligned_size + MINIMUM_OVERHEAD + core::mem::size_of::<DataStackChunk>() + ALIGN;
         while chunk_size < needed {
             chunk_size *= 2;
         }
@@ -142,10 +143,7 @@ impl DataStack {
     unsafe fn pop_slow(&mut self, base: *mut u8) {
         let old_chunk = self.chunk;
         let prev = unsafe { (*old_chunk).previous };
-        debug_assert!(
-            !prev.is_null(),
-            "tried to pop past the root chunk"
-        );
+        debug_assert!(!prev.is_null(), "tried to pop past the root chunk");
         unsafe { Self::free_chunk(old_chunk) };
         self.chunk = prev;
         self.top = base;