Fix vectorcall compatibility with LocalsPlus API

youknowone · youknowone · commit 01459eca1940 · 2026-03-04T00:17:52.000+09:00
Update vectorcall dispatch functions to use localsplus stack
accessors instead of direct stack field access. Add
stack_truncate method to LocalsPlus. Update vectorcall fast
path in function.rs to use datastack and fastlocals_mut().
diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs
@@ -1307,26 +1307,33 @@ pub(crate) fn vectorcall_function(
             zelf.builtins.clone(),
             zelf.closure.as_ref().map_or(&[], |c| c.as_slice()),
             Some(zelf.to_owned().into()),
+            true, // Always use datastack (is_simple excludes gen/coro)
             vm,
         )
         .into_ref(&vm.ctx);
 
         {
-            let fastlocals = unsafe { frame.fastlocals.borrow_mut() };
+            let fastlocals = unsafe { frame.fastlocals_mut() };
             for (slot, arg) in fastlocals.iter_mut().zip(args.drain(..nargs)) {
                 *slot = Some(arg);
             }
         }
 
         if let Some(cell2arg) = code.cell2arg.as_deref() {
-            let fastlocals = unsafe { frame.fastlocals.borrow_mut() };
+            let fastlocals = unsafe { frame.fastlocals_mut() };
             for (cell_idx, arg_idx) in cell2arg.iter().enumerate().filter(|(_, i)| **i != -1) {
                 let x = fastlocals[*arg_idx as usize].take();
                 frame.set_cell_contents(cell_idx, x);
             }
         }
 
-        return vm.run_frame(frame);
+        let result = vm.run_frame(frame.clone());
+        unsafe {
+            if let Some(base) = frame.materialize_localsplus() {
+                vm.datastack_pop(base);
+            }
+        }
+        return result;
     }
 
     // SLOW PATH: construct FuncArgs from owned Vec and delegate to invoke()
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -380,6 +380,14 @@ impl LocalsPlus {
         data.swap(base + a, base + b);
     }
 
+    /// Truncate the stack to `new_len` elements, dropping excess values.
+    fn stack_truncate(&mut self, new_len: usize) {
+        debug_assert!(new_len <= self.stack_top as usize);
+        while self.stack_top as usize > new_len {
+            let _ = self.stack_pop();
+        }
+    }
+
     /// Clear the stack, dropping all values.
     fn stack_clear(&mut self) {
         while self.stack_top > 0 {
@@ -6050,13 +6058,15 @@ impl ExecutingFrame<'_> {
     #[inline]
     fn execute_call_vectorcall(&mut self, nargs: u32, vm: &VirtualMachine) -> FrameResult {
         let nargs_usize = nargs as usize;
-        let stack_len = self.stack.len();
+        let stack_len = self.localsplus.stack_len();
         let callable_idx = stack_len - nargs_usize - 2;
         let self_or_null_idx = stack_len - nargs_usize - 1;
         let args_start = stack_len - nargs_usize;
 
         // Check if callable has vectorcall slot
-        let has_vectorcall = self.stack[callable_idx]
+        let has_vectorcall = self
+            .localsplus
+            .stack_index(callable_idx)
             .as_ref()
             .is_some_and(|sr| sr.as_object().class().slots.vectorcall.load().is_some());
 
@@ -6067,7 +6077,9 @@ impl ExecutingFrame<'_> {
         }
 
         // Build args slice: [self_or_null?, arg1, ..., argN]
-        let self_or_null = self.stack[self_or_null_idx]
+        let self_or_null = self
+            .localsplus
+            .stack_index_mut(self_or_null_idx)
             .take()
             .map(|sr| sr.to_pyobj());
         let has_self = self_or_null.is_some();
@@ -6082,12 +6094,12 @@ impl ExecutingFrame<'_> {
             args_vec.push(self_val);
         }
         for stack_idx in args_start..stack_len {
-            let val = self.stack[stack_idx].take().unwrap().to_pyobj();
+            let val = self.localsplus.stack_index_mut(stack_idx).take().unwrap().to_pyobj();
             args_vec.push(val);
         }
 
-        let callable_obj = self.stack[callable_idx].take().unwrap().to_pyobj();
-        self.stack.truncate(callable_idx);
+        let callable_obj = self.localsplus.stack_index_mut(callable_idx).take().unwrap().to_pyobj();
+        self.localsplus.stack_truncate(callable_idx);
 
         let result = callable_obj.vectorcall(args_vec, effective_nargs, None, vm)?;
         self.push_value(result);
@@ -6106,13 +6118,15 @@ impl ExecutingFrame<'_> {
             .expect("kwarg names should be tuple");
         let kw_count = kwarg_names_tuple.len();
 
-        let stack_len = self.stack.len();
+        let stack_len = self.localsplus.stack_len();
         let callable_idx = stack_len - nargs_usize - 2;
         let self_or_null_idx = stack_len - nargs_usize - 1;
         let args_start = stack_len - nargs_usize;
 
         // Check if callable has vectorcall slot
-        let has_vectorcall = self.stack[callable_idx]
+        let has_vectorcall = self
+            .localsplus
+            .stack_index(callable_idx)
             .as_ref()
             .is_some_and(|sr| sr.as_object().class().slots.vectorcall.load().is_some());
 
@@ -6143,7 +6157,9 @@ impl ExecutingFrame<'_> {
         }
 
         // Build args: [self?, pos_arg1, ..., pos_argM, kw_val1, ..., kw_valK]
-        let self_or_null = self.stack[self_or_null_idx]
+        let self_or_null = self
+            .localsplus
+            .stack_index_mut(self_or_null_idx)
             .take()
             .map(|sr| sr.to_pyobj());
         let has_self = self_or_null.is_some();
@@ -6158,12 +6174,12 @@ impl ExecutingFrame<'_> {
             args_vec.push(self_val);
         }
         for stack_idx in args_start..stack_len {
-            let val = self.stack[stack_idx].take().unwrap().to_pyobj();
+            let val = self.localsplus.stack_index_mut(stack_idx).take().unwrap().to_pyobj();
             args_vec.push(val);
         }
 
-        let callable_obj = self.stack[callable_idx].take().unwrap().to_pyobj();
-        self.stack.truncate(callable_idx);
+        let callable_obj = self.localsplus.stack_index_mut(callable_idx).take().unwrap().to_pyobj();
+        self.localsplus.stack_truncate(callable_idx);
 
         let kwnames = kwarg_names_tuple.as_slice();
         let result = callable_obj.vectorcall(args_vec, effective_nargs, Some(kwnames), vm)?;

Original file line number	Diff line number	Diff line change
`@@ -1307,26 +1307,33 @@ pub(crate) fn vectorcall_function(`
`1307`	`1307`	`zelf.builtins.clone(),`
`1308`	`1308`	`zelf.closure.as_ref().map_or(&[], \|c\| c.as_slice()),`
`1309`	`1309`	`Some(zelf.to_owned().into()),`
	`1310`	`+ true, // Always use datastack (is_simple excludes gen/coro)`
`1310`	`1311`	`vm,`
`1311`	`1312`	`)`
`1312`	`1313`	`.into_ref(&vm.ctx);`
`1313`	`1314`
`1314`	`1315`	`{`
`1315`		`- let fastlocals = unsafe { frame.fastlocals.borrow_mut() };`
	`1316`	`+ let fastlocals = unsafe { frame.fastlocals_mut() };`
`1316`	`1317`	`for (slot, arg) in fastlocals.iter_mut().zip(args.drain(..nargs)) {`
`1317`	`1318`	`*slot = Some(arg);`
`1318`	`1319`	`}`
`1319`	`1320`	`}`
`1320`	`1321`
`1321`	`1322`	`if let Some(cell2arg) = code.cell2arg.as_deref() {`
`1322`		`- let fastlocals = unsafe { frame.fastlocals.borrow_mut() };`
	`1323`	`+ let fastlocals = unsafe { frame.fastlocals_mut() };`
`1323`	`1324`	`for (cell_idx, arg_idx) in cell2arg.iter().enumerate().filter(\|(_, i)\| **i != -1) {`
`1324`	`1325`	`let x = fastlocals[*arg_idx as usize].take();`
`1325`	`1326`	`frame.set_cell_contents(cell_idx, x);`
`1326`	`1327`	`}`
`1327`	`1328`	`}`
`1328`	`1329`
`1329`		`- return vm.run_frame(frame);`
	`1330`	`+ let result = vm.run_frame(frame.clone());`
	`1331`	`+ unsafe {`
	`1332`	`+ if let Some(base) = frame.materialize_localsplus() {`
	`1333`	`+ vm.datastack_pop(base);`
	`1334`	`+ }`
	`1335`	`+ }`
	`1336`	`+ return result;`
`1330`	`1337`	`}`
`1331`	`1338`
`1332`	`1339`	`// SLOW PATH: construct FuncArgs from owned Vec and delegate to invoke()`