Implement PyStackRef borrowed stack references for LOAD_FAST_BORROW (#7290)

* Implement PyStackRef borrowed stack references for LOAD_FAST_BORROW

* Fix PyStackRef: use NonZeroUsize for niche optimization, revert borrowed refs

- Change PyStackRef.bits from usize to NonZeroUsize so Option<PyStackRef>
  has the same size as Option<PyObjectRef> via niche optimization
- Revert LoadFastBorrow to use clone instead of actual borrowed refs
  to avoid borrowed refs remaining on stack at yield points
- Add static size assertions for Option<PyStackRef>
- Add stackref, fastlocal to cspell dictionary
- Remove debug eprintln statements
- Fix clippy warning for unused push_borrowed

* Add borrowed ref debug_assert to InstrumentedYieldValue, clean up comments

Author: youknowone

.cspell.dict/cpython.txt

@@ -54,6 +54,8 @@ eofs
 evalloop
 excepthandler
 exceptiontable
+fastlocal
+fastlocals
 fblock
 fblocks
 fdescr
@@ -171,6 +173,7 @@ SMALLBUF
 SOABI
 SSLEOF
 stackdepth
+stackref
 staticbase
 stginfo
 storefast

crates/vm/src/frame.rs

@@ -1,7 +1,8 @@
 #[cfg(feature = "flame")]
 use crate::bytecode::InstructionMetadata;
 use crate::{
-    AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
+    AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, PyStackRef, TryFromObject,
+    VirtualMachine,
     builtins::{
         PyBaseException, PyBaseExceptionRef, PyCode, PyCoroutine, PyDict, PyDictRef, PyGenerator,
         PyInterpolation, PyList, PySet, PySlice, PyStr, PyStrInterned, PyTemplate, PyTraceback,
@@ -60,7 +61,7 @@ enum UnwindReason {
 struct FrameState {
     // We need 1 stack per frame
     /// The main data frame of the stack machine
-    stack: BoxVec<Option<PyObjectRef>>,
+    stack: BoxVec<Option<PyStackRef>>,
     /// Cell and free variable references (cellvars + freevars).
     cells_frees: Box<[PyCellRef]>,
     /// Previous line number for LINE event suppression.
@@ -1280,14 +1281,9 @@ impl ExecutingFrame<'_> {
                 // This is 1-indexed to match CPython
                 let idx = index.get(arg) as usize;
                 let stack_len = self.state.stack.len();
-                if stack_len < idx {
-                    eprintln!("CopyItem ERROR: stack_len={}, idx={}", stack_len, idx);
-                    eprintln!("  code: {}", self.code.obj_name);
-                    eprintln!("  lasti: {}", self.lasti());
-                    panic!("CopyItem: stack underflow");
-                }
+                debug_assert!(stack_len >= idx, "CopyItem: stack underflow");
                 let value = self.state.stack[stack_len - idx].clone();
-                self.push_value_opt(value);
+                self.push_stackref_opt(value);
                 Ok(None)
             }
             Instruction::CopyFreeVars { .. } => {
@@ -1868,8 +1864,9 @@ impl ExecutingFrame<'_> {
                 self.push_value(x2);
                 Ok(None)
             }
-            // TODO: Implement true borrow optimization (skip Arc::clone).
-            // Currently this just clones like LoadFast.
+            // Borrow optimization not yet active; falls back to clone.
+            // push_borrowed() is available but disabled until stack
+            // lifetime issues at yield/exception points are resolved.
             Instruction::LoadFastBorrow(idx) => {
                 let idx = idx.get(arg) as usize;
                 let x = unsafe { self.fastlocals.borrow() }[idx]
@@ -2502,6 +2499,14 @@ impl ExecutingFrame<'_> {
                 Ok(None)
             }
             Instruction::YieldValue { arg: oparg } => {
+                debug_assert!(
+                    self.state
+                        .stack
+                        .iter()
+                        .flatten()
+                        .all(|sr| !sr.is_borrowed()),
+                    "borrowed refs on stack at yield point"
+                );
                 let value = self.pop_value();
                 // arg=0: direct yield (wrapped for async generators)
                 // arg=1: yield from await/yield-from (NOT wrapped)
@@ -2681,6 +2686,14 @@ impl ExecutingFrame<'_> {
                 self.unwind_blocks(vm, UnwindReason::Returning { value })
             }
             Instruction::InstrumentedYieldValue => {
+                debug_assert!(
+                    self.state
+                        .stack
+                        .iter()
+                        .flatten()
+                        .all(|sr| !sr.is_borrowed()),
+                    "borrowed refs on stack at yield point"
+                );
                 let value = self.pop_value();
                 if self.monitoring_mask & monitoring::EVENT_PY_YIELD != 0 {
                     let offset = (self.lasti() - 1) * 2;
@@ -3604,18 +3617,24 @@ impl ExecutingFrame<'_> {
 
         let mut elements = elements;
         // Elements on stack from right-to-left:
-        self.state
-            .stack
-            .extend(elements.drain(before + middle..).rev().map(Some));
+        self.state.stack.extend(
+            elements
+                .drain(before + middle..)
+                .rev()
+                .map(|e| Some(PyStackRef::new_owned(e))),
+        );
 
         let middle_elements = elements.drain(before..).collect();
         let t = vm.ctx.new_list(middle_elements);
         self.push_value(t.into());
 
         // Lastly the first reversed values:
-        self.state
-            .stack
-            .extend(elements.into_iter().rev().map(Some));
+        self.state.stack.extend(
+            elements
+                .into_iter()
+                .rev()
+                .map(|e| Some(PyStackRef::new_owned(e))),
+        );
 
         Ok(None)
     }
@@ -3854,9 +3873,12 @@ impl ExecutingFrame<'_> {
         if let Some(elements) = fast_elements {
             return match elements.len().cmp(&size) {
                 core::cmp::Ordering::Equal => {
-                    self.state
-                        .stack
-                        .extend(elements.into_iter().rev().map(Some));
+                    self.state.stack.extend(
+                        elements
+                            .into_iter()
+                            .rev()
+                            .map(|e| Some(PyStackRef::new_owned(e))),
+                    );
                     Ok(None)
                 }
                 core::cmp::Ordering::Greater => Err(vm.new_value_error(format!(
@@ -3922,9 +3944,12 @@ impl ExecutingFrame<'_> {
                 Err(vm.new_value_error(msg))
             }
             PyIterReturn::StopIteration(_) => {
-                self.state
-                    .stack
-                    .extend(elements.into_iter().rev().map(Some));
+                self.state.stack.extend(
+                    elements
+                        .into_iter()
+                        .rev()
+                        .map(|e| Some(PyStackRef::new_owned(e))),
+                );
                 Ok(None)
             }
         }
@@ -4057,43 +4082,75 @@ impl ExecutingFrame<'_> {
     // Block stack functions removed - exception table handles all exception/cleanup
 
     #[inline]
-    #[track_caller] // not a real track_caller but push_value is less useful for debugging
-    fn push_value_opt(&mut self, obj: Option<PyObjectRef>) {
+    #[track_caller]
+    fn push_stackref_opt(&mut self, obj: Option<PyStackRef>) {
         match self.state.stack.try_push(obj) {
             Ok(()) => {}
             Err(_e) => self.fatal("tried to push value onto stack but overflowed max_stackdepth"),
         }
     }
 
+    #[inline]
+    #[track_caller] // not a real track_caller but push_value is less useful for debugging
+    fn push_value_opt(&mut self, obj: Option<PyObjectRef>) {
+        self.push_stackref_opt(obj.map(PyStackRef::new_owned));
+    }
+
     #[inline]
     #[track_caller]
     fn push_value(&mut self, obj: PyObjectRef) {
-        self.push_value_opt(Some(obj));
+        self.push_stackref_opt(Some(PyStackRef::new_owned(obj)));
+    }
+
+    /// Push a borrowed reference onto the stack (no refcount increment).
+    ///
+    /// # Safety
+    /// The object must remain alive until the borrowed ref is consumed.
+    /// The compiler guarantees consumption within the same basic block.
+    #[inline]
+    #[track_caller]
+    #[allow(dead_code)]
+    unsafe fn push_borrowed(&mut self, obj: &PyObject) {
+        self.push_stackref_opt(Some(unsafe { PyStackRef::new_borrowed(obj) }));
     }
 
     #[inline]
     fn push_null(&mut self) {
-        self.push_value_opt(None);
+        self.push_stackref_opt(None);
     }
 
-    /// Pop a value from the stack, returning None if the stack slot is NULL
+    /// Pop a raw stackref from the stack, returning None if the stack slot is NULL.
     #[inline]
-    fn pop_value_opt(&mut self) -> Option<PyObjectRef> {
+    fn pop_stackref_opt(&mut self) -> Option<PyStackRef> {
         match self.state.stack.pop() {
-            Some(slot) => slot, // slot is Option<PyObjectRef>
+            Some(slot) => slot,
             None => self.fatal("tried to pop from empty stack"),
         }
     }
 
+    /// Pop a raw stackref from the stack. Panics if NULL.
     #[inline]
     #[track_caller]
-    fn pop_value(&mut self) -> PyObjectRef {
+    fn pop_stackref(&mut self) -> PyStackRef {
         expect_unchecked(
-            self.pop_value_opt(),
-            "pop value but null found. This is a compiler bug.",
+            self.pop_stackref_opt(),
+            "pop stackref but null found. This is a compiler bug.",
         )
     }
 
+    /// Pop a value from the stack, returning None if the stack slot is NULL.
+    /// Automatically promotes borrowed refs to owned.
+    #[inline]
+    fn pop_value_opt(&mut self) -> Option<PyObjectRef> {
+        self.pop_stackref_opt().map(|sr| sr.to_pyobj())
+    }
+
+    #[inline]
+    #[track_caller]
+    fn pop_value(&mut self) -> PyObjectRef {
+        self.pop_stackref().to_pyobj()
+    }
+
     fn call_intrinsic_1(
         &mut self,
         func: bytecode::IntrinsicFunction1,
@@ -4262,22 +4319,23 @@ impl ExecutingFrame<'_> {
             );
         }
         self.state.stack.drain(stack_len - count..).map(|obj| {
-            expect_unchecked(obj, "pop_multiple but null found. This is a compiler bug.")
+            expect_unchecked(obj, "pop_multiple but null found. This is a compiler bug.").to_pyobj()
         })
     }
 
     #[inline]
-    fn replace_top(&mut self, mut top: Option<PyObjectRef>) -> Option<PyObjectRef> {
+    fn replace_top(&mut self, top: Option<PyObjectRef>) -> Option<PyObjectRef> {
+        let mut slot = top.map(PyStackRef::new_owned);
         let last = self.state.stack.last_mut().unwrap();
-        core::mem::swap(last, &mut top);
-        top
+        core::mem::swap(last, &mut slot);
+        slot.map(|sr| sr.to_pyobj())
     }
 
     #[inline]
     #[track_caller]
     fn top_value(&self) -> &PyObject {
         match &*self.state.stack {
-            [.., Some(last)] => last,
+            [.., Some(last)] => last.as_object(),
             [.., None] => self.fatal("tried to get top of stack but got NULL"),
             [] => self.fatal("tried to get top of stack but stack is empty"),
         }
@@ -4288,7 +4346,7 @@ impl ExecutingFrame<'_> {
     fn nth_value(&self, depth: u32) -> &PyObject {
         let stack = &self.state.stack;
         match &stack[stack.len() - depth as usize - 1] {
-            Some(obj) => obj,
+            Some(obj) => obj.as_object(),
             None => unsafe { core::hint::unreachable_unchecked() },
         }
     }
@@ -4415,7 +4473,7 @@ fn is_module_initializing(module: &PyObject, vm: &VirtualMachine) -> bool {
     initializing_attr.try_to_bool(vm).unwrap_or(false)
 }
 
-fn expect_unchecked(optional: Option<PyObjectRef>, err_msg: &'static str) -> PyObjectRef {
+fn expect_unchecked<T: fmt::Debug>(optional: Option<T>, err_msg: &'static str) -> T {
     if cfg!(debug_assertions) {
         optional.expect(err_msg)
     } else {

crates/vm/src/lib.rs

@@ -98,7 +98,7 @@ pub mod windows;
 pub use self::convert::{TryFromBorrowedObject, TryFromObject};
 pub use self::object::{
     AsObject, Py, PyAtomicRef, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact,
-    PyResult, PyWeakRef,
+    PyResult, PyStackRef, PyWeakRef,
 };
 pub use self::vm::{Context, Interpreter, InterpreterBuilder, Settings, VirtualMachine};