fix BINARY_OP_INPLACE_ADD_UNICODE check

Author: cocolato

Include/internal/pycore_opcode_metadata.h

@@ -4407,6 +4407,67 @@ def g():
         PYTHON_JIT="1", PYTHON_JIT_STRESS="1")
         self.assertEqual(result[0].rc, 0, result)
 
+    def test_143820(self):
+        # https://github.com/python/cpython/issues/143358
+
+        result = script_helper.run_python_until_end('-c', textwrap.dedent(f"""
+        import sys
+        import random
+
+        int_v1 = 981679
+        int_v2 = -3791744241805517
+        any_v3 = 939.217
+
+        def f1(): int_v1 ^ int_v1
+
+        for i_f1 in range(300):
+            f1()
+
+        def f2():
+            class Int(int):
+                def __index__(self):...
+
+            inf = float('inf')
+            nzero = -0
+            zero = 0.0
+            dummy = 0
+            print('', file=sys.stderr)
+
+            def f_0_dc_6103(p): return p + 1
+            def f_1_dc_6103(p): return f_0_dc_6103(p) + 1
+            def f_2_dc_6103(p): return f_1_dc_6103(p) + 1
+            def f_3_dc_6103(p): return f_2_dc_6103(p) + 1
+            def f_4_dc_6103(p): return f_3_dc_6103(p) + 1
+            def f_5_dc_6103(p): return f_4_dc_6103(p) + 1
+            def f_6_dc_6103(p): return f_5_dc_6103(p) + 1
+            def f_7_dc_6103(p): return f_6_dc_6103(p) + 1
+            def f_8_dc_6103(p): return f_7_dc_6103(p) + 1
+            def f_9_dc_6103(p): return f_8_dc_6103(p) + 1
+
+            if inf == inf: dummy += 1
+            s = ''
+            try:
+                for _ in range(10):
+                    s += ''
+                s += 'y'
+            except Exception: pass
+            int_v1 ^ int_v1
+            int_v1 ^ int_v1
+            int_v1 ^ int_v1
+            int_v2 - int_v1
+            int_v2 - int_v1
+            int_v2 - int_v1
+            int_v2 - int_v1
+            int_v2 - int_v1
+            not any_v3
+            not any_v3
+            not any_v3
+
+        for i_f2 in range(300):
+            f2()
+        """), PYTHON_JIT="1", PYTHON_JIT_STRESS="1")
+        self.assertEqual(result[0].rc, 0, result)
+
 def global_identity(x):
     return x
 

Lib/test/test_capi/test_opt.py

@@ -794,18 +794,12 @@ dummy_func(
         // This is a subtle one. We write NULL to the local
         // of the following STORE_FAST and leave the result for STORE_FAST
         // later to store.
-        op(_BINARY_OP_INPLACE_ADD_UNICODE, (left, right -- res)) {
+        op(_BINARY_OP_INPLACE_ADD_UNICODE, (next_oparg_idx/1, left, right -- res)) {
             PyObject *left_o = PyStackRef_AsPyObjectBorrow(left);
             assert(PyUnicode_CheckExact(left_o));
             assert(PyUnicode_CheckExact(PyStackRef_AsPyObjectBorrow(right)));
 
-            int next_oparg;
-        #if TIER_ONE
-            assert(next_instr->op.code == STORE_FAST);
-            next_oparg = next_instr->op.arg;
-        #else
-            next_oparg = (int)CURRENT_OPERAND0_16();
-        #endif
+            int next_oparg = (int)next_oparg_idx;
             _PyStackRef *target_local = &GETLOCAL(next_oparg);
             assert(PyUnicode_CheckExact(left_o));
             EXIT_IF(PyStackRef_AsPyObjectBorrow(*target_local) != left_o);
@@ -870,7 +864,7 @@ dummy_func(
             unused/1 + _GUARD_BINARY_OP_EXTEND + rewind/-4 + _BINARY_OP_EXTEND + POP_TOP + POP_TOP;
 
         macro(BINARY_OP_INPLACE_ADD_UNICODE) =
-            _GUARD_TOS_UNICODE + _GUARD_NOS_UNICODE + unused/5 + _BINARY_OP_INPLACE_ADD_UNICODE;
+            _GUARD_TOS_UNICODE + _GUARD_NOS_UNICODE + unused/4 + _BINARY_OP_INPLACE_ADD_UNICODE;
 
         specializing op(_SPECIALIZE_BINARY_SLICE, (container, start, stop -- container, start, stop)) {
             // Placeholder until we implement BINARY_SLICE specialization

Modules/_testinternalcapi/test_cases.c.h

@@ -355,7 +355,7 @@ dummy_func(void) {
         r = right;
     }
 
-    op(_BINARY_OP_INPLACE_ADD_UNICODE, (left, right -- res)) {
+    op(_BINARY_OP_INPLACE_ADD_UNICODE, (next_oparg_idx/1, left, right -- res)) {
         if (sym_is_const(ctx, left) && sym_is_const(ctx, right)) {
             assert(PyUnicode_CheckExact(sym_get_const(ctx, left)));
             assert(PyUnicode_CheckExact(sym_get_const(ctx, right)));
@@ -369,7 +369,7 @@ dummy_func(void) {
         else {
             res = sym_new_type(ctx, &PyUnicode_Type);
         }
-        GETLOCAL(this_instr->operand0) = sym_new_null(ctx);
+        GETLOCAL(next_oparg_idx) = sym_new_null(ctx);
     }
 
     op(_BINARY_OP_SUBSCR_CHECK_FUNC, (container, unused -- container, unused, getitem)) {

Python/bytecodes.c

@@ -2251,10 +2251,22 @@ _Py_Specialize_BinaryOp(_PyStackRef lhs_st, _PyStackRef rhs_st, _Py_CODEUNIT *in
                 break;
             }
             if (PyUnicode_CheckExact(lhs)) {
-                _Py_CODEUNIT next = instr[INLINE_CACHE_ENTRIES_BINARY_OP + 1];
+                int offset = INLINE_CACHE_ENTRIES_BINARY_OP + 1;
+                int locals_index = 0;
+                // Walk through EXTENDED_ARG instructions to compute full oparg
+                while (instr[offset].op.code == EXTENDED_ARG) {
+                    locals_index = (locals_index | instr[offset].op.arg) << 8;
+                    offset++;
+                }
+                _Py_CODEUNIT next = instr[offset];
                 bool to_store = (next.op.code == STORE_FAST);
-                if (to_store && PyStackRef_AsPyObjectBorrow(locals[next.op.arg]) == lhs) {
+                locals_index |= next.op.arg;
+                if (to_store && PyStackRef_AsPyObjectBorrow(locals[locals_index]) == lhs) {
                     specialize(instr, BINARY_OP_INPLACE_ADD_UNICODE);
+                    // Store local index in the last cache slot
+                    // so we don't need to peek at next_instr at runtime.
+                    instr[INLINE_CACHE_ENTRIES_BINARY_OP].cache =
+                        (uint16_t)locals_index;
                     return;
                 }
                 specialize(instr, BINARY_OP_ADD_UNICODE);