gh-131798: JIT optimizer: Support custom binary op and property frames (GH-143735)

Author: Fidget-Spinner

Lib/test/test_capi/test_opt.py

@@ -3575,6 +3575,42 @@ def testfunc(n):
         # _POP_TOP_NOP is a sign the optimizer ran and didn't hit bottom.
         self.assertGreaterEqual(count_ops(ex, "_POP_TOP_NOP"), 1)
 
+    def test_binary_op_subscr_init_frame(self):
+        class B:
+            def __getitem__(self, other):
+                return other + 1
+        def testfunc(*args):
+            n, b = args[0]
+            for _ in range(n):
+                y = b[2]
+
+        res, ex = self._run_with_optimizer(testfunc, (TIER2_THRESHOLD, B()))
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_BINARY_OP_SUBSCR_INIT_CALL", uops)
+        # _POP_TOP_NOP is a sign the optimizer ran and didn't hit contradiction.
+        self.assertGreaterEqual(count_ops(ex, "_POP_TOP_NOP"), 1)
+
+    def test_load_attr_property_frame(self):
+        class B:
+            @property
+            def prop(self):
+                return 3
+        def testfunc(*args):
+            n, b = args[0]
+            for _ in range(n):
+                y = b.prop + b.prop
+
+        testfunc((3, B()))
+        res, ex = self._run_with_optimizer(testfunc, (TIER2_THRESHOLD, B()))
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_LOAD_ATTR_PROPERTY_FRAME", uops)
+        # This is a sign the optimizer ran and didn't hit contradiction.
+        self.assertIn("_INSERT_2_LOAD_CONST_INLINE_BORROW", uops)
+
     def test_unary_negative(self):
         def testfunc(n):
             a = 3

Python/optimizer_bytecodes.c

@@ -327,9 +327,20 @@ dummy_func(void) {
         GETLOCAL(this_instr->operand0) = sym_new_null(ctx);
     }
 
-    op(_BINARY_OP_SUBSCR_INIT_CALL, (container, sub, getitem  -- new_frame)) {
-        new_frame = PyJitRef_NULL;
-        ctx->done = true;
+    op(_BINARY_OP_SUBSCR_INIT_CALL, (container, sub, getitem -- new_frame)) {
+        assert((this_instr + 1)->opcode == _PUSH_FRAME);
+        PyCodeObject *co = get_code_with_logging(this_instr + 1);
+        if (co == NULL) {
+            ctx->done = true;
+            break;
+        }
+        _Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0);
+        if (f == NULL) {
+            break;
+        }
+        f->locals[0] = container;
+        f->locals[1] = sub;
+        new_frame = PyJitRef_Wrap((JitOptSymbol *)f);
     }
 
     op(_BINARY_OP_SUBSCR_STR_INT, (str_st, sub_st -- res, s, i)) {
@@ -761,9 +772,19 @@ dummy_func(void) {
     }
 
     op(_LOAD_ATTR_PROPERTY_FRAME, (fget/4, owner -- new_frame)) {
-        (void)fget;
-        new_frame = PyJitRef_NULL;
-        ctx->done = true;
+        // + 1 for _SAVE_RETURN_OFFSET
+        assert((this_instr + 2)->opcode == _PUSH_FRAME);
+        PyCodeObject *co = get_code_with_logging(this_instr + 2);
+        if (co == NULL) {
+            ctx->done = true;
+            break;
+        }
+        _Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0);
+        if (f == NULL) {
+            break;
+        }
+        f->locals[0] = owner;
+        new_frame = PyJitRef_Wrap((JitOptSymbol *)f);
     }
 
     op(_INIT_CALL_BOUND_METHOD_EXACT_ARGS, (callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) {