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

Author: BobTheBuidler

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,45 @@ 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:
+        builder = self.builder
+        line = self.line
+        # First, get the next item from the sub-generator
+        self.gen.gen_condition()
+        # Now, filter: only enter the body if func(item) is truthy
+        filter_block = BasicBlock()
+        builder.activate_block(filter_block)
+        self.gen.begin_body()
+        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)
+        builder.add_bool_branch(result, self.body_block, self.loop_exit)
+
+    def begin_body(self) -> None:
+        # The item is already assigned to self.index by the sub-generator.
+        pass
+
+    def gen_step(self) -> None:
+        self.gen.gen_step()
+
+    def gen_cleanup(self) -> None:
+        self.gen.gen_cleanup()

mypyc/test-data/irbuild-basic.test

@@ -3546,3 +3546,73 @@ 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 L3 else goto L5 :: 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 L3 else goto L5 :: bool
+L3:
+    r15 = CPyTagged_Add(s, x)
+    s = r15
+L4:
+    r16 = r3 + 1
+    r3 = r16
+    goto L1
+L5:
+L6:
+    return s
+

mypyc/test-data/run-loops.test

@@ -571,3 +571,33 @@ 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
+
+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
+
+print(f([5, 15, 25]))
+print(f([]))
+
+[out]
+40
+0