Fix shared reference data loss: update memo after BUILD

When the same non-Persistent Python object is referenced in two places
in a pickle, BINPUT clones the stack top into memo before BUILD
transforms it (e.g. Reduce → Instance). BINGET then returns the stale
Reduce clone, losing instance state on the second occurrence.

After BUILD, scan memo for entries matching the old pre-BUILD value
and replace them with the post-BUILD result. This mirrors CPython's
pickle VM which uses object identity (shared references).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: maethu

CHANGES.md

@@ -2,7 +2,7 @@
 
 ## 1.2.1 (unreleased)
 
-- nothing yet
+- Fix shared reference data loss: update memo after BUILD
 
 ## 1.2.0 (2026-02-10)
 

src/decode.rs

@@ -414,6 +414,13 @@ impl<'a> Decoder<'a> {
                 BUILD => {
                     let state = self.pop_value()?;
                     let obj = self.pop_value()?;
+                    // Save pre-BUILD value so we can update stale memo entries.
+                    // BINPUT clones the stack top into memo *before* BUILD runs,
+                    // so memo entries still reference the old (e.g. Reduce) value.
+                    // After BUILD transforms it (e.g. to Instance), we must
+                    // propagate the change to memo — mirroring how CPython's
+                    // pickle VM uses object identity (shared references).
+                    let old_obj = obj.clone();
                     match obj {
                         PickleValue::Global { module, name } => {
                             self.push(PickleValue::Instance {
@@ -495,24 +502,24 @@ impl<'a> Decoder<'a> {
                             });
                         }
                     }
+                    // Update memo: replace stale pre-BUILD entries with the
+                    // new post-BUILD value so BINGET returns the correct form.
+                    let new_val = self.peek_value()?.clone();
+                    if old_obj != new_val {
+                        for entry in self.memo.iter_mut() {
+                            if *entry == old_obj {
+                                *entry = new_val.clone();
+                            }
+                        }
+                    }
                 }
                 NEWOBJ => {
                     let args = self.pop_value()?;
                     let cls = self.pop_value()?;
-                    match cls {
-                        PickleValue::Global { module, name } => {
-                            self.push(PickleValue::Reduce {
-                                callable: Box::new(PickleValue::Global { module, name }),
-                                args: Box::new(args),
-                            });
-                        }
-                        _ => {
-                            self.push(PickleValue::Reduce {
-                                callable: Box::new(cls),
-                                args: Box::new(args),
-                            });
-                        }
-                    }
+                    self.push(PickleValue::Reduce {
+                        callable: Box::new(cls),
+                        args: Box::new(args),
+                    });
                 }
                 NEWOBJ_EX => {
                     let kwargs = self.pop_value()?;
@@ -809,4 +816,69 @@ mod tests {
             )])
         );
     }
+
+    #[test]
+    fn test_memo_updated_after_build() {
+        // Reproduces the shared-reference bug: NEWOBJ + BINPUT creates a
+        // Reduce in memo, BUILD transforms the stack top to Instance, but
+        // memo was stale. After the fix, BINGET should return Instance.
+        //
+        // Pickle VM sequence:
+        //   PROTO 3
+        //   GLOBAL "mymod\nMyClass\n"
+        //   EMPTY_TUPLE
+        //   NEWOBJ           -> Reduce{Global{mymod,MyClass}, ()}
+        //   BINPUT 0         -> memo[0] = clone of Reduce
+        //   EMPTY_DICT
+        //   SHORT_BINUNICODE "name"
+        //   SHORT_BINUNICODE "test"
+        //   SETITEM          -> {"name": "test"}
+        //   BUILD            -> Instance{mymod, MyClass, {"name": "test"}}
+        //   BINGET 0         -> should be Instance (was stale Reduce before fix)
+        //   TUPLE2           -> (Instance_from_build, memo_copy)
+        //   STOP
+        let data: &[u8] = &[
+            0x80, 0x03,                                     // PROTO 3
+            b'c', b'm', b'y', b'm', b'o', b'd', b'\n',     // GLOBAL module="mymod"
+            b'M', b'y', b'C', b'l', b's', b'\n',            // GLOBAL name="MyCls"
+            b')',                                            // EMPTY_TUPLE
+            0x81,                                            // NEWOBJ
+            b'q', 0x00,                                      // BINPUT 0
+            b'}',                                            // EMPTY_DICT
+            0x8c, 0x04, b'n', b'a', b'm', b'e',             // SHORT_BINUNICODE "name"
+            0x8c, 0x04, b't', b'e', b's', b't',             // SHORT_BINUNICODE "test"
+            b's',                                            // SETITEM
+            b'b',                                            // BUILD
+            b'h', 0x00,                                      // BINGET 0
+            0x86,                                            // TUPLE2
+            b'.',                                            // STOP
+        ];
+        let result = decode_pickle(data).unwrap();
+
+        // Result should be a tuple of two elements
+        if let PickleValue::Tuple(items) = &result {
+            assert_eq!(items.len(), 2, "expected 2-tuple");
+
+            let expected = PickleValue::Instance {
+                module: "mymod".to_string(),
+                name: "MyCls".to_string(),
+                state: Box::new(PickleValue::Dict(vec![(
+                    PickleValue::String("name".to_string()),
+                    PickleValue::String("test".to_string()),
+                )])),
+            };
+
+            // First element: the Instance from BUILD (stack top)
+            assert_eq!(items[0], expected, "BUILD result should be Instance");
+
+            // Second element: the memo copy from BINGET — must also be
+            // Instance after the fix (was stale Reduce before)
+            assert_eq!(
+                items[1], expected,
+                "BINGET after BUILD should return updated Instance, not stale Reduce"
+            );
+        } else {
+            panic!("expected Tuple, got {:?}", result);
+        }
+    }
 }

tests/test_zodb_records.py

@@ -301,6 +301,76 @@ def test_datetime_in_state(self, zodb):
         assert result == result2
 
 
+class SharedRefHelper:
+    """Non-Persistent helper class for testing shared-reference pickling.
+
+    When the same instance appears in two dict values, pickle uses a memo
+    backreference (BINGET) for the second occurrence. The codec must preserve
+    state for both occurrences.
+    """
+
+    def __init__(self, name="<undefined>"):
+        self.name = name
+
+
+class TestSharedReferences:
+    """Test that objects referenced multiple times via pickle memo preserve state."""
+
+    def test_shared_object_in_dict_values(self):
+        """Same object in two dict values: both should have full state."""
+        obj = SharedRefHelper("acl_users")
+        state = {"first": obj, "second": obj}
+        data = pickle.dumps(state, protocol=3)
+
+        result = zodb_json_codec.pickle_to_dict(data)
+
+        first = result["first"]
+        second = result["second"]
+
+        assert "@cls" in first, (
+            "first occurrence should be Instance (@cls), got: %s" % first
+        )
+        assert "@cls" in second, (
+            "second occurrence should be Instance (@cls), not @reduce. "
+            "This fails when memo is not updated after BUILD. Got: %s" % second
+        )
+        assert first["@s"]["name"] == "acl_users"
+        assert second["@s"]["name"] == "acl_users"
+
+    def test_shared_object_roundtrip(self):
+        """Shared reference round-trips through JSON correctly."""
+        obj = SharedRefHelper("test_hook")
+        state = {"a": obj, "b": obj}
+        data = pickle.dumps(state, protocol=3)
+
+        json_str = zodb_json_codec.pickle_to_json(data)
+        parsed = json.loads(json_str)
+
+        assert parsed["a"]["@cls"] == [__name__, "SharedRefHelper"]
+        assert parsed["b"]["@cls"] == [__name__, "SharedRefHelper"]
+        assert parsed["a"]["@s"]["name"] == "test_hook"
+        assert parsed["b"]["@s"]["name"] == "test_hook"
+
+    def test_shared_object_in_zodb_record(self):
+        """Shared references inside a ZODB record state are preserved."""
+        obj = SharedRefHelper("hook_name")
+        record = make_zodb_record(
+            "myapp", "Container",
+            {"hook_a": obj, "hook_b": obj},
+        )
+        result = zodb_json_codec.decode_zodb_record(record)
+
+        hook_a = result["@s"]["hook_a"]
+        hook_b = result["@s"]["hook_b"]
+
+        assert "@cls" in hook_a, "hook_a should be Instance"
+        assert "@cls" in hook_b, (
+            "hook_b (shared ref) should be Instance, not @reduce"
+        )
+        assert hook_a["@s"]["name"] == "hook_name"
+        assert hook_b["@s"]["name"] == "hook_name"
+
+
 class TestDecodeZodbRecordForPg:
     """Test decode_zodb_record_for_pg: single-pass decode + refs + null sanitize."""