[mypyc] Speed up for loop over native generator

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,10 +52,13 @@
     is_short_int_rprimitive,
     is_str_rprimitive,
     is_tuple_rprimitive,
+    object_pointer_rprimitive,
+    object_rprimitive,
     pointer_rprimitive,
     short_int_rprimitive,
 )
 from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple
 from mypyc.primitives.dict_ops import (
     dict_check_size_op,
@@ -62,7 +69,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 +518,15 @@ 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)
+
+    it = iterable_expr_reg.type
+    for_obj: ForNativeGenerator | ForIterable
+    if isinstance(it, RInstance) and it.class_ir.has_method(GENERATOR_HELPER_NAME):
+        # 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,63 @@ 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 target to contains the generator expression. It's also the iterator.
+        # If we are inside 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
+        self.return_value = Register(object_rprimitive)
+        err = builder.add(LoadErrorValue(object_rprimitive, undefines=True))
+        builder.assign(self.return_value, err, line)
+
+        # Call generated generator helper method, passing a PyObject ** as the final
+        # argument that will be used to store the return value in the return value
+        # register. We ignore the return value but the presence of a return value
+        # indicates that the generator has finished. This is faster than raising
+        # and catching StopIteration, which is the non-native way of doing this.
+        ptr = builder.add(LoadAddress(object_pointer_rprimitive, self.return_value))
+        nn = builder.none_object()
+        helper_call = MethodCall(
+            builder.read(self.iter_target), GENERATOR_HELPER_NAME, [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 the generator helper method 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:
+        # If return value is NULL (it wasn't assigned to by the generator helper method),
+        # an exception was raised that we need to propagate.
+        self.builder.primitive_op(propagate_if_error_op, [self.return_value], self.line)
+
+
 class ForAsyncIterable(ForGenerator):
     """Generate IR for an async for loop."""
 

mypyc/irbuild/prepare.py

@@ -15,7 +15,7 @@
 
 from collections import defaultdict
 from collections.abc import Iterable
-from typing import NamedTuple
+from typing import Final, NamedTuple
 
 from mypy.build import Graph
 from mypy.nodes import (
@@ -71,7 +71,7 @@
 from mypyc.options import CompilerOptions
 from mypyc.sametype import is_same_type
 
-GENERATOR_HELPER_NAME = "__mypyc_generator_helper__"
+GENERATOR_HELPER_NAME: Final = "__mypyc_generator_helper__"
 
 
 def build_type_map(