weakref subclass payload, and the pickler's stale pinned hooks

pyre/extra_tests/parity_tests/pickle_dispatch_table_across_collection.py

@@ -0,0 +1,71 @@
+# CPython-suite gap: pickle tests do not keep dump hooks live across collection.
+# parity-tests reason: this guards pyre/PyPy moving-GC roots during pickling.
+
+"""Pickler hooks and the effective dispatch table survive collection."""
+
+import copyreg
+import gc
+import io
+import pickle
+
+
+class Payload:
+    def __init__(self, value):
+        self.value = value
+
+
+def rebuild_payload(value):
+    return Payload(value)
+
+
+def reduce_payload(obj):
+    return rebuild_payload, (obj.value,)
+
+
+class CollectingPickler(pickle.Pickler):
+    def __init__(self, file, protocol):
+        super().__init__(file, protocol=protocol)
+        self.churned = False
+        self.persistent_before_churn = 0
+        self.persistent_after_churn = 0
+        self.override_after_churn = 0
+
+    def persistent_id(self, obj):
+        if self.churned:
+            self.persistent_after_churn += 1
+        else:
+            self.persistent_before_churn += 1
+            junk = [[object() for _ in range(64)] for _ in range(128)]
+            assert len(junk) == 128
+            gc.collect()
+            self.churned = True
+        return None
+
+    def reducer_override(self, obj):
+        assert self.churned
+        self.override_after_churn += 1
+        return NotImplemented
+
+
+original_dispatch_table = copyreg.dispatch_table
+try:
+    copyreg.dispatch_table = {Payload: reduce_payload}
+    expected = [Payload(i) for i in range(40)]
+    expected.extend([{"padding": [str(i) for i in range(80)]}, Payload(999)])
+
+    for protocol in (0, 2, pickle.HIGHEST_PROTOCOL):
+        stream = io.BytesIO()
+        pickler = CollectingPickler(stream, protocol)
+        pickler.dump(expected)
+        restored = pickle.loads(stream.getvalue())
+
+        assert [obj.value for obj in restored[:40]] == list(range(40))
+        assert restored[-1].value == 999
+        assert restored[-2] == expected[-2]
+        assert pickler.persistent_before_churn == 1
+        assert pickler.persistent_after_churn > 0
+        assert pickler.override_after_churn >= 41
+finally:
+    copyreg.dispatch_table = original_dispatch_table
+
+print("OK")

pyre/extra_tests/parity_tests/weakref_ref_subclass_layout.py

@@ -0,0 +1,102 @@
+# CPython-suite gap: weakref subclass tests do not combine callbacks with all storage shapes.
+# parity-tests reason: this guards the translated W_Weakref payload and subclass carrier layout.
+
+import gc
+from weakref import WeakValueDictionary, ref
+
+
+class Value:
+    def __init__(self, value):
+        self.value = value
+
+    def __eq__(self, other):
+        return isinstance(other, Value) and self.value == other.value
+
+    def __hash__(self):
+        return hash(self.value)
+
+
+class PlainRef(ref):
+    pass
+
+
+class EmptyRef(ref):
+    __slots__ = ()
+
+
+class KeyedRef(ref):
+    __slots__ = ("key",)
+
+
+PRIVATE_FIELDS = ("w_obj_weak", "w_callable", "w_hash", "w_slots")
+
+
+def collect():
+    for _ in range(3):
+        gc.collect()
+
+
+def check_ref_subclass(ref_type):
+    callbacks = []
+
+    def callback(w_ref):
+        callbacks.append(w_ref() is None)
+
+    value = Value((ref_type.__name__, 42))
+    equal_value = Value((ref_type.__name__, 42))
+    w_ref = ref_type(value, callback)
+    equal_ref = ref(equal_value)
+
+    assert w_ref() is value
+    assert w_ref.__callback__ is callback
+    assert hash(w_ref) == hash(value)
+    assert w_ref == equal_ref
+
+    if ref_type is PlainRef:
+        w_ref.marker = "plain"
+        assert w_ref.__dict__ == {"marker": "plain"}
+    else:
+        assert not hasattr(w_ref, "__dict__")
+
+    if ref_type is KeyedRef:
+        w_ref.key = ("module", "name")
+        assert w_ref.key == ("module", "name")
+        w_ref.key = ("module", "renamed")
+        assert w_ref.key == ("module", "renamed")
+        assert w_ref() is value
+        del w_ref.key
+        try:
+            w_ref.key
+        except AttributeError:
+            pass
+        else:
+            raise AssertionError("deleted slot remained readable")
+
+    for name in PRIVATE_FIELDS:
+        assert not hasattr(w_ref, name)
+        if hasattr(w_ref, "__dict__"):
+            assert name not in w_ref.__dict__
+
+    del value
+    collect()
+    assert w_ref() is None
+    assert callbacks == [True]
+
+
+for ref_type in (PlainRef, EmptyRef, KeyedRef):
+    check_ref_subclass(ref_type)
+
+
+values = WeakValueDictionary()
+key = ("namespace", "spam")
+value = Value("live")
+values[key] = value
+assert values[key] is value
+stored_ref = values.data[key]
+assert stored_ref.key == key
+assert stored_ref() is value
+del value
+collect()
+assert key not in values
+
+print("OK")

pyre/pyre-interpreter/src/baseobjspace.rs

@@ -4725,6 +4725,9 @@ pub(crate) fn native_slot_get(
     if unsafe { pyre_object::is_list(obj) } {
         return Ok(unsafe { pyre_object::listobject::w_list_slot_get(obj, index as usize) });
     }
+    if unsafe { pyre_object::weakref::is_typed_weakref(obj) } {
+        return Ok(unsafe { pyre_object::weakref::w_weakref_object_slot_get(obj, index as usize) });
+    }
     let w_dict = getdict(obj)?;
     if w_dict.is_null() {
         return Ok(None);
@@ -4766,6 +4769,10 @@ pub(crate) fn native_slot_set(
         unsafe { pyre_object::listobject::w_list_slot_set(obj, index as usize, value) };
         return Ok(true);
     }
+    if unsafe { pyre_object::weakref::is_typed_weakref(obj) } {
+        unsafe { pyre_object::weakref::w_weakref_object_slot_set(obj, index as usize, value) };
+        return Ok(true);
+    }
     let w_dict = getdict(obj)?;
     if w_dict.is_null() {
         return Ok(false);
@@ -4798,6 +4805,9 @@ pub(crate) fn native_slot_del(obj: PyObjectRef, name: &str, index: u32) -> Resul
     if unsafe { pyre_object::is_list(obj) } {
         return Ok(unsafe { pyre_object::listobject::w_list_slot_del(obj, index as usize) });
     }
+    if unsafe { pyre_object::weakref::is_typed_weakref(obj) } {
+        return Ok(unsafe { pyre_object::weakref::w_weakref_object_slot_del(obj, index as usize) });
+    }
     let w_dict = getdict(obj)?;
     if w_dict.is_null() {
         return Ok(false);

pyre/pyre-interpreter/src/module/_pickle/pickler.rs

@@ -30,6 +30,35 @@ fn cur_pickler(slot: usize) -> &'static mut W_Pickler {
     unsafe { &mut *(pyre_object::gc_roots::shadow_stack_get(slot) as *mut W_Pickler) }
 }
 
+/// A reference kept for the duration of `pickle_core_impl`. Movable values
+/// live in a shadow-stack slot and are re-read on every access so relocation
+/// is observed; `PY_NULL` and immovable singletons can be held verbatim.
+struct PinnedRef {
+    slot: Option<usize>,
+    immovable: PyObjectRef,
+}
+
+impl PinnedRef {
+    fn new(value: PyObjectRef, movable: bool) -> Self {
+        let slot = if movable {
+            pyre_object::gc_roots::pin_root(value);
+            Some(pyre_object::gc_roots::shadow_stack_len() - 1)
+        } else {
+            None
+        };
+        Self {
+            slot,
+            immovable: value,
+        }
+    }
+
+    fn get(&self) -> PyObjectRef {
+        self.slot.map_or(self.immovable, |slot| {
+            pyre_object::gc_roots::shadow_stack_get(slot)
+        })
+    }
+}
+
 /// Old-to-young / incremental-mark write barrier for a store into the
 /// pickler's own GC-pointer fields (`w_memo` / `w_dispatch_table` /
 /// `w_pers_func` / …). The GC transform emits `ll_writebarrier` after every
@@ -97,9 +126,9 @@ struct PickleCtx {
     index: HashMap<usize, Vec<usize>>,
     /// `persistent_id` callable resolved off the pickler (subclass override
     /// or set attribute), or `PY_NULL` when not defined.
-    pers_func: PyObjectRef,
+    pers_func: PinnedRef,
     /// `buffer_callback` for proto-5 out-of-band buffers, or `None`/`PY_NULL`.
-    buffer_callback: PyObjectRef,
+    buffer_callback: PinnedRef,
     /// `fast` mode — when set, memoization is skipped (no PUT/GET); a
     /// shallow cyclic-object guard (`fast_nesting` / `fast_memo`) still fires
     /// past `FAST_NESTING_LIMIT`.
@@ -116,10 +145,10 @@ struct PickleCtx {
     /// Effective `dispatch_table` (the pickler's, else `copyreg.dispatch_table`)
     /// consulted by `type` for the reduce of an otherwise-unhandled object;
     /// `None`/`PY_NULL` when unavailable.
-    dispatch_table: PyObjectRef,
+    dispatch_table: PinnedRef,
     /// `reducer_override` callable (a subclass hook) consulted for every
     /// object, or `PY_NULL` when not defined.
-    reducer_override: PyObjectRef,
+    reducer_override: PinnedRef,
 }
 
 impl PickleCtx {
@@ -1026,18 +1055,16 @@ fn pickle_core_impl(
 ) -> Result<PyObjectRef, PyError> {
     let _roots = pyre_object::gc_roots::push_roots();
     pyre_object::gc_roots::pin_root(w_obj);
-    if !pers_func.is_null() {
-        pyre_object::gc_roots::pin_root(pers_func);
-    }
-    if !buffer_callback.is_null() && !unsafe { pyre_object::is_none(buffer_callback) } {
-        pyre_object::gc_roots::pin_root(buffer_callback);
-    }
-    if !dispatch_table.is_null() && !unsafe { pyre_object::is_none(dispatch_table) } {
-        pyre_object::gc_roots::pin_root(dispatch_table);
-    }
-    if !reducer_override.is_null() {
-        pyre_object::gc_roots::pin_root(reducer_override);
-    }
+    let pers_func = PinnedRef::new(pers_func, !pers_func.is_null());
+    let buffer_callback = PinnedRef::new(
+        buffer_callback,
+        !buffer_callback.is_null() && !unsafe { pyre_object::is_none(buffer_callback) },
+    );
+    let dispatch_table = PinnedRef::new(
+        dispatch_table,
+        !dispatch_table.is_null() && !unsafe { pyre_object::is_none(dispatch_table) },
+    );
+    let reducer_override = PinnedRef::new(reducer_override, !reducer_override.is_null());
     // Pin the memo list and index its existing entries (a reused `Pickler`
     // carries memo state across `dump` calls until `clear_memo`).
     pyre_object::gc_roots::pin_root(w_memo);
@@ -1179,11 +1206,9 @@ fn save(ctx: &mut PickleCtx, buf: &mut Framer, w_obj: PyObjectRef) -> Result<(),
     pyre_object::gc_roots::pin_root(w_obj);
     let slot = pyre_object::gc_roots::shadow_stack_len() - 1;
     buf.commit_frame(false)?;
-    if !ctx.pers_func.is_null() {
-        let w_pid = call_fn(
-            ctx.pers_func,
-            &[pyre_object::gc_roots::shadow_stack_get(slot)],
-        )?;
+    let pers_func = ctx.pers_func.get();
+    if !pers_func.is_null() {
+        let w_pid = call_fn(pers_func, &[pyre_object::gc_roots::shadow_stack_get(slot)])?;
         if !unsafe { pyre_object::is_none(w_pid) } {
             save_pers(ctx, buf, w_pid)
         } else {
@@ -1306,12 +1331,13 @@ fn dispatch_save(ctx: &mut PickleCtx, buf: &mut Framer, w_obj: PyObjectRef) -> R
     // default reduction. (interp_pickle.py:619-625 calls it earlier; CPython
     // 3.14 — the behaviour target — dispatches built-in types first, so the
     // hook never sees a list/dict/str and a repeated object hits the memo.)
-    if !ctx.reducer_override.is_null() {
+    let reducer_override = ctx.reducer_override.get();
+    if !reducer_override.is_null() {
         let _roots = pyre_object::gc_roots::push_roots();
         pyre_object::gc_roots::pin_root(w_obj);
         let slot = pyre_object::gc_roots::shadow_stack_len() - 1;
         let w_rv = call_fn(
-            ctx.reducer_override,
+            reducer_override,
             &[pyre_object::gc_roots::shadow_stack_get(slot)],
         )?;
         if !unsafe { pyre_object::is_not_implemented(w_rv) } {
@@ -1482,7 +1508,7 @@ fn dispatch_table_reduce(
     ctx: &PickleCtx,
     w_obj: PyObjectRef,
 ) -> Result<Option<PyObjectRef>, PyError> {
-    let dt = ctx.dispatch_table;
+    let dt = ctx.dispatch_table.get();
     if dt.is_null() {
         return Ok(None);
     }
@@ -2056,8 +2082,9 @@ fn save_picklebuffer(
     }
     let (data, readonly) = crate::module::__pypy__::interp_buffer::buffer_view(wrapped)?;
     let mut in_band = true;
-    if !unsafe { pyre_object::is_none(ctx.buffer_callback) } {
-        let w_ret = call_fn(ctx.buffer_callback, &[w_obj])?;
+    let buffer_callback = ctx.buffer_callback.get();
+    if !unsafe { pyre_object::is_none(buffer_callback) } {
+        let w_ret = call_fn(buffer_callback, &[w_obj])?;
         in_band = crate::baseobjspace::is_true(w_ret)?;
     }
     if in_band {