pyre: let a raised dict probe fail to match by identity too

pyre/extra_tests/parity_tests/module_dict_invariants.py

@@ -61,4 +61,57 @@
 del g["_ccc"]
 
 
+# (6) A store that raises mid-probe leaves the module dict untouched.
+#
+# `ll_dict_lookup` propagates at the first raising comparison
+# (`rordereddict.py:1055`), so nothing later in the bucket is found and the
+# store never lands.  Reproduce the shape that stresses it: one key whose
+# `__eq__` raises and, in the SAME bucket, the incoming key's own object —
+# so the raise is seen before the entry that would otherwise match.
+_armed = False
+
+
+class _Raiser:
+    def __hash__(self):
+        return 7
+
+    def __eq__(self, other):
+        if _armed:
+            raise ValueError("boom")
+        return self is other
+
+
+class _Quiet:
+    def __hash__(self):
+        return 7
+
+    def __eq__(self, other):
+        return self is other
+
+
+def _store_must_not_mutate_on_raise():
+    global _armed
+    quiet = _Quiet()
+    g[_Raiser()] = "raiser"
+    g[quiet] = "quiet-old"
+    before = list(g.items())
+    _armed = True
+    try:
+        g[quiet] = "quiet-new"
+    except ValueError:
+        pass
+    else:
+        assert False, "raising __eq__ must propagate out of the store"
+    finally:
+        _armed = False
+    live = {id(k) for k, _ in g.items()}
+    dropped = [k for k, _ in before if id(k) not in live]
+    assert not dropped, f"store dropped unrelated keys: {dropped!r}"
+    assert len(g) == len(before), f"len {len(before)} -> {len(g)}"
+    assert g[quiet] == "quiet-old", f"value applied despite raise: {g[quiet]!r}"
+
+
+_store_must_not_mutate_on_raise()
+
+
 print("OK")

pyre/pyre-object/src/dictmultiobject.rs

@@ -1920,22 +1920,31 @@ unsafe fn key_equality_is_builtin(key: PyObjectRef) -> bool {
 /// to the limited-type builtin equality below — sufficient for the
 /// hashable-builtin smoke tests but not for arbitrary user types.
 pub(crate) unsafe fn dict_keys_equal(a: PyObjectRef, b: PyObjectRef) -> bool {
+    // Once `space.eq_w` has raised earlier in this probe, nothing may match —
+    // not by user comparison and not by identity.  The Rust `Eq` callback
+    // cannot abort the `IndexMap` scan, so answering `false` for the rest of it
+    // is how a raised probe is made to look like the `r_dict(space.eq_w,
+    // space.hash_w)` lookup it stands for: `ll_dict_lookup` propagates at the
+    // first comparison, so no later key — however it would have compared — can
+    // still be found.  The flag is cleared per op in
+    // `object_key_for(_checked)`, so this only fires after a raise within the
+    // current probe.
+    //
+    // This has to sit above the identity shortcut, not below it.  A bucket can
+    // hold both a key whose `__eq__` raises and the incoming key's own object;
+    // if the raising one is probed first and identity still answered `true`
+    // afterwards, `insert` would REPLACE rather than append, and the checked
+    // store's `dict_entries_pop_last` undo would then drop an unrelated last
+    // entry while leaving the replaced value overwritten.
+    if crate::dict_eq_hook::eq_error_pending() {
+        return false;
+    }
     if std::ptr::eq(a, b) {
         return true;
     }
     if a.is_null() || b.is_null() {
         return false;
     }
-    // Once `space.eq_w` has raised earlier in this probe, skip every
-    // remaining user comparison.  The Rust `Eq` callback cannot abort the
-    // `IndexMap` scan, but suppressing further `__eq__` calls means no extra
-    // comparison runs and the first exception is the one that propagates —
-    // matching `r_dict(space.eq_w, space.hash_w)` raising at the first
-    // comparison.  The flag is cleared per op in `object_key_for(_checked)`,
-    // so this only fires after a raise within the current probe.
-    if crate::dict_eq_hook::eq_error_pending() {
-        return false;
-    }
     if crate::dict_eq_hook::callback_free_probe_active() {
         // A callback-free probe promises that no user code runs while the
         // container is being read, so the ladder below must answer on its