Update for_helpers.py

Author: BobTheBuidler

mypyc/irbuild/for_helpers.py

@@ -1184,13 +1184,36 @@ def init(self, indexes: list[Lvalue], exprs: list[Expression]) -> None:
             )
             self.gens.append(gen)
 
+        self.conditions, self.cond_blocks = self.__sort_conditions()
+
     def gen_condition(self) -> None:
+        for i, gen in enumerate(ordered + leftovers + for_iterable):
+            gen.gen_condition()
+            if i < len(self.gens) - 1:
+                self.builder.activate_block(self.cond_blocks[i])
+
+    def begin_body(self) -> None:
+        for gen in self.gens:
+            gen.begin_body()
+
+    def gen_step(self) -> None:
+        for gen in self.gens:
+            gen.gen_step()
+
+    def gen_cleanup(self) -> None:
+        for gen in self.gens:
+            gen.gen_cleanup()
+
+    def __sort_conditions(self) -> List[ForSequence]:
         # We don't necessarily need to check the gens in order,
         # we just need to know which gen ends first. Some gens
         # are quicker to check than others, so we will check the
         # specialized ForHelpers before we check any generic
         # ForIterable
+        
         gens = self.gens
+        cond_blocks = self.cond_blocks[:]
+        cond_blocks.remove(self.body_block)
 
         def check_type(obj: Any, typ: type[ForGenerator]) -> bool:
             # ForEnumerate gen_condition is as fast as it's underlying generator's
@@ -1201,46 +1224,42 @@ def check_type(obj: Any, typ: type[ForGenerator]) -> bool:
             )
 
         # these are slowest, they invoke Python's iteration protocol
-        for_iterable = [g for g in gens if check_type(g, ForSequence)]
+        for_iterable = [(g, block) for g, block in zip(gens, cond_blocks) if check_type(g, ForSequence)]
 
         # These aren't the slowest but they're slow, we need to pack an RTuple and then get and item and do a comparison
-        for_dict = [g for g in gens if check_type(g, ForDictionaryCommon)]
+        for_dict = [(g, block) for g, block in zip(gens, cond_blocks) if check_type(g, ForDictionaryCommon)]
 
         # These are faster than ForIterable but not as fast as others (faster than ForDict?)
-        for_native = [g for g in gens if check_type(g, ForNativeGenerator)]
+        for_native = [(g, block) for g, block in zip(gens, cond_blocks) if check_type(g, ForNativeGenerator)]
 
         # forward involves in the best case one pyssize_t comparison, else one length check + the comparison
         # reverse is slightly slower than forward, with one extra check
         for_sequence_reverse_with_len_check = [
-            g
-            for g in gens
+            (g, block) for g, block in zip(gens, cond_blocks)
             if check_type(g, ForSequence)
             and (
                 for_seq := cast(ForSequence, g.main_gen if isinstance(g, ForEnumerate) else g)
             ).reverse
             and for_seq.length_reg is not None
         ]
         for_sequence_reverse_no_len_check = [
-            g
-            for g in gens
+            (g, block) for g, block in zip(gens, cond_blocks)
             if check_type(g, ForSequence)
             and (
                 for_seq := cast(ForSequence, g.main_gen if isinstance(g, ForEnumerate) else g)
             ).reverse
             and for_seq.length_reg is None
         ]
         for_sequence_forward_with_len_check = [
-            g
-            for g in gens
+            (g, block) for g, block in zip(gens, cond_blocks)
             if check_type(g, ForSequence)
             and not (
                 for_seq := cast(ForSequence, g.main_gen if isinstance(g, ForEnumerate) else g)
             ).reverse
             and for_seq.length_reg is not None
         ]
         for_sequence_forward_no_len_check = [
-            g
-            for g in gens
+            (g, block) for g, block in zip(gens, cond_blocks)
             if check_type(g, ForSequence)
             and not (
                 for_seq := cast(ForSequence, g.main_gen if isinstance(g, ForEnumerate) else g)
@@ -1249,7 +1268,7 @@ def check_type(obj: Any, typ: type[ForGenerator]) -> bool:
         ]
 
         # these are really fast, just a C int equality check
-        for_range = [g for g in gens if isinstance(g, ForRange)]
+        for_range = [(g, block) for g, block in zip(gens, cond_blocks) if isinstance(g, ForRange)]
 
         ordered = (
             for_range
@@ -1262,24 +1281,13 @@ def check_type(obj: Any, typ: type[ForGenerator]) -> bool:
         )
 
         # this is a failsafe for ForHelper classes which might have been added after this commit but not added to this function's code
-        leftovers = [g for g in gens if g not in ordered + for_iterable]
-
-        for i, gen in enumerate(ordered + leftovers + for_iterable):
-            gen.gen_condition()
-            if i < len(self.gens) - 1:
-                self.builder.activate_block(self.cond_blocks[i])
-
-    def begin_body(self) -> None:
-        for gen in self.gens:
-            gen.begin_body()
+        leftovers = [(g, block) for g, block in zip(gens, cond_blocks) if g not in ordered + for_iterable]
 
-    def gen_step(self) -> None:
-        for gen in self.gens:
-            gen.gen_step()
+        gens_and_blocks = ordered + leftovers + for_iterable
+        conditons = [g for (g, block) in gens_and_blocks]
+        cond_blocks = [block for (g, block) in gens_and_blocks] + [self.body_block]
 
-    def gen_cleanup(self) -> None:
-        for gen in self.gens:
-            gen.gen_cleanup()
+        return conditions, cond_blocks 
 
 
 def get_expr_length(builder: IRBuilder, expr: Expression) -> int | None: