vm: add GC traverse/clear for specialization_getitem and revalidate version under lock

youknowone · youknowone · commit e97dab9ece70 · 2026-03-08T18:08:59.000+09:00
diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs
@@ -235,6 +235,7 @@ unsafe impl crate::object::Traverse for PyType {
             .count();
         if let Some(ext) = self.heaptype_ext.as_ref() {
             ext.specialization_init.read().traverse(tracer_fn);
+            ext.specialization_getitem.read().traverse(tracer_fn);
         }
     }
 
@@ -263,11 +264,19 @@ unsafe impl crate::object::Traverse for PyType {
                 out.push(val);
             }
         }
-        if let Some(ext) = self.heaptype_ext.as_ref()
-            && let Some(mut guard) = ext.specialization_init.try_write()
-            && let Some(init) = guard.take()
-        {
-            out.push(init.into());
+        if let Some(ext) = self.heaptype_ext.as_ref() {
+            if let Some(mut guard) = ext.specialization_init.try_write()
+                && let Some(init) = guard.take()
+            {
+                out.push(init.into());
+            }
+            if let Some(mut guard) = ext.specialization_getitem.try_write()
+                && let Some(getitem) = guard.take()
+            {
+                out.push(getitem.into());
+                ext.specialization_getitem_version
+                    .store(0, Ordering::Release);
+            }
         }
     }
 }
@@ -860,14 +869,18 @@ impl PyType {
         let Some(ext) = self.heaptype_ext.as_ref() else {
             return false;
         };
-        if tp_version == 0 || self.tp_version_tag.load(Ordering::Acquire) != tp_version {
+        if tp_version == 0 {
             return false;
         }
         let func_version = getitem.get_version_for_current_state();
         if func_version == 0 {
             return false;
         }
-        *ext.specialization_getitem.write() = Some(getitem);
+        let mut guard = ext.specialization_getitem.write();
+        if self.tp_version_tag.load(Ordering::Acquire) != tp_version {
+            return false;
+        }
+        *guard = Some(getitem);
         ext.specialization_getitem_version
             .store(func_version, Ordering::Release);
         true
@@ -876,15 +889,15 @@ impl PyType {
     /// Read cached __getitem__ for BINARY_OP_SUBSCR_GETITEM specialization.
     pub(crate) fn get_cached_getitem_for_specialization(&self) -> Option<(PyRef<PyFunction>, u32)> {
         let ext = self.heaptype_ext.as_ref()?;
-        if self.tp_version_tag.load(Ordering::Acquire) == 0 {
-            return None;
-        }
         let cached_version = ext.specialization_getitem_version.load(Ordering::Acquire);
         if cached_version == 0 {
             return None;
         }
-        ext.specialization_getitem
-            .read()
+        let guard = ext.specialization_getitem.read();
+        if self.tp_version_tag.load(Ordering::Acquire) == 0 {
+            return None;
+        }
+        guard
             .as_ref()
             .map(|getitem| (getitem.to_owned(), cached_version))
     }
