|
24 | 24 | TypeAlias,
|
25 | 25 | )
|
26 | 26 | from mypyc.ir.ops import (
|
| 27 | + ERR_NEVER, |
27 | 28 | BasicBlock,
|
28 | 29 | Branch,
|
29 | 30 | Integer,
|
30 | 31 | IntOp,
|
31 | 32 | LoadAddress,
|
| 33 | + LoadErrorValue, |
32 | 34 | LoadMem,
|
| 35 | + MethodCall, |
33 | 36 | RaiseStandardError,
|
34 | 37 | Register,
|
35 | 38 | TupleGet,
|
36 | 39 | TupleSet,
|
37 | 40 | Value,
|
38 | 41 | )
|
39 | 42 | from mypyc.ir.rtypes import (
|
| 43 | + RInstance, |
40 | 44 | RTuple,
|
41 | 45 | RType,
|
42 | 46 | bool_rprimitive,
|
|
48 | 52 | is_short_int_rprimitive,
|
49 | 53 | is_str_rprimitive,
|
50 | 54 | is_tuple_rprimitive,
|
| 55 | + object_pointer_rprimitive, |
| 56 | + object_rprimitive, |
51 | 57 | pointer_rprimitive,
|
52 | 58 | short_int_rprimitive,
|
53 | 59 | )
|
54 | 60 | from mypyc.irbuild.builder import IRBuilder
|
| 61 | +from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME |
55 | 62 | from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple
|
56 | 63 | from mypyc.primitives.dict_ops import (
|
57 | 64 | dict_check_size_op,
|
|
62 | 69 | dict_next_value_op,
|
63 | 70 | dict_value_iter_op,
|
64 | 71 | )
|
65 |
| -from mypyc.primitives.exc_ops import no_err_occurred_op |
| 72 | +from mypyc.primitives.exc_ops import no_err_occurred_op, propagate_if_error_op |
66 | 73 | from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op
|
67 | 74 | from mypyc.primitives.list_ops import list_append_op, list_get_item_unsafe_op, new_list_set_item_op
|
68 | 75 | from mypyc.primitives.misc_ops import stop_async_iteration_op
|
@@ -511,7 +518,15 @@ def make_for_loop_generator(
|
511 | 518 | # Default to a generic for loop.
|
512 | 519 | if iterable_expr_reg is None:
|
513 | 520 | iterable_expr_reg = builder.accept(expr)
|
514 |
| - for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested) |
| 521 | + |
| 522 | + it = iterable_expr_reg.type |
| 523 | + for_obj: ForNativeGenerator | ForIterable |
| 524 | + if isinstance(it, RInstance) and it.class_ir.has_method(GENERATOR_HELPER_NAME): |
| 525 | + # Directly call generator object methods if iterating over a native generator. |
| 526 | + for_obj = ForNativeGenerator(builder, index, body_block, loop_exit, line, nested) |
| 527 | + else: |
| 528 | + # Generic implementation that works of arbitrary iterables. |
| 529 | + for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested) |
515 | 530 | item_type = builder._analyze_iterable_item_type(expr)
|
516 | 531 | item_rtype = builder.type_to_rtype(item_type)
|
517 | 532 | for_obj.init(iterable_expr_reg, item_rtype)
|
@@ -623,6 +638,63 @@ def gen_cleanup(self) -> None:
|
623 | 638 | self.builder.call_c(no_err_occurred_op, [], self.line)
|
624 | 639 |
|
625 | 640 |
|
| 641 | +class ForNativeGenerator(ForGenerator): |
| 642 | + """Generate IR for a for loop over a native generator.""" |
| 643 | + |
| 644 | + def need_cleanup(self) -> bool: |
| 645 | + # Create a new cleanup block for when the loop is finished. |
| 646 | + return True |
| 647 | + |
| 648 | + def init(self, expr_reg: Value, target_type: RType) -> None: |
| 649 | + # Define target to contains the generator expression. It's also the iterator. |
| 650 | + # If we are inside a generator function, spill these into the environment class. |
| 651 | + builder = self.builder |
| 652 | + self.iter_target = builder.maybe_spill(expr_reg) |
| 653 | + self.target_type = target_type |
| 654 | + |
| 655 | + def gen_condition(self) -> None: |
| 656 | + builder = self.builder |
| 657 | + line = self.line |
| 658 | + self.return_value = Register(object_rprimitive) |
| 659 | + err = builder.add(LoadErrorValue(object_rprimitive, undefines=True)) |
| 660 | + builder.assign(self.return_value, err, line) |
| 661 | + |
| 662 | + # Call generated generator helper method, passing a PyObject ** as the final |
| 663 | + # argument that will be used to store the return value in the return value |
| 664 | + # register. We ignore the return value but the presence of a return value |
| 665 | + # indicates that the generator has finished. This is faster than raising |
| 666 | + # and catching StopIteration, which is the non-native way of doing this. |
| 667 | + ptr = builder.add(LoadAddress(object_pointer_rprimitive, self.return_value)) |
| 668 | + nn = builder.none_object() |
| 669 | + helper_call = MethodCall( |
| 670 | + builder.read(self.iter_target), GENERATOR_HELPER_NAME, [nn, nn, nn, nn, ptr], line |
| 671 | + ) |
| 672 | + # We provide custom handling for error values. |
| 673 | + helper_call.error_kind = ERR_NEVER |
| 674 | + |
| 675 | + self.next_reg = builder.add(helper_call) |
| 676 | + builder.add(Branch(self.next_reg, self.loop_exit, self.body_block, Branch.IS_ERROR)) |
| 677 | + |
| 678 | + def begin_body(self) -> None: |
| 679 | + # Assign the value obtained from the generator helper method to the |
| 680 | + # lvalue so that it can be referenced by code in the body of the loop. |
| 681 | + builder = self.builder |
| 682 | + line = self.line |
| 683 | + # We unbox here so that iterating with tuple unpacking generates a tuple based |
| 684 | + # unpack instead of an iterator based one. |
| 685 | + next_reg = builder.coerce(self.next_reg, self.target_type, line) |
| 686 | + builder.assign(builder.get_assignment_target(self.index), next_reg, line) |
| 687 | + |
| 688 | + def gen_step(self) -> None: |
| 689 | + # Nothing to do here, since we get the next item as part of gen_condition(). |
| 690 | + pass |
| 691 | + |
| 692 | + def gen_cleanup(self) -> None: |
| 693 | + # If return value is NULL (it wasn't assigned to by the generator helper method), |
| 694 | + # an exception was raised that we need to propagate. |
| 695 | + self.builder.primitive_op(propagate_if_error_op, [self.return_value], self.line) |
| 696 | + |
| 697 | + |
626 | 698 | class ForAsyncIterable(ForGenerator):
|
627 | 699 | """Generate IR for an async for loop."""
|
628 | 700 |
|
|
0 commit comments