add more tests

Author: kumaraditya303

Lib/test/test_capi/test_opt.py

@@ -2378,6 +2378,24 @@ 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_store_subscr_dict_known_hash(self):
         # str, int, bytes, float, complex, tuple and any python object which has generic hash
         def testfunc(n):
@@ -2399,6 +2417,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

Python/bytecodes.c

@@ -1248,7 +1248,7 @@ dummy_func(
         tier2 op(_BINARY_OP_SUBSCR_DICT_KNOWN_HASH, (dict_st, sub_st, hash/4 -- res, ds, ss)) {
             PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
             PyObject *dict = PyStackRef_AsPyObjectBorrow(dict_st);
-
+            assert(Py_TYPE(dict)->tp_as_mapping->mp_subscript == _PyDict_Subscript);
             STAT_INC(BINARY_OP, hit);
             PyObject *res_o = _PyDict_SubscriptKnownHash(dict, sub, (Py_hash_t)hash);
             if (res_o == NULL) {
@@ -1401,8 +1401,7 @@ dummy_func(
 
         tier2 op(_STORE_SUBSCR_DICT_KNOWN_HASH, (value, dict_st, sub, hash/4 -- st)) {
             PyObject *dict = PyStackRef_AsPyObjectBorrow(dict_st);
-
-            assert(PyDict_CheckExact(dict));
+            assert(Py_TYPE(dict)->tp_as_mapping->mp_ass_subscript == _PyDict_StoreSubscript);
             STAT_INC(STORE_SUBSCR, hit);
             int err = _PyDict_SetItem_Take2_KnownHash((PyDictObject *)dict,
                                                 PyStackRef_AsPyObjectSteal(sub),