[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,72 @@ 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
+
+        for gen in self.gens:
+            gen.begin_body()
+
+        call_expr = CallExpr(
+            self.func_expr,
+            [gen.index for gen in self.gens],
+            [ARG_POS] * len(self.gens),
+            [None] * len(self.gens),
+        )
+
+        # TODO: debug redundant box->unbox op in builder.accept and then replace this
+        from mypyc.irbuild.expression import transform_call_expr
+
+        result = transform_call_expr(builder, call_expr)
+        builder.assign(builder.get_assignment_target(self.index), result, self.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/irbuild/specialize.py

@@ -75,6 +75,7 @@
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.for_helpers import (
     comprehension_helper,
+    for_loop_helper,
     sequence_from_generator_preallocate_helper,
     translate_list_comprehension,
     translate_set_comprehension,
@@ -95,7 +96,7 @@
 )
 from mypyc.primitives.float_ops import isinstance_float
 from mypyc.primitives.int_ops import isinstance_int
-from mypyc.primitives.list_ops import isinstance_list, new_list_set_item_op
+from mypyc.primitives.list_ops import isinstance_list, list_append_op, new_list_set_item_op
 from mypyc.primitives.misc_ops import isinstance_bool
 from mypyc.primitives.set_ops import isinstance_frozenset, isinstance_set
 from mypyc.primitives.str_ops import (
@@ -262,7 +263,7 @@ def dict_methods_fast_path(builder: IRBuilder, expr: CallExpr, callee: RefExpr)
 
 
 @specialize_function("builtins.list")
-def translate_list_from_generator_call(
+def translate_list_from_generator_expr(
     builder: IRBuilder, expr: CallExpr, callee: RefExpr
 ) -> Value | None:
     """Special case for simplest list comprehension.
@@ -286,6 +287,50 @@ def translate_list_from_generator_call(
     return None
 
 
+@specialize_function("builtins.list")
+def translate_list_from_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for simplest list construction using one of our for_helpers.
+
+    For example:
+        list(map(f, some_list/some_tuple/some_str))
+    """
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], CallExpr)
+        and isinstance(expr.args[0].callee, NameExpr)
+        and expr.args[0].callee.fullname
+        in (
+            # TODO: make constant for these vals
+            "builtins.map",
+            "builtins.filter",
+            "builtins.filterfalse",
+        )
+    ):
+        call_expr = expr.args[0]
+        itemtype = builder._analyze_iterable_item_type(call_expr)
+        indextype = builder.type_to_rtype(itemtype)
+        index = Register(indextype, "__mypyc_list_helper__", line=call_expr.line)
+
+        result = builder.new_list_op([], expr.line)
+
+        def body_insts() -> None:
+            builder.primitive_op(list_append_op, [result, index], expr.line)
+
+        for_loop_helper(
+            builder=builder,
+            index=index,
+            expr=call_expr,
+            body_insts=body_insts,
+            else_insts=None,
+            is_async=False,
+            line=expr.line,
+        )
+    return None
+
+
 @specialize_function("builtins.tuple")
 def translate_tuple_from_generator_call(
     builder: IRBuilder, expr: CallExpr, callee: RefExpr

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)
@@ -407,3 +409,17 @@ 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: ...