[mypyc] feat: support RTuple in sequence_from_generator_preallocate_helper (#19931)

This PR extends `sequence_from_generator_preallocate_helper` to work
with RTuple types. This is ready for review.

Author: BobTheBuidler

mypyc/irbuild/for_helpers.py

@@ -28,6 +28,7 @@
     TypeAlias,
     Var,
 )
+from mypy.types import LiteralType, TupleType, get_proper_type, get_proper_types
 from mypyc.ir.ops import (
     ERR_NEVER,
     BasicBlock,
@@ -36,6 +37,7 @@
     IntOp,
     LoadAddress,
     LoadErrorValue,
+    LoadLiteral,
     LoadMem,
     MethodCall,
     RaiseStandardError,
@@ -170,7 +172,7 @@ def for_loop_helper_with_index(
         body_insts: a function that generates the body of the loop.
                     It needs a index as parameter.
     """
-    assert is_sequence_rprimitive(expr_reg.type)
+    assert is_sequence_rprimitive(expr_reg.type), (expr_reg, expr_reg.type)
     target_type = builder.get_sequence_type(expr)
 
     body_block = BasicBlock()
@@ -217,10 +219,9 @@ def sequence_from_generator_preallocate_helper(
     there is no condition list in the generator and only one original sequence with
     one index is allowed.
 
-    e.g.  (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes)
-          (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes)
-          (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes]
-    RTuple as an original sequence is not supported yet.
+    e.g.  (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple)
+          (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple)
+          (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple]
 
     Args:
         empty_op_llbuilder: A function that can generate an empty sequence op when
@@ -235,23 +236,41 @@ def sequence_from_generator_preallocate_helper(
             implementation.
     """
     if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0:
-        rtype = builder.node_type(gen.sequences[0])
-        if is_sequence_rprimitive(rtype):
-            sequence = builder.accept(gen.sequences[0])
-            length = get_expr_length_value(
-                builder, gen.sequences[0], sequence, gen.line, use_pyssize_t=True
-            )
-            target_op = empty_op_llbuilder(length, gen.line)
-
-            def set_item(item_index: Value) -> None:
-                e = builder.accept(gen.left_expr)
-                builder.call_c(set_item_op, [target_op, item_index, e], gen.line)
-
-            for_loop_helper_with_index(
-                builder, gen.indices[0], gen.sequences[0], sequence, set_item, gen.line, length
-            )
+        line = gen.line
+        sequence_expr = gen.sequences[0]
+        rtype = builder.node_type(sequence_expr)
+        if not (is_sequence_rprimitive(rtype) or isinstance(rtype, RTuple)):
+            return None
+        sequence = builder.accept(sequence_expr)
+        length = get_expr_length_value(builder, sequence_expr, sequence, line, use_pyssize_t=True)
+        if isinstance(rtype, RTuple):
+            # If input is RTuple, box it to tuple_rprimitive for generic iteration
+            # TODO: this can be optimized a bit better with an unrolled ForRTuple helper
+            proper_type = get_proper_type(builder.types[sequence_expr])
+            assert isinstance(proper_type, TupleType), proper_type
+
+            get_item_ops = [
+                (
+                    LoadLiteral(typ.value, object_rprimitive)
+                    if isinstance(typ, LiteralType)
+                    else TupleGet(sequence, i, line)
+                )
+                for i, typ in enumerate(get_proper_types(proper_type.items))
+            ]
+            items = list(map(builder.add, get_item_ops))
+            sequence = builder.new_tuple(items, line)
+
+        target_op = empty_op_llbuilder(length, line)
+
+        def set_item(item_index: Value) -> None:
+            e = builder.accept(gen.left_expr)
+            builder.call_c(set_item_op, [target_op, item_index, e], line)
+
+        for_loop_helper_with_index(
+            builder, gen.indices[0], sequence_expr, sequence, set_item, line, length
+        )
 
-            return target_op
+        return target_op
     return None
 
 
@@ -804,7 +823,7 @@ class ForSequence(ForGenerator):
     def init(
         self, expr_reg: Value, target_type: RType, reverse: bool, length: Value | None = None
     ) -> None:
-        assert is_sequence_rprimitive(expr_reg.type), expr_reg
+        assert is_sequence_rprimitive(expr_reg.type), (expr_reg, expr_reg.type)
         builder = self.builder
         # Record a Value indicating the length of the sequence, if known at compile time.
         self.length = length
@@ -834,7 +853,6 @@ def init(
     def gen_condition(self) -> None:
         builder = self.builder
         line = self.line
-        # TODO: Don't reload the length each time when iterating an immutable sequence?
         if self.reverse:
             # If we are iterating in reverse order, we obviously need
             # to check that the index is still positive. Somewhat less
@@ -1216,7 +1234,7 @@ def get_expr_length_value(
     builder: IRBuilder, expr: Expression, expr_reg: Value, line: int, use_pyssize_t: bool
 ) -> Value:
     rtype = builder.node_type(expr)
-    assert is_sequence_rprimitive(rtype), rtype
+    assert is_sequence_rprimitive(rtype) or isinstance(rtype, RTuple), rtype
     length = get_expr_length(expr)
     if length is None:
         # We cannot compute the length at compile time, so we will fetch it.

mypyc/test-data/irbuild-tuple.test

@@ -694,36 +694,50 @@ L0:
     return r1
 def test():
     r0, source :: tuple[int, int, int]
-    r1 :: list
-    r2, r3, r4 :: object
-    r5, x :: int
-    r6 :: bool
-    r7 :: object
-    r8 :: i32
-    r9, r10 :: bit
-    r11, a :: tuple
+    r1 :: object
+    r2 :: native_int
+    r3 :: bit
+    r4, r5, r6 :: int
+    r7, r8, r9 :: object
+    r10, r11 :: tuple
+    r12 :: native_int
+    r13 :: bit
+    r14 :: object
+    r15, x :: int
+    r16 :: bool
+    r17 :: object
+    r18 :: native_int
+    a :: tuple
 L0:
     r0 = (2, 4, 6)
     source = r0
-    r1 = PyList_New(0)
-    r2 = box(tuple[int, int, int], source)
-    r3 = PyObject_GetIter(r2)
+    r1 = box(tuple[int, int, int], source)
+    r2 = PyObject_Size(r1)
+    r3 = r2 >= 0 :: signed
+    r4 = source[0]
+    r5 = source[1]
+    r6 = source[2]
+    r7 = box(int, r4)
+    r8 = box(int, r5)
+    r9 = box(int, r6)
+    r10 = PyTuple_Pack(3, r7, r8, r9)
+    r11 = PyTuple_New(r2)
+    r12 = 0
 L1:
-    r4 = PyIter_Next(r3)
-    if is_error(r4) goto L4 else goto L2
+    r13 = r12 < r2 :: signed
+    if r13 goto L2 else goto L4 :: bool
 L2:
-    r5 = unbox(int, r4)
-    x = r5
-    r6 = f(x)
-    r7 = box(bool, r6)
-    r8 = PyList_Append(r1, r7)
-    r9 = r8 >= 0 :: signed
+    r14 = CPySequenceTuple_GetItemUnsafe(r10, r12)
+    r15 = unbox(int, r14)
+    x = r15
+    r16 = f(x)
+    r17 = box(bool, r16)
+    CPySequenceTuple_SetItemUnsafe(r11, r12, r17)
 L3:
+    r18 = r12 + 1
+    r12 = r18
     goto L1
 L4:
-    r10 = CPy_NoErrOccurred()
-L5:
-    r11 = PyList_AsTuple(r1)
     a = r11
     return 1
 
@@ -746,42 +760,56 @@ L0:
     r1 = int_eq r0, 0
     return r1
 def test():
-    r0 :: list
-    r1 :: tuple[int, int, int]
-    r2 :: bool
-    r3, r4, r5 :: object
-    r6, x :: int
-    r7 :: bool
-    r8 :: object
-    r9 :: i32
-    r10, r11 :: bit
-    r12, a :: tuple
-L0:
-    r0 = PyList_New(0)
-    r1 = __main__.source :: static
-    if is_error(r1) goto L1 else goto L2
+    r0 :: tuple[int, int, int]
+    r1 :: bool
+    r2 :: object
+    r3 :: native_int
+    r4 :: bit
+    r5, r6, r7 :: int
+    r8, r9, r10 :: object
+    r11, r12 :: tuple
+    r13 :: native_int
+    r14 :: bit
+    r15 :: object
+    r16, x :: int
+    r17 :: bool
+    r18 :: object
+    r19 :: native_int
+    a :: tuple
+L0:
+    r0 = __main__.source :: static
+    if is_error(r0) goto L1 else goto L2
 L1:
-    r2 = raise NameError('value for final name "source" was not set')
+    r1 = raise NameError('value for final name "source" was not set')
     unreachable
 L2:
-    r3 = box(tuple[int, int, int], r1)
-    r4 = PyObject_GetIter(r3)
+    r2 = box(tuple[int, int, int], r0)
+    r3 = PyObject_Size(r2)
+    r4 = r3 >= 0 :: signed
+    r5 = r0[0]
+    r6 = r0[1]
+    r7 = r0[2]
+    r8 = box(int, r5)
+    r9 = box(int, r6)
+    r10 = box(int, r7)
+    r11 = PyTuple_Pack(3, r8, r9, r10)
+    r12 = PyTuple_New(r3)
+    r13 = 0
 L3:
-    r5 = PyIter_Next(r4)
-    if is_error(r5) goto L6 else goto L4
+    r14 = r13 < r3 :: signed
+    if r14 goto L4 else goto L6 :: bool
 L4:
-    r6 = unbox(int, r5)
-    x = r6
-    r7 = f(x)
-    r8 = box(bool, r7)
-    r9 = PyList_Append(r0, r8)
-    r10 = r9 >= 0 :: signed
+    r15 = CPySequenceTuple_GetItemUnsafe(r11, r13)
+    r16 = unbox(int, r15)
+    x = r16
+    r17 = f(x)
+    r18 = box(bool, r17)
+    CPySequenceTuple_SetItemUnsafe(r12, r13, r18)
 L5:
+    r19 = r13 + 1
+    r13 = r19
     goto L3
 L6:
-    r11 = CPy_NoErrOccurred()
-L7:
-    r12 = PyList_AsTuple(r0)
     a = r12
     return 1