frame helper in PyFuncion

youknowone · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025
commit ae9e6634eaeb8900caadd121d00588fd3ff71fd7
diff --git a/vm/src/builtins/function.rs b/vm/src/builtins/function.rs
@@ -317,89 +317,11 @@ impl PyFunction {
 
         Ok(())
     }
-}
-
-impl Py<PyFunction> {
-    pub fn invoke_with_locals(
-        &self,
-        func_args: FuncArgs,
-        locals: Option<ArgMapping>,
-        vm: &VirtualMachine,
-    ) -> PyResult {
-        #[cfg(feature = "jit")]
-        if let Some(jitted_code) = self.jitted_code.get() {
-            use crate::convert::ToPyObject;
-            match jit::get_jit_args(self, &func_args, jitted_code, vm) {
-                Ok(args) => {
-                    return Ok(args.invoke().to_pyobject(vm));
-                }
-                Err(err) => info!(
-                    "jit: function `{}` is falling back to being interpreted because of the \
-                    error: {}",
-                    self.code.obj_name, err
-                ),
-            }
-        }
-
-        let code = &self.code;
-
-        let locals = if self.code.flags.contains(bytecode::CodeFlags::NEW_LOCALS) {
-            ArgMapping::from_dict_exact(vm.ctx.new_dict())
-        } else if let Some(locals) = locals {
-            locals
-        } else {
-            ArgMapping::from_dict_exact(self.globals.clone())
-        };
-
-        // Construct frame:
-        let frame = Frame::new(
-            code.clone(),
-            Scope::new(Some(locals), self.globals.clone()),
-            vm.builtins.dict(),
-            self.closure.deref().as_ref().map_or(&[], |tuple| {
-                // SAFETY: We know closure contains only cells from construction
-                unsafe {
-                    std::slice::from_raw_parts(
-                        tuple.as_slice().as_ptr() as *const PyCellRef,
-                        tuple.len(),
-                    )
-                }
-            }),
-            Some(self.to_owned().into()),
-            vm,
-        )
-        .into_ref(&vm.ctx);
-
-        self.fill_locals_from_args(&frame, func_args, vm)?;
-
-        // If we have a generator, create a new generator
-        let is_gen = code.flags.contains(bytecode::CodeFlags::IS_GENERATOR);
-        let is_coro = code.flags.contains(bytecode::CodeFlags::IS_COROUTINE);
-        match (is_gen, is_coro) {
-            (true, false) => Ok(PyGenerator::new(frame, self.__name__()).into_pyobject(vm)),
-            (false, true) => Ok(PyCoroutine::new(frame, self.__name__()).into_pyobject(vm)),
-            (true, true) => Ok(PyAsyncGen::new(frame, self.__name__()).into_pyobject(vm)),
-            (false, false) => vm.run_frame(frame),
-        }
-    }
-
-    #[inline(always)]
-    pub fn invoke(&self, func_args: FuncArgs, vm: &VirtualMachine) -> PyResult {
-        self.invoke_with_locals(func_args, None, vm)
-    }
-
-    /// SAFETY: This method should only be called from frame.rs for SET_FUNCTION_ATTRIBUTE
-    /// It allows setting the closure after function creation, which is required for
-    /// CPython 3.13 compatibility. The closure must be a valid tuple of cells.
-    #[doc(hidden)]
-    pub(crate) fn set_closure(&self, closure: Option<PyTupleRef>) {
-        let _ = unsafe { self.closure.swap(closure) };
-    }
 
     /// Set function attribute based on MakeFunctionFlags
     /// This is used by SET_FUNCTION_ATTRIBUTE instruction in CPython 3.13 style
     pub(crate) fn set_function_attribute(
-        &self,
+        &mut self,
         attr: bytecode::MakeFunctionFlags,
         attr_value: PyObjectRef,
         vm: &VirtualMachine,
@@ -458,7 +380,7 @@ impl Py<PyFunction> {
                 .collect();
             let cells_tuple = PyTuple::new_ref(cells_objects, &vm.ctx);
 
-            self.set_closure(Some(cells_tuple));
+            let _ = unsafe { self.closure.swap(Some(cells_tuple)) };
         } else if attr.contains(bytecode::MakeFunctionFlags::TYPE_PARAMS) {
             let type_params = attr_value.clone().downcast::<PyTuple>().map_err(|_| {
                 vm.new_type_error(format!(
@@ -472,6 +394,76 @@ impl Py<PyFunction> {
     }
 }
 
+impl Py<PyFunction> {
+    pub fn invoke_with_locals(
+        &self,
+        func_args: FuncArgs,
+        locals: Option<ArgMapping>,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        #[cfg(feature = "jit")]
+        if let Some(jitted_code) = self.jitted_code.get() {
+            use crate::convert::ToPyObject;
+            match jit::get_jit_args(self, &func_args, jitted_code, vm) {
+                Ok(args) => {
+                    return Ok(args.invoke().to_pyobject(vm));
+                }
+                Err(err) => info!(
+                    "jit: function `{}` is falling back to being interpreted because of the \
+                    error: {}",
+                    self.code.obj_name, err
+                ),
+            }
+        }
+
+        let code = &self.code;
+
+        let locals = if self.code.flags.contains(bytecode::CodeFlags::NEW_LOCALS) {
+            ArgMapping::from_dict_exact(vm.ctx.new_dict())
+        } else if let Some(locals) = locals {
+            locals
+        } else {
+            ArgMapping::from_dict_exact(self.globals.clone())
+        };
+
+        // Construct frame:
+        let frame = Frame::new(
+            code.clone(),
+            Scope::new(Some(locals), self.globals.clone()),
+            vm.builtins.dict(),
+            self.closure.deref().as_ref().map_or(&[], |tuple| {
+                // SAFETY: We know closure contains only cells from construction
+                unsafe {
+                    std::slice::from_raw_parts(
+                        tuple.as_slice().as_ptr() as *const PyCellRef,
+                        tuple.len(),
+                    )
+                }
+            }),
+            Some(self.to_owned().into()),
+            vm,
+        )
+        .into_ref(&vm.ctx);
+
+        self.fill_locals_from_args(&frame, func_args, vm)?;
+
+        // If we have a generator, create a new generator
+        let is_gen = code.flags.contains(bytecode::CodeFlags::IS_GENERATOR);
+        let is_coro = code.flags.contains(bytecode::CodeFlags::IS_COROUTINE);
+        match (is_gen, is_coro) {
+            (true, false) => Ok(PyGenerator::new(frame, self.__name__()).into_pyobject(vm)),
+            (false, true) => Ok(PyCoroutine::new(frame, self.__name__()).into_pyobject(vm)),
+            (true, true) => Ok(PyAsyncGen::new(frame, self.__name__()).into_pyobject(vm)),
+            (false, false) => vm.run_frame(frame),
+        }
+    }
+
+    #[inline(always)]
+    pub fn invoke(&self, func_args: FuncArgs, vm: &VirtualMachine) -> PyResult {
+        self.invoke_with_locals(func_args, None, vm)
+    }
+}
+
 impl PyPayload for PyFunction {
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {

diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -1888,7 +1888,13 @@ impl ExecutingFrame<'_> {
             .downcast_ref::<PyFunction>()
             .expect("SET_FUNCTION_ATTRIBUTE expects function on stack");
 
-        func_ref.set_function_attribute(attr, attr_value, vm)?;
+        let payload: &PyFunction = func_ref.payload();
+        // SetFunctionAttribute always follows MakeFunction, so at this point
+        // there are no other references to func. It is therefore safe to treat it as mutable.
+        unsafe {
+            let payload_ptr = payload as *const PyFunction as *mut PyFunction;
+            (*payload_ptr).set_function_attribute(attr, attr_value, vm)?;
+        };
 
         // Push function back onto stack
         self.push_value(func);

diff --git a/vm/src/object/core.rs b/vm/src/object/core.rs
@@ -906,6 +906,11 @@ impl<T: PyObjectPayload> Py<T> {
             _marker: PhantomData,
         })
     }
+
+    pub fn payload(&self) -> &T {
+        // SAFETY: we know the payload is T because of the type parameter
+        unsafe { self.as_object().payload_unchecked() }
+    }
 }
 
 impl<T> ToOwned for Py<T> {