[mypyc] Speed up for loop over native generator

Author: JukkaL

mypyc/irbuild/for_helpers.py

@@ -24,19 +24,23 @@
     TypeAlias,
 )
 from mypyc.ir.ops import (
+    ERR_NEVER,
     BasicBlock,
     Branch,
     Integer,
     IntOp,
     LoadAddress,
+    LoadErrorValue,
     LoadMem,
+    MethodCall,
     RaiseStandardError,
     Register,
     TupleGet,
     TupleSet,
     Value,
 )
 from mypyc.ir.rtypes import (
+    RInstance,
     RTuple,
     RType,
     bool_rprimitive,
@@ -48,6 +52,8 @@
     is_short_int_rprimitive,
     is_str_rprimitive,
     is_tuple_rprimitive,
+    object_pointer_rprimitive,
+    object_rprimitive,
     pointer_rprimitive,
     short_int_rprimitive,
 )
@@ -62,7 +68,7 @@
     dict_next_value_op,
     dict_value_iter_op,
 )
-from mypyc.primitives.exc_ops import no_err_occurred_op
+from mypyc.primitives.exc_ops import no_err_occurred_op, propagate_if_error_op
 from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op
 from mypyc.primitives.list_ops import list_append_op, list_get_item_unsafe_op, new_list_set_item_op
 from mypyc.primitives.misc_ops import stop_async_iteration_op
@@ -511,7 +517,16 @@ def make_for_loop_generator(
     # Default to a generic for loop.
     if iterable_expr_reg is None:
         iterable_expr_reg = builder.accept(expr)
-    for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested)
+
+    helper_method = "__mypyc_generator_helper__"
+    it = iterable_expr_reg.type
+    for_obj: ForNativeGenerator | ForIterable
+    if isinstance(it, RInstance) and it.class_ir.has_method(helper_method):
+        # Directly call generator object methods if iterating over a native generator.
+        for_obj = ForNativeGenerator(builder, index, body_block, loop_exit, line, nested)
+    else:
+        # Generic implementation that works of arbitrary iterables.
+        for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested)
     item_type = builder._analyze_iterable_item_type(expr)
     item_rtype = builder.type_to_rtype(item_type)
     for_obj.init(iterable_expr_reg, item_rtype)
@@ -623,6 +638,56 @@ def gen_cleanup(self) -> None:
         self.builder.call_c(no_err_occurred_op, [], self.line)
 
 
+class ForNativeGenerator(ForGenerator):
+    """Generate IR for a for loop over a native generator."""
+
+    def need_cleanup(self) -> bool:
+        # Create a new cleanup block for when the loop is finished.
+        return True
+
+    def init(self, expr_reg: Value, target_type: RType) -> None:
+        # Define targets to contain the expression, along with the iterator that will be used
+        # for the for-loop. If we are inside of a generator function, spill these into the
+        # environment class.
+        builder = self.builder
+        self.iter_target = builder.maybe_spill(expr_reg)
+        self.target_type = target_type
+
+    def gen_condition(self) -> None:
+        builder = self.builder
+        line = self.line
+        helper_method = "__mypyc_generator_helper__"
+        self.return_value = Register(object_rprimitive)
+        err = builder.add(LoadErrorValue(object_rprimitive, undefines=True))
+        builder.assign(self.return_value, err, line)
+        ptr = builder.add(LoadAddress(object_pointer_rprimitive, self.return_value))
+        nn = builder.none_object()
+        helper_call =(
+            MethodCall(builder.read(self.iter_target), helper_method, [nn, nn, nn, nn, ptr], line)
+        )
+        # We provide custom handling for error values.
+        helper_call.error_kind = ERR_NEVER
+        self.next_reg = builder.add(helper_call)
+        builder.add(Branch(self.next_reg, self.loop_exit, self.body_block, Branch.IS_ERROR))
+
+    def begin_body(self) -> None:
+        # Assign the value obtained from __next__ to the
+        # lvalue so that it can be referenced by code in the body of the loop.
+        builder = self.builder
+        line = self.line
+        # We unbox here so that iterating with tuple unpacking generates a tuple based
+        # unpack instead of an iterator based one.
+        next_reg = builder.coerce(self.next_reg, self.target_type, line)
+        builder.assign(builder.get_assignment_target(self.index), next_reg, line)
+
+    def gen_step(self) -> None:
+        # Nothing to do here, since we get the next item as part of gen_condition().
+        pass
+
+    def gen_cleanup(self) -> None:
+        self.builder.primitive_op(propagate_if_error_op, [self.return_value], self.line)
+
+
 class ForAsyncIterable(ForGenerator):
     """Generate IR for an async for loop."""