From b5b06349eb71b7cf9e5082e26e6fe0145875f95b Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@microsoft.com>
Date: Mon, 28 Oct 2024 14:37:16 -0700
Subject: [PATCH] GH-125912: Teach the JIT's optimizer about
 _BINARY_OP_INPLACE_ADD_UNICODE (GH-125935)

---
 Python/optimizer_bytecodes.c | 18 ++++++++++++++++++
 Python/optimizer_cases.c.h   | 19 +++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index d71b55cbe1c..f40ad1e5744 100644
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -331,6 +331,24 @@ dummy_func(void) {
         }
     }
 
+    op(_BINARY_OP_INPLACE_ADD_UNICODE, (left, right -- )) {
+        _Py_UopsSymbol *res;
+        if (sym_is_const(left) && sym_is_const(right) &&
+            sym_matches_type(left, &PyUnicode_Type) && sym_matches_type(right, &PyUnicode_Type)) {
+            PyObject *temp = PyUnicode_Concat(sym_get_const(left), sym_get_const(right));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            Py_DECREF(temp);
+        }
+        else {
+            res = sym_new_type(ctx, &PyUnicode_Type);
+        }
+        // _STORE_FAST:
+        GETLOCAL(this_instr->operand) = res;
+    }
+
     op(_BINARY_SUBSCR_INIT_CALL, (container, sub -- new_frame: _Py_UOpsAbstractFrame *)) {
         (void)container;
         (void)sub;
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 6ec9e69d1db..243b3efa41b 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -484,6 +484,25 @@
         }
 
         case _BINARY_OP_INPLACE_ADD_UNICODE: {
+            _Py_UopsSymbol *right;
+            _Py_UopsSymbol *left;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            _Py_UopsSymbol *res;
+            if (sym_is_const(left) && sym_is_const(right) &&
+                sym_matches_type(left, &PyUnicode_Type) && sym_matches_type(right, &PyUnicode_Type)) {
+                PyObject *temp = PyUnicode_Concat(sym_get_const(left), sym_get_const(right));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                Py_DECREF(temp);
+            }
+            else {
+                res = sym_new_type(ctx, &PyUnicode_Type);
+            }
+            // _STORE_FAST:
+            GETLOCAL(this_instr->operand) = res;
             stack_pointer += -2;
             assert(WITHIN_STACK_BOUNDS());
             break;