GH-118093: Add tier two support for BINARY_OP_INPLACE_ADD_UNICODE (GH-122253)

5 files changed, 87 insertions, 5 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index d356fc9..eb8b66f 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -581,12 +581,18 @@ dummy_func(
         // So the inputs are the same as for all BINARY_OP
         // specializations, but there is no output.
         // At the end we just skip over the STORE_FAST.
-        tier1 op(_BINARY_OP_INPLACE_ADD_UNICODE, (left, right --)) {
+        op(_BINARY_OP_INPLACE_ADD_UNICODE, (left, right --)) {
             PyObject *left_o = PyStackRef_AsPyObjectBorrow(left);
             PyObject *right_o = PyStackRef_AsPyObjectBorrow(right);
 
+            int next_oparg;
+        #if TIER_ONE
             assert(next_instr->op.code == STORE_FAST);
-            _PyStackRef *target_local = &GETLOCAL(next_instr->op.arg);
+            next_oparg = next_instr->op.arg;
+        #else
+            next_oparg = CURRENT_OPERAND();
+        #endif
+            _PyStackRef *target_local = &GETLOCAL(next_oparg);
             DEOPT_IF(!PyStackRef_Is(*target_local, left));
             STAT_INC(BINARY_OP, hit);
             /* Handle `left = left + right` or `left += right` for str.
@@ -607,9 +613,12 @@ dummy_func(
             *target_local = PyStackRef_FromPyObjectSteal(temp);
             _Py_DECREF_SPECIALIZED(right_o, _PyUnicode_ExactDealloc);
             ERROR_IF(PyStackRef_IsNull(*target_local), error);
-            // The STORE_FAST is already done.
+        #if TIER_ONE
+            // The STORE_FAST is already done. This is done here in tier one,
+            // and during trace projection in tier two:
             assert(next_instr->op.code == STORE_FAST);
             SKIP_OVER(1);
+        #endif
         }
 
         macro(BINARY_OP_INPLACE_ADD_UNICODE) =
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index b8343f9..3379f0b 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -693,6 +693,55 @@
             break;
         }
 
+        case _BINARY_OP_INPLACE_ADD_UNICODE: {
+            _PyStackRef right;
+            _PyStackRef left;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            PyObject *left_o = PyStackRef_AsPyObjectBorrow(left);
+            PyObject *right_o = 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 = CURRENT_OPERAND();
+            #endif
+            _PyStackRef *target_local = &GETLOCAL(next_oparg);
+            if (!PyStackRef_Is(*target_local, left)) {
+                UOP_STAT_INC(uopcode, miss);
+                JUMP_TO_JUMP_TARGET();
+            }
+            STAT_INC(BINARY_OP, hit);
+            /* Handle `left = left + right` or `left += right` for str.
+             *
+             * When possible, extend `left` in place rather than
+             * allocating a new PyUnicodeObject. This attempts to avoid
+             * quadratic behavior when one neglects to use str.join().
+             *
+             * If `left` has only two references remaining (one from
+             * the stack, one in the locals), DECREFing `left` leaves
+             * only the locals reference, so PyUnicode_Append knows
+             * that the string is safe to mutate.
+             */
+            assert(Py_REFCNT(left_o) >= 2);
+            _Py_DECREF_NO_DEALLOC(left_o);
+            PyObject *temp = PyStackRef_AsPyObjectBorrow(*target_local);
+            PyUnicode_Append(&temp, right_o);
+            *target_local = PyStackRef_FromPyObjectSteal(temp);
+            _Py_DECREF_SPECIALIZED(right_o, _PyUnicode_ExactDealloc);
+            if (PyStackRef_IsNull(*target_local)) JUMP_TO_ERROR();
+            #if TIER_ONE
+            // The STORE_FAST is already done. This is done here in tier one,
+            // and during trace projection in tier two:
+            assert(next_instr->op.code == STORE_FAST);
+            SKIP_OVER(1);
+            #endif
+            stack_pointer += -2;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
         case _BINARY_SUBSCR: {
             _PyStackRef sub;
             _PyStackRef container;
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 6f996f9..c990743 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -181,8 +181,14 @@
             {
                 PyObject *left_o = PyStackRef_AsPyObjectBorrow(left);
                 PyObject *right_o = PyStackRef_AsPyObjectBorrow(right);
+                int next_oparg;
+                #if TIER_ONE
                 assert(next_instr->op.code == STORE_FAST);
-                _PyStackRef *target_local = &GETLOCAL(next_instr->op.arg);
+                next_oparg = next_instr->op.arg;
+                #else
+                next_oparg = CURRENT_OPERAND();
+                #endif
+                _PyStackRef *target_local = &GETLOCAL(next_oparg);
                 DEOPT_IF(!PyStackRef_Is(*target_local, left), BINARY_OP);
                 STAT_INC(BINARY_OP, hit);
                 /* Handle `left = left + right` or `left += right` for str.
@@ -203,9 +209,12 @@
                 *target_local = PyStackRef_FromPyObjectSteal(temp);
                 _Py_DECREF_SPECIALIZED(right_o, _PyUnicode_ExactDealloc);
                 if (PyStackRef_IsNull(*target_local)) goto pop_2_error;
-                // The STORE_FAST is already done.
+                #if TIER_ONE
+                // The STORE_FAST is already done. This is done here in tier one,
+                // and during trace projection in tier two:
                 assert(next_instr->op.code == STORE_FAST);
                 SKIP_OVER(1);
+                #endif
             }
             stack_pointer += -2;
             assert(WITHIN_STACK_BOUNDS());
diff --git a/Python/optimizer.c b/Python/optimizer.c
index e08c1dc..f0793b8 100644
--- a/Python/optimizer.c
+++ b/Python/optimizer.c
@@ -875,6 +875,15 @@ top:  // Jump here after _PUSH_FRAME or likely branches
                             goto done;
                         }
 
+                        if (uop == _BINARY_OP_INPLACE_ADD_UNICODE) {
+                            assert(i + 1 == nuops);
+                            _Py_CODEUNIT *next_instr = instr + 1 + _PyOpcode_Caches[_PyOpcode_Deopt[opcode]];
+                            assert(next_instr->op.code == STORE_FAST);
+                            operand = next_instr->op.arg;
+                            // Skip the STORE_FAST:
+                            instr++;
+                        }
+
                         // All other instructions
                         ADD_TO_TRACE(uop, oparg, operand, target);
                     }
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 8c2b1ac..33af855 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -473,6 +473,12 @@
             break;
         }
 
+        case _BINARY_OP_INPLACE_ADD_UNICODE: {
+            stack_pointer += -2;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
         case _BINARY_SUBSCR: {
             _Py_UopsSymbol *res;
             res = sym_new_not_null(ctx);