gh-100239: Specialize binary operations using BINARY_OP_EXTEND

Include/internal/pycore_code.h

@@ -503,6 +503,11 @@ typedef struct {
        aliased to either operand). Used by the tier 2 optimizer to enable
        inplace follow-up ops. */
     int result_unique;
+    /* Expected types of the left and right operands. Used by the tier 2
+       optimizer to eliminate _GUARD_BINARY_OP_EXTEND when the operand
+       types are already known. NULL means unknown/don't eliminate. */
+    PyTypeObject *lhs_type;
+    PyTypeObject *rhs_type;
 } _PyBinaryOpSpecializationDescr;
 
 /* Comparison bit masks. */

Include/internal/pycore_list.h

@@ -15,6 +15,7 @@ extern "C" {
 PyAPI_FUNC(PyObject*) _PyList_Extend(PyListObject *, PyObject *);
 PyAPI_FUNC(PyObject) *_PyList_SliceSubscript(PyObject*, PyObject*);
 PyAPI_FUNC(PyObject *) _PyList_BinarySlice(PyObject *, PyObject *, PyObject *);
+PyAPI_FUNC(PyObject *) _PyList_Concat(PyObject *, PyObject *);
 extern void _PyList_DebugMallocStats(FILE *out);
 // _PyList_GetItemRef should be used only when the object is known as a list
 // because it doesn't raise TypeError when the object is not a list, whereas PyList_GetItemRef does.

Include/internal/pycore_tuple.h

@@ -28,6 +28,7 @@ PyAPI_FUNC(void) _PyStolenTuple_Free(PyObject *self);
 PyAPI_FUNC(PyObject *)_PyTuple_FromStackRefStealOnSuccess(const union _PyStackRef *, Py_ssize_t);
 PyAPI_FUNC(PyObject *)_PyTuple_FromArraySteal(PyObject *const *, Py_ssize_t);
 PyAPI_FUNC(PyObject *) _PyTuple_BinarySlice(PyObject *, PyObject *, PyObject *);
+PyAPI_FUNC(PyObject *) _PyTuple_Concat(PyObject *, PyObject *);
 
 PyAPI_FUNC(PyObject *) _PyTuple_FromPair(PyObject *, PyObject *);
 PyAPI_FUNC(PyObject *) _PyTuple_FromPairSteal(PyObject *, PyObject *);

Lib/test/test_capi/test_opt.py

@@ -3836,6 +3836,72 @@ def testfunc(n):
         # propagates PyFloat_Type.
         self.assertNotIn("_GUARD_NOS_FLOAT", uops)
 
+    def test_binary_op_extend_list_concat_type_propagation(self):
+        # list + list is specialized via BINARY_OP_EXTEND. The tier 2 optimizer
+        # should learn that the result is a list and eliminate subsequent
+        # list-type guards.
+        def testfunc(n):
+            a = [1, 2]
+            b = [3, 4]
+            x = True
+            for _ in range(n):
+                c = a + b
+                if c[0]:
+                    x = False
+            return x
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, False)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_BINARY_OP_EXTEND", uops)
+        # The c[0] subscript emits _GUARD_NOS_LIST before _BINARY_OP_SUBSCR_LIST_INT;
+        # since _BINARY_OP_EXTEND now propagates PyList_Type, that guard is gone.
+        self.assertIn("_BINARY_OP_SUBSCR_LIST_INT", uops)
+        self.assertNotIn("_GUARD_NOS_LIST", uops)
+
+    def test_binary_op_extend_tuple_concat_type_propagation(self):
+        # tuple + tuple is specialized via BINARY_OP_EXTEND. The tier 2 optimizer
+        # should learn the result is a tuple and eliminate subsequent tuple guards.
+        def testfunc(n):
+            t1 = (1, 2)
+            t2 = (3, 4)
+            for _ in range(n):
+                a, b, c, d = t1 + t2
+            return a + b + c + d
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, 10)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_BINARY_OP_EXTEND", uops)
+        self.assertIn("_UNPACK_SEQUENCE_TUPLE", uops)
+        self.assertNotIn("_GUARD_TOS_TUPLE", uops)
+
+    def test_binary_op_extend_guard_elimination(self):
+        # When both operands have known types (e.g., from a prior
+        # _BINARY_OP_EXTEND result), the _GUARD_BINARY_OP_EXTEND
+        # should be eliminated.
+        def testfunc(n):
+            a = [1, 2]
+            b = [3, 4]
+            total = 0
+            for _ in range(n):
+                c = a + b    # first: guard stays, result type = list
+                d = c + c    # second: both operands are list -> guard eliminated
+                total += d[0]
+            return total
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        # Both list additions use _BINARY_OP_EXTEND
+        self.assertEqual(uops.count("_BINARY_OP_EXTEND"), 2)
+        # But the second guard is eliminated because both operands
+        # are known to be lists from the first _BINARY_OP_EXTEND.
+        self.assertEqual(uops.count("_GUARD_BINARY_OP_EXTEND"), 1)
+
     def test_unary_invert_long_type(self):
         def testfunc(n):
             for _ in range(n):

Lib/test/test_opcache.py

@@ -1423,6 +1423,21 @@ def binary_op_add_extend():
         self.assert_specialized(binary_op_add_extend, "BINARY_OP_EXTEND")
         self.assert_no_opcode(binary_op_add_extend, "BINARY_OP")
 
+        def binary_op_add_extend_sequences():
+            l1 = [1, 2]
+            l2 = [None]
+            t1 = (1, 2)
+            t2 = (None,)
+            for _ in range(100):
+                list_sum = l1 + l2
+                self.assertEqual(list_sum, [1, 2, None])
+                tuple_sum = t1 + t2
+                self.assertEqual(tuple_sum, (1, 2, None))
+
+        binary_op_add_extend_sequences()
+        self.assert_specialized(binary_op_add_extend_sequences, "BINARY_OP_EXTEND")
+        self.assert_no_opcode(binary_op_add_extend_sequences, "BINARY_OP")
+
         def binary_op_zero_division():
             def compactlong_lhs(arg):
                 42 / arg

Misc/NEWS.d/next/Core_and_Builtins/2025-01-17-19-48-28.gh-issue-100239.7pbTEA.rst

@@ -0,0 +1,3 @@
+Specialize ``BINARY_OP`` for concatenation of lists and tuples, and
+propagate the result type through ``_BINARY_OP_EXTEND`` in the tier 2
+optimizer so that follow-up type guards can be eliminated.

Objects/listobject.c

@@ -798,8 +798,8 @@ list_concat_lock_held(PyListObject *a, PyListObject *b)
     return (PyObject *)np;
 }
 
-static PyObject *
-list_concat(PyObject *aa, PyObject *bb)
+PyObject *
+_PyList_Concat(PyObject *aa, PyObject *bb)
 {
     if (!PyList_Check(bb)) {
         PyErr_Format(PyExc_TypeError,
@@ -3617,7 +3617,7 @@ static PyMethodDef list_methods[] = {
 
 static PySequenceMethods list_as_sequence = {
     list_length,                                /* sq_length */
-    list_concat,                                /* sq_concat */
+    _PyList_Concat,                             /* sq_concat */
     list_repeat,                                /* sq_repeat */
     list_item,                                  /* sq_item */
     0,                                          /* sq_slice */

Objects/tupleobject.c

@@ -547,8 +547,8 @@ PyTuple_GetSlice(PyObject *op, Py_ssize_t i, Py_ssize_t j)
     return tuple_slice((PyTupleObject *)op, i, j);
 }
 
-static PyObject *
-tuple_concat(PyObject *aa, PyObject *bb)
+PyObject *
+_PyTuple_Concat(PyObject *aa, PyObject *bb)
 {
     PyTupleObject *a = _PyTuple_CAST(aa);
     if (Py_SIZE(a) == 0 && PyTuple_CheckExact(bb)) {
@@ -864,7 +864,7 @@ tuple_subtype_new(PyTypeObject *type, PyObject *iterable)
 
 static PySequenceMethods tuple_as_sequence = {
     tuple_length,                               /* sq_length */
-    tuple_concat,                               /* sq_concat */
+    _PyTuple_Concat,                            /* sq_concat */
     tuple_repeat,                               /* sq_repeat */
     tuple_item,                                 /* sq_item */
     0,                                          /* sq_slice */

Python/bytecodes.c

@@ -983,6 +983,8 @@ dummy_func(
             if (res_o == NULL) {
                 ERROR_NO_POP();
             }
+            assert(d->result_type == NULL || Py_TYPE(res_o) == d->result_type);
+            assert(!d->result_unique || Py_REFCNT(res_o) == 1 || _Py_IsImmortal(res_o));
             res = PyStackRef_FromPyObjectSteal(res_o);
             l = left;
             r = right;

Python/optimizer_bytecodes.c

@@ -409,6 +409,16 @@ dummy_func(void) {
         r = right;
     }
 
+    op(_GUARD_BINARY_OP_EXTEND, (descr/4, left, right -- left, right)) {
+        _PyBinaryOpSpecializationDescr *d = (_PyBinaryOpSpecializationDescr *)descr;
+        if (d != NULL && d->lhs_type != NULL && d->rhs_type != NULL) {
+            if (sym_matches_type(left, d->lhs_type) &&
+                sym_matches_type(right, d->rhs_type)) {
+                REPLACE_OP(this_instr, _NOP, 0, 0);
+            }
+        }
+    }
+
     op(_BINARY_OP_EXTEND, (descr/4, left, right -- res, l, r)) {
         _PyBinaryOpSpecializationDescr *d = (_PyBinaryOpSpecializationDescr *)descr;
         if (d != NULL && d->result_type != NULL) {