gh-143732: allow dict subclasses to be specialized

Include/internal/pycore_dict.h

@@ -44,6 +44,10 @@ extern int _PyDict_Next(
 
 extern int _PyDict_HasOnlyStringKeys(PyObject *mp);
 
+PyAPI_FUNC(PyObject *) _PyDict_Subscript(PyObject *self, PyObject *key);
+PyAPI_FUNC(PyObject *) _PyDict_SubscriptKnownHash(PyObject *self, PyObject *key, Py_hash_t hash);
+PyAPI_FUNC(int) _PyDict_StoreSubscript(PyObject *self, PyObject *key, PyObject *value);
+
 // Export for '_ctypes' shared extension
 PyAPI_FUNC(Py_ssize_t) _PyDict_SizeOf(PyDictObject *);
 

Lib/test/test_capi/test_opt.py

@@ -2273,10 +2273,49 @@ def f(n):
         self.assertEqual(res, TIER2_THRESHOLD)
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
-        self.assertEqual(uops.count("_GUARD_NOS_DICT"), 0)
-        self.assertEqual(uops.count("_STORE_SUBSCR_DICT_KNOWN_HASH"), 1)
+        self.assertEqual(uops.count("_GUARD_NOS_DICT_SUBSCRIPT"), 0)
+        self.assertEqual(uops.count("_GUARD_NOS_DICT_STORE_SUBSCRIPT"), 0)
         self.assertEqual(uops.count("_BINARY_OP_SUBSCR_DICT_KNOWN_HASH"), 1)
 
+    def test_dict_subclass_subscr(self):
+        import collections
+
+        def f(n):
+            x = 0
+            d = collections.defaultdict(int)
+            for _ in range(n):
+                d["key"] = 1
+                x += d["key"]
+            return x
+
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertEqual(uops.count("_BINARY_OP_SUBSCR_DICT_KNOWN_HASH"), 1)
+        self.assertEqual(uops.count("_STORE_SUBSCR_DICT_KNOWN_HASH"), 1)
+        self.assertEqual(uops.count("_GUARD_NOS_DICT_SUBSCRIPT"), 0)
+        self.assertEqual(uops.count("_GUARD_NOS_DICT_STORE_SUBSCRIPT"), 0)
+        self.assertEqual(uops.count("_GUARD_TYPE"), 1)
+
+    def test_dict_subclass_subscr_with_override(self):
+        class MyDict(dict):
+            def __getitem__(self, key):
+                return 42
+
+        def f(n):
+            d = MyDict()
+            x = 0
+            for _ in range(n):
+                x += d["anything"]
+            return x
+
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(res, 42 * TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertEqual(uops.count("_BINARY_OP_SUBSCR_INIT_CALL"), 1)
+
     def test_remove_guard_for_known_type_list(self):
         def f(n):
             x = 0
@@ -2599,6 +2638,23 @@ def testfunc(n):
         self.assertIn("_BINARY_OP_SUBSCR_DICT_KNOWN_HASH", uops)
         self.assertNotIn("_BINARY_OP_SUBSCR_DICT", uops)
 
+    def test_binary_op_subscr_defaultdict_known_hash(self):
+        # str, int, bytes, float, complex, tuple and any python object which has generic hash
+        import collections
+
+        def testfunc(n):
+            x = 0
+            d = collections.defaultdict(lambda: 1)
+            for _ in range(n):
+                x += d['a'] + d[1] + d[b'b'] + d[(1, 2)] + d[_GENERIC_KEY] + d[1.5] + d[1+2j]
+            return x
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, 7 * TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_BINARY_OP_SUBSCR_DICT_KNOWN_HASH", uops)
+        self.assertNotIn("_BINARY_OP_SUBSCR_DICT", uops)
 
     def test_binary_op_subscr_constant_frozendict_known_hash(self):
         def testfunc(n):
@@ -2635,6 +2691,28 @@ def testfunc(n):
         self.assertIn("_STORE_SUBSCR_DICT_KNOWN_HASH", uops)
         self.assertNotIn("_STORE_SUBSCR_DICT", uops)
 
+    def test_store_subscr_defaultdict_known_hash(self):
+        import collections
+
+        def testfunc(n):
+            d = collections.defaultdict(lambda: 0)
+            for _ in range(n):
+                d['a'] += 1
+                d[1] += 2
+                d[b'b'] += 3
+                d[(1, 2)] += 4
+                d[_GENERIC_KEY] += 5
+                d[1.5] += 6
+                d[1+2j] += 7
+            return d['a'] + d[1] + d[b'b'] + d[(1, 2)] + d[_GENERIC_KEY] + d[1.5] + d[1+2j]
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, 28 * TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_STORE_SUBSCR_DICT_KNOWN_HASH", uops)
+        self.assertNotIn("_STORE_SUBSCR_DICT", uops)
+
     def test_contains_op(self):
         def testfunc(n):
             x = 0

Lib/test/test_opcache.py

@@ -1,3 +1,4 @@
+import collections
 import copy
 import pickle
 import dis
@@ -1863,7 +1864,43 @@ class MyFrozenDict(frozendict):
                 self.assertEqual(a[2], 3)
 
         binary_subscr_frozen_dict_subclass()
-        self.assert_no_opcode(binary_subscr_frozen_dict_subclass, "BINARY_OP_SUBSCR_DICT")
+        self.assert_specialized(binary_subscr_frozen_dict_subclass, "BINARY_OP_SUBSCR_DICT")
+        self.assert_no_opcode(binary_subscr_frozen_dict_subclass, "BINARY_OP")
+
+        def binary_subscr_defaultdict():
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = collections.defaultdict(lambda: 42, {1: 2, 2: 3})
+                self.assertEqual(a[1], 2)
+                self.assertEqual(a[2], 3)
+                self.assertEqual(a[7], 42)
+
+        binary_subscr_defaultdict()
+        self.assert_specialized(binary_subscr_defaultdict, "BINARY_OP_SUBSCR_DICT")
+        self.assert_no_opcode(binary_subscr_defaultdict, "BINARY_OP")
+
+        def binary_subscr_counter():
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = collections.Counter('abcdeabcdabcaba')
+                self.assertEqual(a['a'], 5)
+                self.assertEqual(a['b'], 4)
+                self.assertEqual(a['m'], 0)
+
+        binary_subscr_counter()
+        self.assert_specialized(binary_subscr_counter, "BINARY_OP_SUBSCR_DICT")
+        self.assert_no_opcode(binary_subscr_counter, "BINARY_OP")
+
+        def binary_subscr_dict_subclass_override():
+            class MyDict(dict):
+                def __getitem__(self, key):
+                    return 42
+
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = MyDict()
+                self.assertEqual(a['a'], 42)
+                self.assertEqual(a['b'], 42)
+
+        binary_subscr_dict_subclass_override()
+        self.assert_no_opcode(binary_subscr_dict_subclass_override, "BINARY_OP_SUBSCR_DICT")
 
         def binary_subscr_str_int():
             for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
@@ -1924,6 +1961,29 @@ def store_subscr_frozen_dict():
         self.assert_specialized(store_subscr_frozen_dict, "STORE_SUBSCR_DICT")
         self.assert_no_opcode(store_subscr_frozen_dict, "STORE_SUBSCR")
 
+        def store_subscr_defaultdict():
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = collections.defaultdict(int)
+                a[1] = 4
+                self.assertEqual(a[1], 4)
+
+        store_subscr_defaultdict()
+        self.assert_specialized(store_subscr_defaultdict, "STORE_SUBSCR_DICT")
+        self.assert_no_opcode(store_subscr_defaultdict, "STORE_SUBSCR")
+
+        def store_subscr_dict_subclass_override():
+            class MyDict(dict):
+                def __setitem__(self, key, value):
+                    super().__setitem__(key, value * 2)
+
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                a = MyDict()
+                a['x'] = 5
+                self.assertEqual(a['x'], 10)
+
+        store_subscr_dict_subclass_override()
+        self.assert_no_opcode(store_subscr_dict_subclass_override, "STORE_SUBSCR_DICT")
+
     @cpython_only
     @requires_specialization
     def test_compare_op(self):