optimize floats and complex as well

Author: kumaraditya303

Lib/test/test_capi/test_opt.py

@@ -2321,35 +2321,37 @@ def testfunc(n):
         self.assertLessEqual(count_ops(ex, "_POP_TOP"), 2)
 
     def test_binary_op_subscr_dict_known_hash(self):
-        # str, int, bytes, tuple and any python object which has generic hash
+        # str, int, bytes, float, complex, tuple and any python object which has generic hash
         def testfunc(n):
             x = 0
-            d = {'a': 1, 1: 2, b'b': 3, (1, 2): 4, _GENERIC_KEY: 5}
+            d = {'a': 1, 1: 2, b'b': 3, (1, 2): 4, _GENERIC_KEY: 5, 1.5: 6, 1+2j: 7}
             for _ in range(n):
-                x += d['a'] + d[1] + d[b'b'] + d[(1, 2)] + d[_GENERIC_KEY]
+                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, 15 * TIER2_THRESHOLD)
+        self.assertEqual(res, 28 * 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, tuple and any python object which has generic hash
+        # str, int, bytes, float, complex, tuple and any python object which has generic hash
         def testfunc(n):
-            d = {'a': 0, 1: 0, b'b': 0, (1, 2): 0, _GENERIC_KEY: 0}
+            d = {'a': 0, 1: 0, b'b': 0, (1, 2): 0, _GENERIC_KEY: 0, 1.5: 0, 1+2j: 0}
             for _ in range(n):
                 d['a'] += 1
                 d[1] += 2
                 d[b'b'] += 3
                 d[(1, 2)] += 4
                 d[_GENERIC_KEY] += 5
-            return d['a'] + d[1] + d[b'b'] + d[(1, 2)] + d[_GENERIC_KEY]
+                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, 15 * TIER2_THRESHOLD)
+        self.assertEqual(res, 28 * TIER2_THRESHOLD)
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
         self.assertIn("_STORE_SUBSCR_DICT_KNOWN_HASH", uops)

Python/optimizer_bytecodes.c

@@ -167,7 +167,8 @@ dummy_func(void) {
     op(_STORE_SUBSCR_DICT, (value, dict_st, sub -- st)) {
         PyObject *sub_o = sym_get_const(ctx, sub);
         if (sub_o != NULL) {
-            if (PyUnicode_CheckExact(sub_o) || PyLong_CheckExact(sub_o) || PyBytes_CheckExact(sub_o)) {
+            if (PyUnicode_CheckExact(sub_o) || PyLong_CheckExact(sub_o) || PyBytes_CheckExact(sub_o)
+                || PyFloat_CheckExact(sub_o) || PyComplex_CheckExact(sub_o)) {
                 // PyObject_Hash can't fail on these types
                 ADD_OP(_STORE_SUBSCR_DICT_KNOWN_HASH, 0, PyObject_Hash(sub_o));
             }
@@ -506,7 +507,8 @@ dummy_func(void) {
     op(_BINARY_OP_SUBSCR_DICT, (dict_st, sub_st -- res, ds, ss)) {
         PyObject *sub = sym_get_const(ctx, sub_st);
         if (sub != NULL) {
-            if (PyUnicode_CheckExact(sub) || PyLong_CheckExact(sub) || PyBytes_CheckExact(sub)) {
+            if (PyUnicode_CheckExact(sub) || PyLong_CheckExact(sub) || PyBytes_CheckExact(sub)
+                    || PyFloat_CheckExact(sub) || PyComplex_CheckExact(sub)) {
                 // PyObject_Hash can't fail on these types
                 ADD_OP(_BINARY_OP_SUBSCR_DICT_KNOWN_HASH, 0, PyObject_Hash(sub));
             }