[mypyc] feat: ForFilter generator helper for builtins.filter

mypyc/irbuild/for_helpers.py

@@ -490,6 +490,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.filter"
+            and len(expr.args) == 2
+            and all(k == ARG_POS for k in expr.arg_kinds)
+        ):
+            for_filter = ForFilter(builder, index, body_block, loop_exit, line, nested)
+            for_filter.init(index, expr.args[0], expr.args[1])
+            return for_filter
+
     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)
@@ -1147,3 +1157,56 @@ def gen_step(self) -> None:
     def gen_cleanup(self) -> None:
         for gen in self.gens:
             gen.gen_cleanup()
+
+
+class ForFilter(ForGenerator):
+    """Generate optimized IR for a for loop over filter(f, iterable)."""
+
+    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, index: Lvalue, func: Expression, iterable: Expression) -> None:
+        self.filter_func = self.builder.accept(func)
+        self.iterable = iterable
+        self.index = index
+
+        self.gen = make_for_loop_generator(
+            self.builder,
+            self.index,
+            self.iterable,
+            self.body_block,
+            self.loop_exit,
+            self.line,
+            is_async=False,
+            nested=True,
+        )
+
+    def gen_condition(self) -> None:
+        self.gen.gen_condition()
+
+    def begin_body(self) -> None:
+        # 1. Assign the next item to the loop variable
+        self.gen.begin_body()
+
+        # 2. Call the filter function
+        builder = self.builder
+        line = self.line
+        item = builder.read(builder.get_assignment_target(self.index), line)
+        # TODO: implement logic to handle c calls of native functions
+        result = builder.py_call(self.filter_func, [item], line)
+
+        # Now, filter: only enter the body if func(item) is truthy
+        cont_block, rest_block = BasicBlock(), BasicBlock()
+        builder.add_bool_branch(result, rest_block, cont_block)
+        builder.activate_block(cont_block)
+        builder.nonlocal_control[-1].gen_continue(builder, line)
+        builder.goto_and_activate(rest_block)
+        # At this point, the rest of the loop body (user code) will be emitted
+
+    def gen_step(self) -> None:
+        self.gen.gen_step()
+
+    def gen_cleanup(self) -> None:
+        self.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)
@@ -405,3 +407,22 @@ class classmethod: pass
 class staticmethod: pass
 
 NotImplemented: Any = ...
+
+_T1 = TypeVar("_T1")
+_T2 = TypeVar("_T2")
+
+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: ...
+    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: ...

mypyc/test-data/irbuild-basic.test

@@ -3546,3 +3546,74 @@ L0:
     r2 = PyObject_Vectorcall(r1, 0, 0, 0)
     r3 = box(None, 1)
     return r3
+
+[case testForFilter]
+def f(x: int) -> bool:
+    return bool(x % 2)
+def g(a: list[int]) -> int:
+    s = 0
+    for x in filter(f, a):
+        s += x
+    return s
+[out]
+def f(x):
+    x, r0 :: int
+    r1 :: bit
+L0:
+    r0 = CPyTagged_Remainder(x, 4)
+    r1 = r0 != 0
+    return r1
+def g(a):
+    a :: list
+    s :: int
+    r0 :: dict
+    r1 :: str
+    r2 :: object
+    r3, r4 :: native_int
+    r5 :: bit
+    r6 :: object
+    r7, x :: int
+    r8 :: object
+    r9 :: object[1]
+    r10 :: object_ptr
+    r11 :: object
+    r12 :: i32
+    r13 :: bit
+    r14 :: bool
+    r15 :: int
+    r16 :: native_int
+L0:
+    s = 0
+    r0 = __main__.globals :: static
+    r1 = 'f'
+    r2 = CPyDict_GetItem(r0, r1)
+    r3 = 0
+L1:
+    r4 = var_object_size a
+    r5 = r3 < r4 :: signed
+    if r5 goto L2 else goto L6 :: bool
+L2:
+    r6 = list_get_item_unsafe a, r3
+    r7 = unbox(int, r6)
+    x = r7
+    r8 = box(int, x)
+    r9 = [r8]
+    r10 = load_address r9
+    r11 = PyObject_Vectorcall(r2, r10, 1, 0)
+    keep_alive r8
+    r12 = PyObject_IsTrue(r11)
+    r13 = r12 >= 0 :: signed
+    r14 = truncate r12: i32 to builtins.bool
+    if r14 goto L4 else goto L3 :: bool
+L3:
+    goto L5
+L4:
+    r15 = CPyTagged_Add(s, x)
+    s = r15
+L5:
+    r16 = r3 + 1
+    r3 = r16
+    goto L1
+L6:
+L7:
+    return s

mypyc/test-data/run-loops.test

@@ -571,3 +571,37 @@ print([x for x in native.Vector2(4, -5.2)])
 [out]
 Vector2(x=-2, y=3.1)
 \[4, -5.2]
+
+[case testRunForFilter]
+def f(a: list[int]) -> int:
+    s = 0
+    for x in filter(lambda x: x % 2 == 0, a):
+        s += x
+    return s
+
+[file driver.py]
+from native import f
+print(f([1, 2, 3, 4, 5, 6]))
+print(f([1, 3, 5]))
+print(f([]))
+
+[out]
+12
+0
+0
+
+[case testRunForFilterEdgeCases]
+def f(a: list[int]) -> int:
+    s = 0
+    for x in filter(lambda x: x > 10, a):
+        s += x
+    return s
+
+[file driver.py]
+from native import f
+print(f([5, 15, 25]))
+print(f([]))
+
+[out]
+40
+0