more stuff

Author: BobTheBuidler

mypyc/irbuild/for_helpers.py

@@ -86,6 +86,7 @@
 
 GenFunc = Callable[[], None]
 
+
 def for_loop_helper(
     builder: IRBuilder,
     index: Lvalue,
@@ -94,7 +95,7 @@ def for_loop_helper(
     else_insts: GenFunc | None,
     is_async: bool,
     line: int,
-    can_unroll: bool = True,
+    can_unroll: bool = False,
 ) -> None:
     """Generate IR for a loop.
 

mypyc/irbuild/statement.py

@@ -13,6 +13,7 @@
 from typing import Callable
 
 import mypy.nodes
+import mypy.traverser
 from mypy.nodes import (
     ARG_NAMED,
     ARG_POS,
@@ -130,6 +131,29 @@
 ValueGenFunc = Callable[[], Value]
 
 
+class ControlFlowDetector(mypy.traverser.TraverserVisitor):
+    """
+    A Visitor class that detects whether a block contains any of the following control flow statements:
+
+    - return
+    - continue
+    - break
+    
+    """
+    def __init__(self) -> None:
+        super().__init__()
+        self.has_control_flow = False
+
+    def visit_break_stmt(self, o: BreakStmt) -> None:
+        self.has_control_flow = True
+
+    def visit_continue_stmt(self, o: ContinueStmt) -> None:
+        self.has_control_flow = True
+
+    def visit_return_stmt(self, o: ReturnStmt) -> None:
+        self.has_control_flow = True
+
+
 def transform_block(builder: IRBuilder, block: Block) -> None:
     if not block.is_unreachable:
         builder.block_reachable_stack.append(True)
@@ -434,16 +458,37 @@ def transform_while_stmt(builder: IRBuilder, s: WhileStmt) -> None:
 
 
 def transform_for_stmt(builder: IRBuilder, s: ForStmt) -> None:
+    # Check for control flow statements in the loop body
+    detector = ControlFlowDetector()
+    s.body.accept(detector)
+    can_unroll = s.else_body is None and not detector.has_control_flow
+
     def body() -> None:
         builder.accept(s.body)
 
     def else_block() -> None:
         assert s.else_body is not None
         builder.accept(s.else_body)
 
-    for_loop_helper(
-        builder, s.index, s.expr, body, else_block if s.else_body else None, s.is_async, s.line
-    )
+    while True:
+        try:
+            # for whatever reason this can blow up with can_unroll=True
+            # but in those cases we can just fall back to the existing for loop logic.
+            for_loop_helper(
+                builder,
+                s.index,
+                s.expr,
+                body,
+                else_block if s.else_body else None,
+                s.is_async,
+                s.line,
+                can_unroll=can_unroll,
+            )
+            return
+        except AssertionError:
+            if not can_unroll:
+                raise
+            can_unroll = False
 
 
 def transform_break_stmt(builder: IRBuilder, node: BreakStmt) -> None:

mypyc/test-data/irbuild-basic.test

@@ -3626,22 +3626,32 @@ def f(s: str) -> str:
     return out
 [out]
 def f(s):
-    s :: str
+    s, r0, out :: str
+    r1, r2 :: native_int
+    r3, r4 :: bit
+    r5, c, r6 :: str
+    r7 :: native_int
     out, c :: str
-    r0 :: int
-    L0:
-        out = ''
-        r0 = 0
-    L1:
-        r1 = r0 < CPyStr_Size(s)
-        if r1 goto L2 else goto L4 :: bool
-    L2:
-        c = CPyStr_GetItemUnsafe(s, r0)
-        out = CPyStr_Append(out, c)
-        r0 = r0 + 1
-        goto L1
-    L4:
-        return out
+L0:
+    r0 = ''
+    out = r0
+    r1 = 0
+L1:
+    r2 = CPyStr_Size_size_t(s)
+    r3 = r2 >= 0 :: signed
+    r4 = r1 < r2 :: signed
+    if r4 goto L2 else goto L2 :: bool
+L2:
+    r5 = CPyStr_GetItemUnsafe(s, r1)
+    c = r5
+    r6 = CPyStr_Append(out, c)
+    out = r6
+L3:
+    r7 = r1 + 1
+    r1 = r7
+    goto L1
+L4:
+    return out
 
 [case TestForOverCompledTupleExpr]
 def f() -> None:
@@ -3650,4 +3660,27 @@ def f() -> None:
         y = x
 [out]
 def f():
-    abc :: tuple[int, int, str]
+    r0 :: str
+    r1, abc :: tuple[int, int, str]
+    r2 :: int
+    r3 :: object
+    x, y :: union[int, str]
+    r4 :: int
+    r5 :: object
+    r6 :: str
+L0:
+    r0 = CPyTagged_Str(6)
+    r1 = (2, 4, r0)
+    abc = r1
+    r2 = abc[0]
+    r3 = box(int, r2)
+    x = r3
+    y = x
+    r4 = abc[1]
+    r5 = box(int, r4)
+    x = r5
+    y = x
+    r6 = abc[2]
+    x = r6
+    y = x
+    return 1

mypyc/test-data/irbuild-set.test

@@ -737,6 +737,7 @@ def not_precomputed() -> None:
 def precomputed():
     r0 :: str
     _ :: object
+    r1, r2 :: str
 L0:
     r0 = 'None'
     _ = r0