fix: rtuple and feat: general literal sequence

BobTheBuidler · BobTheBuidler · commit c3df28bb902e · 2025-07-28T00:56:21.000Z
diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py
@@ -7,14 +7,18 @@
 
 from __future__ import annotations
 
-from typing import Callable, ClassVar
+from typing import Any, Callable, ClassVar, Union
 
 from mypy.nodes import (
     ARG_POS,
+    BytesExpr,
     CallExpr,
     DictionaryComprehension,
     Expression,
+    FloatExpr,
     GeneratorExpr,
+    IntExpr,
+    ListExpr,
     Lvalue,
     MemberExpr,
     NameExpr,
@@ -385,6 +389,27 @@ def is_range_ref(expr: RefExpr) -> bool:
     )
 
 
+def is_literal_expr(expr: Expression) -> bool:
+    # Add other literal types as needed
+    if isinstance(expr, (IntExpr, StrExpr, FloatExpr, BytesExpr)):
+        return True
+    if isinstance(expr, NameExpr) and expr.fullname in {"builtins.None", "builtins.True", "builtins.False"}:
+        return True
+    return False
+
+
+def is_iterable_expr_with_literal_mambers(expr: Expression) -> bool:
+    return (
+        isinstance(expr, (ListExpr, SetExpr, TupleExpr))
+        and not isinstance(expr, MemberExpr)
+        and all(
+            is_literal_expr(item)
+            or is_iterable_expr_with_literal_mambers(item)
+            for item in expr.items
+        )
+    )
+
+
 def make_for_loop_generator(
     builder: IRBuilder,
     index: Lvalue,
@@ -413,21 +438,22 @@ def make_for_loop_generator(
     rtyp = builder.node_type(expr)
 
     # Special case: tuple literal (unroll the loop)
-    if isinstance(expr, TupleExpr):
-        return ForUnrolledLiteral(builder, index, body_block, loop_exit, line, expr.items, expr, body_insts)
+    if is_iterable_expr_with_literal_mambers(expr):
+        return ForUnrolledSequenceLiteral(builder, index, body_block, loop_exit, line, expr, body_insts)
 
-    # Special case: RTuple (known-length tuple, index-based iteration)
+    # Special case: RTuple (known-length tuple, struct field iteration)
     if isinstance(rtyp, RTuple):
         expr_reg = builder.accept(expr)
-        target_type = builder.get_sequence_type(expr)
-        for_tuple = ForSequence(builder, index, body_block, loop_exit, line, nested)
-        for_tuple.init(expr_reg, target_type, reverse=False)
-        return for_tuple
+        return ForUnrolledRTuple(builder, index, body_block, loop_exit, line, rtyp, expr_reg, expr, body_insts)
 
     # Special case: string literal (unroll the loop)
     if isinstance(expr, StrExpr):
         return ForUnrolledStringLiteral(builder, index, body_block, loop_exit, line, expr.value, expr, body_insts)
 
+    # Special case: string literal (unroll the loop)
+    if isinstance(expr, BytesExpr):
+        return ForUnrolledBytesLiteral(builder, index, body_block, loop_exit, line, expr.value, expr, body_insts)
+
     if is_sequence_rprimitive(rtyp):
         # Special case "for x in <list>".
         expr_reg = builder.accept(expr)
@@ -790,7 +816,28 @@ def gen_step(self) -> None:
         pass
 
 
-class ForUnrolledLiteral(ForGenerator):
+class _ForUnrolled(ForGenerator):
+    """Generate IR for a for loop over a value known at compile time by unrolling the loop.
+
+    This class emits the loop body for each element of the value literal directly,
+    avoiding any runtime iteration logic and generator handling.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        if type(self) is _ForUnrolled:
+            raise NotImplementedError("This is a base class and should not be initialized directly.")
+        super().__init__(*args, **kwargs)
+
+    def gen_condition(self) -> None:
+        # Unrolled: nothing to do here.
+        pass
+
+    def gen_step(self) -> None:
+        # Unrolled: nothing to do here.
+        pass
+
+
+class ForUnrolledSequenceLiteral(_ForUnrolled):
     """Generate IR for a for loop over a tuple literal by unrolling the loop.
 
     This class emits the loop body for each element of the tuple literal directly,
@@ -805,31 +852,25 @@ def __init__(
         body_block: BasicBlock,
         loop_exit: BasicBlock,
         line: int,
-        items: list[Expression],
-        expr: Expression,
+        expr: Union[ListExpr, SetExpr, TupleExpr],
         body_insts: GenFunc,
     ) -> None:
         super().__init__(builder, index, body_block, loop_exit, line, nested=False)
-        self.items = items
         self.expr = expr
+        self.items = expr.items
         self.body_insts = body_insts
-
-    def gen_condition(self) -> None:
-        # Unrolled: nothing to do here.
-        pass
+        self.item_types = [builder.node_type(item) for item in self.items]
 
     def begin_body(self) -> None:
         builder = self.builder
-        for item in self.items:
-            builder.assign(builder.get_assignment_target(self.index), builder.accept(item), self.line)
+        for item, item_type in zip(self.items, self.item_types):
+            value = builder.accept(item)
+            value = builder.coerce(value, item_type, self.line)
+            builder.assign(builder.get_assignment_target(self.index), value, self.line)
             self.body_insts()
 
-    def gen_step(self) -> None:
-        # Unrolled: nothing to do here.
-        pass
-
 
-class ForUnrolledStringLiteral(ForGenerator):
+class ForUnrolledStringLiteral(_ForUnrolled):
     """Generate IR for a for loop over a string literal by unrolling the loop.
 
     This class emits the loop body for each character of the string literal directly,
@@ -853,10 +894,6 @@ def __init__(
         self.expr = expr
         self.body_insts = body_insts
 
-    def gen_condition(self) -> None:
-        # Unrolled: nothing to do here.
-        pass
-
     def begin_body(self) -> None:
         builder = self.builder
         for c in self.value:
@@ -867,9 +904,74 @@ def begin_body(self) -> None:
             )
             self.body_insts()
 
-    def gen_step(self) -> None:
-        # Unrolled: nothing to do here.
-        pass
+
+class ForUnrolledBytesLiteral(_ForUnrolled):
+    """Generate IR for a for loop over a string literal by unrolling the loop.
+
+    This class emits the loop body for each character of the string literal directly,
+    avoiding any runtime iteration logic.
+    """
+    handles_body_insts = True
+
+    def __init__(
+        self,
+        builder: IRBuilder,
+        index: Lvalue,
+        body_block: BasicBlock,
+        loop_exit: BasicBlock,
+        line: int,
+        value: bytes,
+        expr: Expression,
+        body_insts: GenFunc,
+    ) -> None:
+        super().__init__(builder, index, body_block, loop_exit, line, nested=False)
+        self.value = value
+        self.expr = expr
+        self.body_insts = body_insts
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        for c in self.value:
+            builder.assign(
+                builder.get_assignment_target(self.index),
+                builder.accept(IntExpr(c)),
+                self.line,
+            )
+            self.body_insts()
+
+
+class ForUnrolledRTuple(_ForUnrolled):
+    """Generate IR for a for loop over an RTuple by directly accessing struct fields."""
+
+    handles_body_insts = True
+
+    def __init__(
+        self,
+        builder: IRBuilder,
+        index: Lvalue,
+        body_block: BasicBlock,
+        loop_exit: BasicBlock,
+        line: int,
+        rtuple_type: RTuple,
+        expr_reg: Value,
+        expr: Expression,
+        body_insts: GenFunc,
+    ) -> None:
+        super().__init__(builder, index, body_block, loop_exit, line, nested=False)
+        self.rtuple_type = rtuple_type
+        self.expr_reg = expr_reg
+        self.expr = expr
+        self.body_insts = body_insts
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+        for i, item_type in enumerate(self.rtuple_type.types):
+            # Directly access the struct field for each RTuple element
+            value = builder.add(TupleGet(self.expr_reg, i, line))
+            value = builder.coerce(value, item_type, line)
+            builder.assign(builder.get_assignment_target(self.index), value, line)
+            self.body_insts()
 
 
 def unsafe_index(builder: IRBuilder, target: Value, index: Value, line: int) -> Value:
diff --git a/mypyc/test-data/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test
@@ -3605,15 +3605,17 @@ def f(t: Tuple[int, int]) -> int:
 [out]
 def f(t):
     t :: tuple[int, int]
-    s, x :: int
-    r0 :: int
+    s, r0, x, r1, r2, r3 :: int
+L0:
     s = 0
-    x = t.f0
-    L1:
-        s = CPyTagged_Add(s, x)
-    x = t.f1
-    L2:
-        s = CPyTagged_Add(s, x)
+    r0 = t[0]
+    x = r0
+    r1 = CPyTagged_Add(s, x)
+    s = r1
+    r2 = t[1]
+    x = r2
+    r3 = CPyTagged_Add(s, x)
+    s = r3
     return s
 
 [case testForOverStringVar]
@@ -3640,3 +3642,12 @@ def f(s):
         goto L1
     L4:
         return out
+
+[case TestForOverCompledTupleExpr]
+def f() -> None:
+    abc = (1, 2, str(3))
+    for x in abc:
+        y = x
+[out]
+def f():
+    abc :: tuple[int, int, str]
diff --git a/mypyc/test-data/irbuild-set.test b/mypyc/test-data/irbuild-set.test
@@ -129,41 +129,41 @@ L4:
 def test2():
     r0, tmp_tuple :: tuple[int, int, int]
     r1 :: set
-    r2 :: native_int
-    r3 :: object
-    r4 :: native_int
-    r5, r6 :: bit
-    r7 :: object
-    r8 :: bit
-    r9, r10, r12 :: int
-    r13 :: object
-    r14, x, r15 :: int
-    r16 :: object
-    r17 :: i32
-    r18 :: bit
-    r19 :: native_int
+    r2, x, r3 :: int
+    r4 :: object
+    r5 :: i32
+    r6 :: bit
+    r7, r8 :: int
+    r9 :: object
+    r10 :: i32
+    r11 :: bit
+    r12, r13 :: int
+    r14 :: object
+    r15 :: i32
+    r16 :: bit
     b :: set
 L0:
     r0 = (2, 6, 10)
     tmp_tuple = r0
     r1 = PySet_New(0)
-    r2 = box(tuple[int, int, int], tmp_tuple)
-    r3 = PyObject_GetIter(r2)
-L1:
-    r4 = PyIter_Next(r3)
-    if is_error(r4) goto L4 else goto L2
-L2:
-    r5 = unbox(int, r4)
-    x = r5
-    r6 = f(x)
-    r7 = box(int, r6)
-    r8 = PySet_Add(r1, r7)
-    r9 = r8 >= 0 :: signed
-L3:
-    goto L1
-L4:
-    r10 = CPy_NoErrOccurred()
-L5:
+    r2 = tmp_tuple[0]
+    x = r2
+    r3 = f(x)
+    r4 = box(int, r3)
+    r5 = PySet_Add(r1, r4)
+    r6 = r5 >= 0 :: signed
+    r7 = tmp_tuple[1]
+    x = r7
+    r8 = f(x)
+    r9 = box(int, r8)
+    r10 = PySet_Add(r1, r9)
+    r11 = r10 >= 0 :: signed
+    r12 = tmp_tuple[2]
+    x = r12
+    r13 = f(x)
+    r14 = box(int, r13)
+    r15 = PySet_Add(r1, r14)
+    r16 = r15 >= 0 :: signed
     b = r1
     return 1
 def test3():
@@ -735,25 +735,15 @@ def not_precomputed() -> None:
 
 [out]
 def precomputed():
-    r0 :: set
-    r1, r2 :: object
-    r3 :: str
+    r0 :: str
     _ :: object
-    r4 :: bit
 L0:
-    r0 = frozenset({'False', 'None', 'True'})
-    r1 = PyObject_GetIter(r0)
-L1:
-    r2 = PyIter_Next(r1)
-    if is_error(r2) goto L4 else goto L2
-L2:
-    r3 = cast(str, r2)
-    _ = r3
-L3:
-    goto L1
-L4:
-    r4 = CPy_NoErrOccurred()
-L5:
+    r0 = 'None'
+    _ = r0
+    r1 = 'True'
+    _ = r1
+    r2 = 'False'
+    _ = r2
     return 1
 def precomputed2():
     r0 :: set
