remove closure lock

youknowone · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025
commit a5291c5da429f1314fae47d276761261602fe570
diff --git a/vm/src/builtins/function.rs b/vm/src/builtins/function.rs
@@ -5,7 +5,6 @@ use super::{
     PyAsyncGen, PyCode, PyCoroutine, PyDictRef, PyGenerator, PyStr, PyStrRef, PyTuple, PyTupleRef,
     PyType, PyTypeRef,
 };
-use crate::PyAtomicRef;
 #[cfg(feature = "jit")]
 use crate::common::lock::OnceCell;
 use crate::common::lock::PyMutex;
@@ -32,23 +31,23 @@ pub struct PyFunction {
     code: PyRef<PyCode>,
     globals: PyDictRef,
     builtins: PyObjectRef,
-    closure: PyAtomicRef<Option<PyTuple>>,
+    closure: Option<PyRef<PyTuple<PyCellRef>>>,
     defaults_and_kwdefaults: PyMutex<(Option<PyTupleRef>, Option<PyDictRef>)>,
     name: PyMutex<PyStrRef>,
     qualname: PyMutex<PyStrRef>,
     type_params: PyMutex<PyTupleRef>,
-    #[cfg(feature = "jit")]
-    jitted_code: OnceCell<CompiledCode>,
     annotations: PyMutex<PyDictRef>,
     module: PyMutex<PyObjectRef>,
     doc: PyMutex<PyObjectRef>,
+    #[cfg(feature = "jit")]
+    jitted_code: OnceCell<CompiledCode>,
 }
 
 unsafe impl Traverse for PyFunction {
     fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
         self.globals.traverse(tracer_fn);
-        if let Some(closure) = self.closure.deref() {
-            closure.traverse(tracer_fn);
+        if let Some(closure) = self.closure.as_ref() {
+            closure.as_untyped().traverse(tracer_fn);
         }
         self.defaults_and_kwdefaults.traverse(tracer_fn);
     }
@@ -60,7 +59,7 @@ impl PyFunction {
     pub(crate) fn new(
         code: PyRef<PyCode>,
         globals: PyDictRef,
-        closure: Option<PyTupleRef>,
+        closure: Option<PyRef<PyTuple<PyCellRef>>>,
         defaults: Option<PyTupleRef>,
         kw_only_defaults: Option<PyDictRef>,
         qualname: PyStrRef,
@@ -84,7 +83,7 @@ impl PyFunction {
             code,
             globals,
             builtins,
-            closure: PyAtomicRef::from(closure),
+            closure,
             defaults_and_kwdefaults: PyMutex::new((defaults, kw_only_defaults)),
             name,
             qualname: PyMutex::new(qualname),
@@ -334,15 +333,15 @@ impl PyFunction {
                     attr_value.class().name()
                 ))
             })?;
-            self.set___defaults__(Some(defaults));
+            self.defaults_and_kwdefaults.lock().0 = Some(defaults);
         } else if attr.contains(bytecode::MakeFunctionFlags::KW_ONLY_DEFAULTS) {
             let kwdefaults = attr_value.clone().downcast::<PyDict>().map_err(|_| {
                 vm.new_type_error(format!(
                     "__kwdefaults__ must be a dict, not {}",
                     attr_value.class().name()
                 ))
             })?;
-            self.set___kwdefaults__(Some(kwdefaults));
+            self.defaults_and_kwdefaults.lock().1 = Some(kwdefaults);
         } else if attr.contains(bytecode::MakeFunctionFlags::ANNOTATIONS) {
             let annotations = attr_value.clone().downcast::<PyDict>().map_err(|_| {
                 vm.new_type_error(format!(
@@ -354,33 +353,18 @@ impl PyFunction {
         } else if attr.contains(bytecode::MakeFunctionFlags::CLOSURE) {
             // For closure, we need special handling
             // The closure tuple contains cell objects
-            let closure_tuple =
-                attr_value
-                    .clone()
-                    .downcast_exact::<PyTuple>(vm)
-                    .map_err(|obj| {
-                        vm.new_type_error(format!(
-                            "closure must be a tuple, not {}",
-                            obj.class().name()
-                        ))
-                    })?;
-
-            // Convert to tuple of cells
-            let cells: Result<Vec<_>, _> = closure_tuple
-                .iter()
-                .map(|cell| cell.clone().downcast_exact::<PyCell>(vm))
-                .collect();
-            let cells = cells
-                .map_err(|_| vm.new_type_error("closure must be a tuple of cells".to_owned()))?;
-
-            // Convert cells to PyTuple
-            let cells_objects: Vec<PyObjectRef> = cells
-                .into_iter()
-                .map(|cell| cell.into_pyref().into())
-                .collect();
-            let cells_tuple = PyTuple::new_ref(cells_objects, &vm.ctx);
-
-            let _ = unsafe { self.closure.swap(Some(cells_tuple)) };
+            let closure_tuple = attr_value
+                .clone()
+                .downcast_exact::<PyTuple>(vm)
+                .map_err(|obj| {
+                    vm.new_type_error(format!(
+                        "closure must be a tuple, not {}",
+                        obj.class().name()
+                    ))
+                })?
+                .into_pyref();
+
+            self.closure = Some(closure_tuple.try_into_typed::<PyCell>(vm)?);
         } else if attr.contains(bytecode::MakeFunctionFlags::TYPE_PARAMS) {
             let type_params = attr_value.clone().downcast::<PyTuple>().map_err(|_| {
                 vm.new_type_error(format!(
@@ -431,15 +415,7 @@ impl Py<PyFunction> {
             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(),
-                    )
-                }
-            }),
+            self.closure.as_ref().map_or(&[], |c| c.as_slice()),
             Some(self.to_owned().into()),
             vm,
         )
@@ -513,8 +489,7 @@ impl PyFunction {
     #[pymember]
     fn __closure__(vm: &VirtualMachine, zelf: PyObjectRef) -> PyResult {
         let zelf = Self::_as_pyref(&zelf, vm)?;
-        let closure = zelf.closure.deref().map(|x| x.as_object().to_owned());
-        Ok(vm.unwrap_or_none(closure))
+        Ok(vm.unwrap_or_none(zelf.closure.clone().map(|x| x.into())))
     }
 
     #[pymember]
@@ -698,9 +673,9 @@ impl Constructor for PyFunction {
                 )));
             }
 
-            // Validate that all items are cells
+            // Validate that all items are cells and create typed tuple
             let typed_closure = closure_tuple.try_into_typed::<PyCell>(vm)?;
-            Some(typed_closure.into_untyped())
+            Some(typed_closure)
         } else if !args.code.freevars.is_empty() {
             return Err(vm.new_type_error("arg 5 (closure) must be tuple"));
         } else {

diff --git a/vm/src/object/core.rs b/vm/src/object/core.rs
@@ -1023,12 +1023,12 @@ impl<T> Clone for PyRef<T> {
 }
 
 impl<T: PyObjectPayload> PyRef<T> {
-    #[inline(always)]
-    pub(crate) const fn into_non_null(self) -> NonNull<Py<T>> {
-        let ptr = self.ptr;
-        std::mem::forget(self);
-        ptr
-    }
+    // #[inline(always)]
+    // pub(crate) const fn into_non_null(self) -> NonNull<Py<T>> {
+    //     let ptr = self.ptr;
+    //     std::mem::forget(self);
+    //     ptr
+    // }
 
     #[inline(always)]
     pub(crate) const unsafe fn from_non_null(ptr: NonNull<Py<T>>) -> Self {