[mypyc] Optimize list.__add__, list.__iadd__, tuple.__add__ (#18845)

https://docs.python.org/3/c-api/sequence.html#c.PySequence_Concat
https://docs.python.org/3/c-api/sequence.html#c.PySequence_InPlaceConcat

Author: cdce8p

mypyc/doc/list_operations.rst

@@ -32,6 +32,7 @@ Operators
 
 * ``lst[n]`` (get item by integer index)
 * ``lst[n:m]``, ``lst[n:]``, ``lst[:m]``, ``lst[:]`` (slicing)
+* ``lst1 + lst2``, ``lst += iter``
 * ``lst * n``, ``n * lst``
 * ``obj in lst``
 

mypyc/doc/tuple_operations.rst

@@ -21,6 +21,7 @@ Operators
 
 * ``tup[n]`` (integer index)
 * ``tup[n:m]``, ``tup[n:]``, ``tup[:m]`` (slicing)
+* ``tup1 + tup2``
 
 Statements
 ----------

mypyc/primitives/list_ops.py

@@ -271,6 +271,24 @@
     error_kind=ERR_MAGIC,
 )
 
+# list + list
+binary_op(
+    name="+",
+    arg_types=[list_rprimitive, list_rprimitive],
+    return_type=list_rprimitive,
+    c_function_name="PySequence_Concat",
+    error_kind=ERR_MAGIC,
+)
+
+# list += list
+binary_op(
+    name="+=",
+    arg_types=[list_rprimitive, object_rprimitive],
+    return_type=list_rprimitive,
+    c_function_name="PySequence_InPlaceConcat",
+    error_kind=ERR_MAGIC,
+)
+
 # list * int
 binary_op(
     name="*",

mypyc/primitives/tuple_ops.py

@@ -15,7 +15,7 @@
     object_rprimitive,
     tuple_rprimitive,
 )
-from mypyc.primitives.registry import custom_op, function_op, load_address_op, method_op
+from mypyc.primitives.registry import binary_op, custom_op, function_op, load_address_op, method_op
 
 # Get the 'builtins.tuple' type object.
 load_address_op(name="builtins.tuple", type=object_rprimitive, src="PyTuple_Type")
@@ -74,6 +74,15 @@
     error_kind=ERR_MAGIC,
 )
 
+# tuple + tuple
+binary_op(
+    name="+",
+    arg_types=[tuple_rprimitive, tuple_rprimitive],
+    return_type=tuple_rprimitive,
+    c_function_name="PySequence_Concat",
+    error_kind=ERR_MAGIC,
+)
+
 # tuple[begin:end]
 tuple_slice_op = custom_op(
     arg_types=[tuple_rprimitive, int_rprimitive, int_rprimitive],

mypyc/test-data/fixtures/ir.py

@@ -208,6 +208,10 @@ def __getitem__(self, i: slice) -> Tuple[T_co, ...]: pass
     def __len__(self) -> int: pass
     def __iter__(self) -> Iterator[T_co]: ...
     def __contains__(self, item: object) -> int: ...
+    @overload
+    def __add__(self, value: Tuple[T_co, ...], /) -> Tuple[T_co, ...]: ...
+    @overload
+    def __add__(self, value: Tuple[_T, ...], /) -> Tuple[T_co | _T, ...]: ...
 
 class function: pass
 
@@ -224,7 +228,11 @@ def __rmul__(self, i: int) -> List[_T]: pass
     def __iter__(self) -> Iterator[_T]: pass
     def __len__(self) -> int: pass
     def __contains__(self, item: object) -> int: ...
-    def __add__(self, x: List[_T]) -> List[_T]: ...
+    @overload
+    def __add__(self, value: List[_T], /) -> List[_T]: ...
+    @overload
+    def __add__(self, value: List[_S], /) -> List[_S | _T]: ...
+    def __iadd__(self, value: Iterable[_T], /) -> List[_T]: ...  # type: ignore[misc]
     def append(self, x: _T) -> None: pass
     def pop(self, i: int = -1) -> _T: pass
     def count(self, _T) -> int: pass

mypyc/test-data/irbuild-lists.test

@@ -145,6 +145,32 @@ L0:
     x = r10
     return 1
 
+[case testListAdd]
+from typing import List
+def f(a: List[int], b: List[int]) -> None:
+    c = a + b
+[out]
+def f(a, b):
+    a, b, r0, c :: list
+L0:
+    r0 = PySequence_Concat(a, b)
+    c = r0
+    return 1
+
+[case testListIAdd]
+from typing import List, Any
+def f(a: List[int], b: Any) -> None:
+    a += b
+[out]
+def f(a, b):
+    a :: list
+    b :: object
+    r0 :: list
+L0:
+    r0 = PySequence_InPlaceConcat(a, b)
+    a = r0
+    return 1
+
 [case testListMultiply]
 from typing import List
 def f(a: List[int]) -> None:

mypyc/test-data/irbuild-tuple.test

@@ -384,3 +384,37 @@ L3:
 L4:
     a = r1
     return 1
+
+[case testTupleAdd]
+from typing import Tuple
+def f(a: Tuple[int, ...], b: Tuple[int, ...]) -> None:
+    c = a + b
+    d = a + (1, 2)
+def g(a: Tuple[int, int], b: Tuple[int, int]) -> None:
+    c = a + b
+[out]
+def f(a, b):
+    a, b, r0, c :: tuple
+    r1 :: tuple[int, int]
+    r2 :: object
+    r3, d :: tuple
+L0:
+    r0 = PySequence_Concat(a, b)
+    c = r0
+    r1 = (2, 4)
+    r2 = box(tuple[int, int], r1)
+    r3 = PySequence_Concat(a, r2)
+    d = r3
+    return 1
+def g(a, b):
+    a, b :: tuple[int, int]
+    r0, r1 :: object
+    r2 :: tuple
+    r3, c :: tuple[int, int, int, int]
+L0:
+    r0 = box(tuple[int, int], a)
+    r1 = box(tuple[int, int], b)
+    r2 = PySequence_Concat(r0, r1)
+    r3 = unbox(tuple[int, int, int, int], r2)
+    c = r3
+    return 1

mypyc/test-data/run-lists.test

@@ -267,6 +267,9 @@ print(g())
 7
 
 [case testListOps]
+from typing import Any, cast
+from testutil import assertRaises
+
 def test_slicing() -> None:
     # Use dummy adds to avoid constant folding
     zero = int()
@@ -289,6 +292,27 @@ def test_slicing() -> None:
     assert s[long_int:] == []
     assert s[-long_int:-1] == ["f", "o", "o", "b", "a"]
 
+def in_place_add(l2: Any) -> list[Any]:
+    l1 = [1, 2]
+    l1 += l2
+    return l1
+
+def test_add() -> None:
+    res = [1, 2, 3, 4]
+    assert [1, 2] + [3, 4] == res
+    with assertRaises(TypeError, 'can only concatenate list (not "tuple") to list'):
+        assert [1, 2] + cast(Any, (3, 4)) == res
+    l1 = [1, 2]
+    id_l1 = id(l1)
+    l1 += [3, 4]
+    assert l1 == res
+    assert id_l1 == id(l1)
+    assert in_place_add([3, 4]) == res
+    assert in_place_add((3, 4)) == res
+    assert in_place_add({3, 4}) == res
+    assert in_place_add({3: "", 4: ""}) == res
+    assert in_place_add(range(3, 5)) == res
+
 [case testOperatorInExpression]
 
 def tuple_in_int0(i: int) -> bool:

mypyc/test-data/run-tuples.test

@@ -146,7 +146,8 @@ assert Record.__annotations__ == {
 }, Record.__annotations__
 
 [case testTupleOps]
-from typing import Tuple, Final, List, Any, Optional
+from typing import Tuple, Final, List, Any, Optional, cast
+from testutil import assertRaises
 
 def f() -> Tuple[()]:
     return ()
@@ -254,3 +255,9 @@ TUPLE: Final[Tuple[str, ...]] = ('x', 'y')
 def test_final_boxed_tuple() -> None:
     t = TUPLE
     assert t == ('x', 'y')
+
+def test_add() -> None:
+    res = (1, 2, 3, 4)
+    assert (1, 2) + (3, 4) == res
+    with assertRaises(TypeError, 'can only concatenate tuple (not "list") to tuple'):
+        assert (1, 2) + cast(Any, [3, 4]) == res