JIT constant fold special method lookup

Author: kumaraditya303

Include/internal/pycore_opcode_metadata.h

@@ -3316,6 +3316,25 @@ def f(n):
         self.assertNotIn("_LOAD_ATTR_METHOD_NO_DICT", uops)
         self.assertIn("_LOAD_CONST_INLINE_BORROW", uops)
 
+    def test_cached_load_special(self):
+        class CM:
+            def __enter__(self):
+                return self
+            def __exit__(self, *args):
+                pass
+        def f(n):
+            cm = CM()
+            x = 0
+            for _ in range(n):
+                with cm:
+                    x += 1
+            return x
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        uops = get_opnames(ex)
+        self.assertNotIn("_LOAD_SPECIAL", uops)
+
     def test_store_fast_refcount_elimination(self):
         def foo(x):
             # Since x is known to be

Lib/test/test_capi/test_opt.py

@@ -3887,6 +3887,7 @@ dummy_func(
         }
 
         macro(LOAD_SPECIAL) =
+            _RECORD_TOS_TYPE +
             _INSERT_NULL +
             _LOAD_SPECIAL;
 

Python/bytecodes.c

@@ -1682,8 +1682,29 @@ dummy_func(void) {
     }
 
     op(_LOAD_SPECIAL, (method_and_self[2] -- method_and_self[2])) {
-        method_and_self[0] = sym_new_not_null(ctx);
-        method_and_self[1] = sym_new_unknown(ctx);
+        PyTypeObject *type = sym_get_probable_type(method_and_self[1]);
+        if (type != NULL) {
+            PyObject *name = _Py_SpecialMethods[oparg].name;
+            PyObject *descr = _PyType_Lookup(type, name);
+            if (descr != NULL && _PyType_HasFeature(Py_TYPE(descr), Py_TPFLAGS_METHOD_DESCRIPTOR))
+            {
+                ADD_OP(_GUARD_TYPE_VERSION, 0, type->tp_version_tag);
+                bool immortal = _Py_IsImmortal(descr) || (type->tp_flags & Py_TPFLAGS_IMMUTABLETYPE);
+                ADD_OP(immortal ? _LOAD_CONST_INLINE_BORROW : _LOAD_CONST_INLINE,
+                       0, (uintptr_t)descr);
+                ADD_OP(_SWAP, 3, 0);
+                ADD_OP(_POP_TOP, 0, 0);
+                if ((type->tp_flags & Py_TPFLAGS_IMMUTABLETYPE) == 0) {
+                    PyType_Watch(TYPE_WATCHER_ID, (PyObject *)type);
+                    _Py_BloomFilter_Add(dependencies, type);
+                }
+                method_and_self[0] = sym_new_const(ctx, descr);
+            }
+        }
+        else {
+            method_and_self[0] = sym_new_not_null(ctx);
+            method_and_self[1] = sym_new_unknown(ctx);
+        }
     }
 
     op(_JUMP_TO_TOP, (--)) {