From 0fd074a002caab0c244452902d5c809a4972462b Mon Sep 17 00:00:00 2001
From: Malthe Borch <mborch@gmail.com>
Date: Fri, 4 Mar 2022 16:06:10 +0100
Subject: [PATCH] Replace key if not identical to old key

---
 Lib/test/test_dict.py    | 15 +++++++++++++++
 Lib/test/test_weakref.py | 13 +++++++++++++
 Objects/dictobject.c     | 10 ++++++++++
 3 files changed, 38 insertions(+)

diff --git a/Lib/test/test_dict.py b/Lib/test/test_dict.py
index e60ae4309cbc18..1b49b7e4f58fe0 100644
--- a/Lib/test/test_dict.py
+++ b/Lib/test/test_dict.py
@@ -618,6 +618,21 @@ def __hash__(self):
         with self.assertRaises(Exc):
             d1 == d2
 
+    def test_eq_replace(self):
+        class AlwaysEqualCmp(object):
+            def __eq__(self, other):
+                return True
+
+            def __hash__(self):
+                return 1
+
+        o1 = AlwaysEqualCmp()
+        d3 = {}
+        d3[o1] = 1
+        o2 = AlwaysEqualCmp()
+        d3[o2] = 2
+        assert tuple(d3.keys())[0] is o2
+
     def test_keys_contained(self):
         self.helper_keys_contained(lambda x: x.keys())
         self.helper_keys_contained(lambda x: x.items())
diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
index 46dac0228295d7..b246521c039f74 100644
--- a/Lib/test/test_weakref.py
+++ b/Lib/test/test_weakref.py
@@ -1797,6 +1797,19 @@ def test_weak_keyed_bad_delitem(self):
         self.assertRaises(TypeError, d.__getitem__,  13)
         self.assertRaises(TypeError, d.__setitem__,  13, 13)
 
+    def test_weak_keyed_equal_replacement(self):
+        d = weakref.WeakKeyDictionary()
+        o1 = Object('1')
+        o2 = Object('1')
+        d[o1] = 1
+        d[o2] = 2
+        assert len(d) == 1
+        assert d[o2] == 2
+        assert tuple(d.keys())[0] is o2
+        del o1
+        assert len(d) == 1
+        assert tuple(d.keys())[0] is o2
+
     def test_weak_keyed_cascading_deletes(self):
         # SF bug 742860.  For some reason, before 2.3 __delitem__ iterated
         # over the keys via self.data.iterkeys().  If things vanished from
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
index 635a738985c01d..f8c55a01c59015 100644
--- a/Objects/dictobject.c
+++ b/Objects/dictobject.c
@@ -1222,6 +1222,7 @@ Consumes key and value references.
 static int
 insertdict(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject *value)
 {
+    PyObject *old_key;
     PyObject *old_value;
 
     if (DK_IS_UNICODE(mp->ma_keys) && !PyUnicode_CheckExact(key)) {
@@ -1279,6 +1280,15 @@ insertdict(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject *value)
         return 0;
     }
 
+    if (!DK_IS_UNICODE(mp->ma_keys)) {
+        PyDictKeyEntry *ep = &DK_ENTRIES(mp->ma_keys)[ix];
+        old_key = ep->me_key;
+        if (old_key != key) {
+            ep->me_key = key;
+            key = old_key;
+        }
+    }
+
     if (old_value != value) {
         if (_PyDict_HasSplitTable(mp)) {
             mp->ma_values->values[ix] = value;