[mypyc] Evaluate args in correct order when builtin takes a genexpr

mypyc/irbuild/for_helpers.py

@@ -106,6 +106,7 @@ def for_loop_helper(
     else_insts: GenFunc | None,
     is_async: bool,
     line: int,
+    after_init: GenFunc | None = None,
 ) -> None:
     """Generate IR for a loop.
 
@@ -114,6 +115,8 @@ def for_loop_helper(
         expr: the expression to iterate over
         body_insts: a function that generates the body of the loop
         else_insts: a function that generates the else block instructions
+        after_init: a function that generates instructions to run after the iterable has
+            been evaluated (and iter() called on it) but before the first iteration
     """
     # Body of the loop
     body_block = BasicBlock()
@@ -130,6 +133,8 @@ def for_loop_helper(
     for_gen = make_for_loop_generator(
         builder, index, expr, body_block, normal_loop_exit, line, is_async=is_async
     )
+    if after_init is not None:
+        after_init()
 
     builder.push_loop_stack(step_block, exit_block)
     condition_block = BasicBlock()
@@ -223,6 +228,7 @@ def sequence_from_generator_preallocate_helper(
     gen: GeneratorExpr,
     empty_op_llbuilder: Callable[[Value, int], Value],
     set_item_op: Callable[[Value, Value, Value, int], None],
+    after_first_init: GenFunc | None = None,
 ) -> Value | None:
     """Generate a new tuple or list from a simple generator expression.
 
@@ -245,6 +251,8 @@ def sequence_from_generator_preallocate_helper(
                 - New Value
             See `new_list_set_item_op` and `new_tuple_set_item_op` for detailed
             implementation.
+        after_first_init: Optional function that generates instructions to run after the
+            sequence has been evaluated (see comprehension_helper).
     """
     if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0:
         line = gen.line
@@ -289,6 +297,9 @@ def sequence_from_generator_preallocate_helper(
         else:
             sequence = builder.accept(sequence_expr)
 
+        if after_first_init is not None:
+            after_first_init()
+
         length = get_expr_length_value(builder, sequence_expr, sequence, line, use_pyssize_t=True)
 
         target_op = empty_op_llbuilder(length, line)
@@ -306,8 +317,16 @@ def set_item(item_index: Value) -> None:
     return None
 
 
-def translate_list_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value:
+def translate_list_comprehension(
+    builder: IRBuilder, gen: GeneratorExpr, *, after_first_init: GenFunc | None = None
+) -> Value:
+    """Generate a list from a comprehension or generator expression.
+
+    See comprehension_helper for the meaning of after_first_init.
+    """
     if raise_error_if_contains_unreachable_names(builder, gen):
+        if after_first_init is not None:
+            after_first_init()
         return builder.none()
 
     def set_item(x: Value, y: Value, z: Value, line: int) -> None:
@@ -319,6 +338,7 @@ def set_item(x: Value, y: Value, z: Value, line: int) -> None:
         gen,
         empty_op_llbuilder=builder.builder.new_list_op_with_length,
         set_item_op=set_item,
+        after_first_init=after_first_init,
     )
     if val is not None:
         return val
@@ -331,10 +351,45 @@ def gen_inner_stmts() -> None:
         e = builder.accept(gen.left_expr)
         builder.primitive_op(list_append_op, [builder.read(list_ops, gen.line), e], gen.line)
 
-    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+    comprehension_helper(
+        builder, loop_params, gen_inner_stmts, gen.line, after_first_init=after_first_init
+    )
     return builder.read(list_ops, gen.line)
 
 
+def translate_generator_to_list(
+    builder: IRBuilder, gen: GeneratorExpr, *, before_loop: GenFunc | None = None
+) -> Value:
+    """Generate a list from a generator expression whose items are all consumed at once.
+
+    This is equivalent to creating and fully consuming a generator object, when the
+    consumer doesn't run other code while iterating.
+
+    If given, before_loop generates code for the other arguments of the consuming call.
+    CPython evaluates them after the first iterable (when the generator object is created)
+    but before the generator runs, so we do the same.
+    """
+    return translate_list_comprehension(builder, gen, after_first_init=before_loop)
+
+
+def generator_loop_helper(
+    builder: IRBuilder,
+    gen: GeneratorExpr,
+    gen_inner_stmts: GenFunc,
+    *,
+    before_loop: GenFunc | None = None,
+) -> None:
+    """Generate loops that iterate a generator expression inline (no generator object).
+
+    gen_inner_stmts generates code for each item (gen.left_expr). See
+    translate_generator_to_list for the meaning of before_loop.
+    """
+    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
+    comprehension_helper(
+        builder, loop_params, gen_inner_stmts, gen.line, after_first_init=before_loop
+    )
+
+
 def raise_error_if_contains_unreachable_names(
     builder: IRBuilder, gen: GeneratorExpr | DictionaryComprehension
 ) -> bool:
@@ -407,6 +462,7 @@ def comprehension_helper(
     loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]],
     gen_inner_stmts: Callable[[], None],
     line: int,
+    after_first_init: GenFunc | None = None,
 ) -> None:
     """Helper function for list comprehensions.
 
@@ -417,9 +473,14 @@ def comprehension_helper(
             - "conditions" is a list of conditions, evaluated in order with short-circuiting,
                 that must all be true for the loop body to be executed
         gen_inner_stmts: function to generate the IR for the body of the innermost loop
+        after_first_init: function to generate IR to run after the first (outermost)
+            iterable has been evaluated, but before the first iteration
     """
 
-    def handle_loop(loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]]) -> None:
+    def handle_loop(
+        loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]],
+        after_init: GenFunc | None = None,
+    ) -> None:
         """Generate IR for a loop.
 
         Given a list of (index, expression, [conditions]) tuples, generate IR
@@ -434,6 +495,7 @@ def handle_loop(loop_params: list[tuple[Lvalue, Expression, list[Expression], bo
             None,
             is_async=is_async,
             line=line,
+            after_init=after_init,
         )
 
     def loop_contents(
@@ -468,7 +530,7 @@ def loop_contents(
             with builder.enter_borrow_scope(line):
                 gen_inner_stmts()
 
-    handle_loop(loop_params)
+    handle_loop(loop_params, after_first_init)
 
 
 def is_range_ref(expr: RefExpr) -> bool:

mypyc/irbuild/specialize.py

@@ -94,9 +94,10 @@
 from mypyc.irbuild.constant_fold import constant_fold_expr
 from mypyc.irbuild.for_helpers import (
     comprehension_helper,
+    generator_loop_helper,
     get_expr_length_value,
     sequence_from_generator_preallocate_helper,
-    translate_list_comprehension,
+    translate_generator_to_list,
     translate_set_comprehension,
 )
 from mypyc.irbuild.format_str_tokenizer import (
@@ -561,28 +562,29 @@ def translate_safe_generator_call(
         and expr.arg_kinds[0] == ARG_POS
         and isinstance(expr.args[0], GeneratorExpr)
     ):
-        if isinstance(callee, MemberExpr):
+        receiver = builder.accept(callee.expr) if isinstance(callee, MemberExpr) else None
+        other_args: list[Value] = []
+
+        def gen_other_args() -> None:
+            other_args.extend(builder.accept(arg) for arg in expr.args[1:])
+
+        items = translate_generator_to_list(
+            builder, expr.args[0], before_loop=gen_other_args if len(expr.args) > 1 else None
+        )
+        args = [items] + other_args
+        if receiver is not None:
+            assert isinstance(callee, MemberExpr)
             return builder.gen_method_call(
-                builder.accept(callee.expr),
+                receiver,
                 callee.name,
-                (
-                    [translate_list_comprehension(builder, expr.args[0])]
-                    + [builder.accept(arg) for arg in expr.args[1:]]
-                ),
+                args,
                 builder.node_type(expr),
                 expr.line,
                 expr.arg_kinds,
                 expr.arg_names,
             )
         else:
-            return builder.call_refexpr_with_args(
-                expr,
-                callee,
-                (
-                    [translate_list_comprehension(builder, expr.args[0])]
-                    + [builder.accept(arg) for arg in expr.args[1:]]
-                ),
-            )
+            return builder.call_refexpr_with_args(expr, callee, args)
     return None
 
 
@@ -667,20 +669,18 @@ def translate_sum_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> V
     gen_expr = expr.args[0]
     target_type = builder.node_type(expr)
     retval = Register(target_type, line=expr.line)
-    builder.assign(
-        retval, builder.coerce(builder.accept(start_expr), target_type, expr.line), expr.line
-    )
+
+    def gen_start() -> None:
+        start = builder.coerce(builder.accept(start_expr), target_type, expr.line)
+        builder.assign(retval, start, expr.line)
 
     def gen_inner_stmts() -> None:
         call_expr = builder.accept(gen_expr.left_expr)
         builder.assign(
             retval, builder.binary_op(retval, call_expr, "+", call_expr.line), call_expr.line
         )
 
-    loop_params = list(
-        zip(gen_expr.indices, gen_expr.sequences, gen_expr.condlists, gen_expr.is_async)
-    )
-    comprehension_helper(builder, loop_params, gen_inner_stmts, gen_expr.line)
+    generator_loop_helper(builder, gen_expr, gen_inner_stmts, before_loop=gen_start)
 
     return retval
 
@@ -720,17 +720,22 @@ def translate_next_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
 
     gen = expr.args[0]
     retval = Register(builder.node_type(expr))
-    default_val = builder.accept(expr.args[1]) if len(expr.args) > 1 else None
+    default_val: Value | None = None
     exit_block = BasicBlock()
 
+    def gen_default() -> None:
+        nonlocal default_val
+        default_val = builder.accept(expr.args[1])
+
     def gen_inner_stmts() -> None:
         # next takes the first element of the generator, so if
         # something gets produced, we are done.
         builder.assign(retval, builder.accept(gen.left_expr), gen.left_expr.line)
         builder.goto(exit_block)
 
-    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
-    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+    generator_loop_helper(
+        builder, gen, gen_inner_stmts, before_loop=gen_default if len(expr.args) > 1 else None
+    )
 
     # Now we need the case for when nothing got hit. If there was
     # a default value, we produce it, and otherwise we raise

mypyc/test-data/fixtures/ir.py

@@ -442,10 +442,22 @@ def pow(base: __SupportsPow2[T_contra, T_co], exp: T_contra, mod: None = None) -
 def pow(base: __SupportsPow3NoneOnly[T_contra, T_co], exp: T_contra, mod: None = None) -> T_co: ...
 @overload
 def pow(base: __SupportsPow3[T_contra, _M, T_co], exp: T_contra, mod: _M) -> T_co: ...
-def sorted(iterable: Iterable[_T]) -> list[_T]: ...
+def sorted(
+    iterable: Iterable[_T], *, key: Callable[[_T], Any] | None = None, reverse: bool = False
+) -> list[_T]: ...
 def exit() -> None: ...
+@overload
 def min(x: _T, y: _T) -> _T: ...
+@overload
+def min(
+    iterable: Iterable[_T], *, key: Callable[[_T], Any] | None = None, default: _T = ...
+) -> _T: ...
+@overload
 def max(x: _T, y: _T) -> _T: ...
+@overload
+def max(
+    iterable: Iterable[_T], *, key: Callable[[_T], Any] | None = None, default: _T = ...
+) -> _T: ...
 def repr(o: object) -> str: ...
 def ascii(o: object) -> str: ...
 def ord(o: object) -> int: ...

mypyc/test-data/irbuild-basic.test

@@ -2944,9 +2944,9 @@ def call_sum(l: Iterable[int], comparison: Callable[[int], bool]) -> int:
 
 [out]
 def call_sum(l, comparison):
-    l, comparison :: object
-    r0 :: int
-    r1, r2 :: object
+    l, comparison, r0 :: object
+    r1 :: int
+    r2 :: object
     r3, x :: int
     r4 :: object
     r5 :: object[1]
@@ -2956,10 +2956,10 @@ def call_sum(l, comparison):
     r10, r11 :: int
     r12 :: bit
 L0:
-    r0 = 0
-    r1 = PyObject_GetIter(l)
+    r0 = PyObject_GetIter(l)
+    r1 = 0
 L1:
-    r2 = PyIter_Next(r1)
+    r2 = PyIter_Next(r0)
     if is_error(r2) goto L4 else goto L2
 L2:
     r3 = unbox(int, r2)
@@ -2972,14 +2972,14 @@ L2:
     r8 = unbox(bool, r7)
     r9 = r8 << 1
     r10 = extend r9: builtins.bool to builtins.int
-    r11 = CPyTagged_Add(r0, r10)
-    r0 = r11
+    r11 = CPyTagged_Add(r1, r10)
+    r1 = r11
 L3:
     goto L1
 L4:
     r12 = CPy_NoErrOccurred()
 L5:
-    return r0
+    return r1
 
 [case testSetAttr1]
 from typing import Any, Dict, List