[mypyc] feat: ForMap generator helper for builtins.map

mypyc/doc/native_operations.rst

@@ -54,3 +54,4 @@ These variants of statements have custom implementations:
 * ``for ... in seq:`` (for loop over a sequence)
 * ``for ... in enumerate(...):``
 * ``for ... in zip(...):``
+* ``for ... in map(...)``

mypyc/irbuild/for_helpers.py

@@ -11,6 +11,7 @@
 
 from mypy.nodes import (
     ARG_POS,
+    LDEF,
     CallExpr,
     DictionaryComprehension,
     Expression,
@@ -22,6 +23,7 @@
     SetExpr,
     TupleExpr,
     TypeAlias,
+    Var,
 )
 from mypyc.ir.ops import (
     ERR_NEVER,
@@ -491,6 +493,16 @@ def make_for_loop_generator(
             for_list = ForSequence(builder, index, body_block, loop_exit, line, nested)
             for_list.init(expr_reg, target_type, reverse=True)
             return for_list
+
+        elif (
+            expr.callee.fullname == "builtins.map"
+            and len(expr.args) >= 2
+            and all(k == ARG_POS for k in expr.arg_kinds)
+        ):
+            for_map = ForMap(builder, index, body_block, loop_exit, line, nested)
+            for_map.init(expr.args[0], expr.args[1:])
+            return for_map
+
     if isinstance(expr, CallExpr) and isinstance(expr.callee, MemberExpr) and not expr.args:
         # Special cases for dictionary iterator methods, like dict.items().
         rtype = builder.node_type(expr.callee.expr)
@@ -1166,3 +1178,73 @@ def gen_step(self) -> None:
     def gen_cleanup(self) -> None:
         for gen in self.gens:
             gen.gen_cleanup()
+
+
+class ForMap(ForGenerator):
+    """Generate optimized IR for a for loop over map(f, ...)."""
+
+    def need_cleanup(self) -> bool:
+        # The wrapped for loops might need cleanup. We might generate a
+        # redundant cleanup block, but that's okay.
+        return True
+
+    def init(self, func: Expression, exprs: list[Expression]) -> None:
+        self.func_expr = func
+        self.func = self.builder.accept(func)
+        self.exprs = exprs
+        self.cond_blocks = [BasicBlock() for _ in range(len(exprs) - 1)] + [self.body_block]
+
+        self.gens: list[ForGenerator] = []
+        for i, iterable_expr in enumerate(exprs):
+            argname = f"_mypyc_map_arg_{i}"
+            var_type = self.builder._analyze_iterable_item_type(iterable_expr)
+            name_expr = NameExpr(argname)
+            name_expr.kind = LDEF
+            name_expr.node = Var(argname, var_type)
+            self.builder.add_local_reg(name_expr.node, self.builder.type_to_rtype(var_type))
+            self.gens.append(
+                make_for_loop_generator(
+                    self.builder,
+                    name_expr,
+                    iterable_expr,
+                    self.cond_blocks[i],
+                    self.loop_exit,
+                    self.line,
+                    is_async=False,
+                    nested=True,
+                )
+            )
+
+    def gen_condition(self) -> None:
+        for i, gen in enumerate(self.gens):
+            gen.gen_condition()
+            if i < len(self.gens) - 1:
+                self.builder.activate_block(self.cond_blocks[i])
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+
+        for gen in self.gens:
+            gen.begin_body()
+
+        # This goes here to prevent a circular import
+        from mypyc.irbuild.expression import transform_call_expr
+
+        call_expr = CallExpr(
+            self.func_expr,
+            [gen.index for gen in self.gens],
+            [ARG_POS] * len(self.gens),
+            [None] * len(self.gens),
+        )
+
+        result = transform_call_expr(builder, call_expr)
+        builder.assign(builder.get_assignment_target(self.index), result, line)
+
+    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()

mypyc/test-data/fixtures/ir.py

@@ -7,6 +7,8 @@
     overload, Mapping, Union, Callable, Sequence, FrozenSet, Protocol
 )
 
+from typing_extensions import Self
+
 _T = TypeVar('_T')
 T_co = TypeVar('T_co', covariant=True)
 T_contra = TypeVar('T_contra', contravariant=True)
@@ -406,3 +408,25 @@ class classmethod: pass
 class staticmethod: pass
 
 NotImplemented: Any = ...
+
+_T1 = TypeVar("_T1")
+_T2 = TypeVar("_T2")
+_T3 = TypeVar("_T3")
+
+class map(Generic[_S]):
+    @overload
+    def __new__(cls, func: Callable[[_T1], _S], iterable: Iterable[_T1], /) -> Self: ...
+    @overload
+    def __new__(cls, func: Callable[[_T1, _T2], _S], iterable: Iterable[_T1], iter2: Iterable[_T2], /) -> Self: ...
+    @overload
+    def __new__(cls, func: Callable[[_T1, _T2, _T3], _S], iterable: Iterable[_T1], iter2: Iterable[_T2], iter3: Iterable[_T3], /) -> Self: ...
+    def __iter__(self) -> Self: ...
+    def __next__(self) -> _S: ...
+
+class filter(Generic[_T]):
+    @overload
+    def __new__(cls, function: None, iterable: Iterable[_T | None], /) -> Self: ...
+    @overload
+    def __new__(cls, function: Callable[[_T], Any], iterable: Iterable[_T], /) -> Self: ...
+    def __iter__(self) -> Self: ...
+    def __next__(self) -> _T: ...