mypy-dev 1.9.0.dev0__py3-none-any.whl

mypyc/irbuild/for_helpers.py

@@ -0,0 +1,1046 @@
+"""Helpers for generating for loops and comprehensions.
+
+We special case certain kinds for loops such as "for x in range(...)"
+for better efficiency.  Each for loop generator class below deals one
+such special case.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, ClassVar
+
+from mypy.nodes import (
+    ARG_POS,
+    CallExpr,
+    Expression,
+    GeneratorExpr,
+    Lvalue,
+    MemberExpr,
+    RefExpr,
+    SetExpr,
+    TupleExpr,
+    TypeAlias,
+)
+from mypyc.ir.ops import (
+    BasicBlock,
+    Branch,
+    Integer,
+    IntOp,
+    LoadAddress,
+    LoadMem,
+    Register,
+    TupleGet,
+    TupleSet,
+    Value,
+)
+from mypyc.ir.rtypes import (
+    RTuple,
+    RType,
+    bool_rprimitive,
+    int_rprimitive,
+    is_dict_rprimitive,
+    is_fixed_width_rtype,
+    is_list_rprimitive,
+    is_sequence_rprimitive,
+    is_short_int_rprimitive,
+    is_str_rprimitive,
+    is_tuple_rprimitive,
+    pointer_rprimitive,
+    short_int_rprimitive,
+)
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple
+from mypyc.primitives.dict_ops import (
+    dict_check_size_op,
+    dict_item_iter_op,
+    dict_key_iter_op,
+    dict_next_item_op,
+    dict_next_key_op,
+    dict_next_value_op,
+    dict_value_iter_op,
+)
+from mypyc.primitives.exc_ops import no_err_occurred_op
+from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op
+from mypyc.primitives.list_ops import list_append_op, list_get_item_unsafe_op, new_list_set_item_op
+from mypyc.primitives.misc_ops import stop_async_iteration_op
+from mypyc.primitives.registry import CFunctionDescription
+from mypyc.primitives.set_ops import set_add_op
+
+GenFunc = Callable[[], None]
+
+
+def for_loop_helper(
+    builder: IRBuilder,
+    index: Lvalue,
+    expr: Expression,
+    body_insts: GenFunc,
+    else_insts: GenFunc | None,
+    is_async: bool,
+    line: int,
+) -> None:
+    """Generate IR for a loop.
+
+    Args:
+        index: the loop index Lvalue
+        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
+    """
+    # Body of the loop
+    body_block = BasicBlock()
+    # Block that steps to the next item
+    step_block = BasicBlock()
+    # Block for the else clause, if we need it
+    else_block = BasicBlock()
+    # Block executed after the loop
+    exit_block = BasicBlock()
+
+    # Determine where we want to exit, if our condition check fails.
+    normal_loop_exit = else_block if else_insts is not None else exit_block
+
+    for_gen = make_for_loop_generator(
+        builder, index, expr, body_block, normal_loop_exit, line, is_async=is_async
+    )
+
+    builder.push_loop_stack(step_block, exit_block)
+    condition_block = BasicBlock()
+    builder.goto_and_activate(condition_block)
+
+    # Add loop condition check.
+    for_gen.gen_condition()
+
+    # Generate loop body.
+    builder.activate_block(body_block)
+    for_gen.begin_body()
+    body_insts()
+
+    # We generate a separate step block (which might be empty).
+    builder.goto_and_activate(step_block)
+    for_gen.gen_step()
+    # Go back to loop condition.
+    builder.goto(condition_block)
+
+    for_gen.add_cleanup(normal_loop_exit)
+    builder.pop_loop_stack()
+
+    if else_insts is not None:
+        builder.activate_block(else_block)
+        else_insts()
+        builder.goto(exit_block)
+
+    builder.activate_block(exit_block)
+
+
+def for_loop_helper_with_index(
+    builder: IRBuilder,
+    index: Lvalue,
+    expr: Expression,
+    expr_reg: Value,
+    body_insts: Callable[[Value], None],
+    line: int,
+) -> None:
+    """Generate IR for a sequence iteration.
+
+    This function only works for sequence type. Compared to for_loop_helper,
+    it would feed iteration index to body_insts.
+
+    Args:
+        index: the loop index Lvalue
+        expr: the expression to iterate over
+        body_insts: a function that generates the body of the loop.
+                    It needs a index as parameter.
+    """
+    assert is_sequence_rprimitive(expr_reg.type)
+    target_type = builder.get_sequence_type(expr)
+
+    body_block = BasicBlock()
+    step_block = BasicBlock()
+    exit_block = BasicBlock()
+    condition_block = BasicBlock()
+
+    for_gen = ForSequence(builder, index, body_block, exit_block, line, False)
+    for_gen.init(expr_reg, target_type, reverse=False)
+
+    builder.push_loop_stack(step_block, exit_block)
+
+    builder.goto_and_activate(condition_block)
+    for_gen.gen_condition()
+
+    builder.activate_block(body_block)
+    for_gen.begin_body()
+    body_insts(builder.read(for_gen.index_target))
+
+    builder.goto_and_activate(step_block)
+    for_gen.gen_step()
+    builder.goto(condition_block)
+
+    for_gen.add_cleanup(exit_block)
+    builder.pop_loop_stack()
+
+    builder.activate_block(exit_block)
+
+
+def sequence_from_generator_preallocate_helper(
+    builder: IRBuilder,
+    gen: GeneratorExpr,
+    empty_op_llbuilder: Callable[[Value, int], Value],
+    set_item_op: CFunctionDescription,
+) -> Value | None:
+    """Generate a new tuple or list from a simple generator expression.
+
+    Currently we only optimize for simplest generator expression, which means that
+    there is no condition list in the generator and only one original sequence with
+    one index is allowed.
+
+    e.g.  (1) tuple(f(x) for x in a_list/a_tuple)
+          (2) list(f(x) for x in a_list/a_tuple)
+          (3) [f(x) for x in a_list/a_tuple]
+    RTuple as an original sequence is not supported yet.
+
+    Args:
+        empty_op_llbuilder: A function that can generate an empty sequence op when
+            passed in length. See `new_list_op_with_length` and `new_tuple_op_with_length`
+            for detailed implementation.
+        set_item_op: A primitive that can modify an arbitrary position of a sequence.
+            The op should have three arguments:
+                - Self
+                - Target position
+                - New Value
+            See `new_list_set_item_op` and `new_tuple_set_item_op` for detailed
+            implementation.
+    """
+    if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0:
+        rtype = builder.node_type(gen.sequences[0])
+        if is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype) or is_str_rprimitive(rtype):
+            sequence = builder.accept(gen.sequences[0])
+            length = builder.builder.builtin_len(sequence, gen.line, use_pyssize_t=True)
+            target_op = empty_op_llbuilder(length, gen.line)
+
+            def set_item(item_index: Value) -> None:
+                e = builder.accept(gen.left_expr)
+                builder.call_c(set_item_op, [target_op, item_index, e], gen.line)
+
+            for_loop_helper_with_index(
+                builder, gen.indices[0], gen.sequences[0], sequence, set_item, gen.line
+            )
+
+            return target_op
+    return None
+
+
+def translate_list_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value:
+    # Try simplest list comprehension, otherwise fall back to general one
+    val = sequence_from_generator_preallocate_helper(
+        builder,
+        gen,
+        empty_op_llbuilder=builder.builder.new_list_op_with_length,
+        set_item_op=new_list_set_item_op,
+    )
+    if val is not None:
+        return val
+
+    list_ops = builder.maybe_spill(builder.new_list_op([], gen.line))
+
+    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
+
+    def gen_inner_stmts() -> None:
+        e = builder.accept(gen.left_expr)
+        builder.call_c(list_append_op, [builder.read(list_ops), e], gen.line)
+
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+    return builder.read(list_ops)
+
+
+def translate_set_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value:
+    set_ops = builder.maybe_spill(builder.new_set_op([], gen.line))
+    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
+
+    def gen_inner_stmts() -> None:
+        e = builder.accept(gen.left_expr)
+        builder.call_c(set_add_op, [builder.read(set_ops), e], gen.line)
+
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+    return builder.read(set_ops)
+
+
+def comprehension_helper(
+    builder: IRBuilder,
+    loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]],
+    gen_inner_stmts: Callable[[], None],
+    line: int,
+) -> None:
+    """Helper function for list comprehensions.
+
+    Args:
+        loop_params: a list of (index, expr, [conditions]) tuples defining nested loops:
+            - "index" is the Lvalue indexing that loop;
+            - "expr" is the expression for the object to be iterated over;
+            - "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
+    """
+
+    def handle_loop(loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]]) -> None:
+        """Generate IR for a loop.
+
+        Given a list of (index, expression, [conditions]) tuples, generate IR
+        for the nested loops the list defines.
+        """
+        index, expr, conds, is_async = loop_params[0]
+        for_loop_helper(
+            builder,
+            index,
+            expr,
+            lambda: loop_contents(conds, loop_params[1:]),
+            None,
+            is_async=is_async,
+            line=line,
+        )
+
+    def loop_contents(
+        conds: list[Expression],
+        remaining_loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]],
+    ) -> None:
+        """Generate the body of the loop.
+
+        Args:
+            conds: a list of conditions to be evaluated (in order, with short circuiting)
+                to gate the body of the loop
+            remaining_loop_params: the parameters for any further nested loops; if it's empty
+                we'll instead evaluate the "gen_inner_stmts" function
+        """
+        # Check conditions, in order, short circuiting them.
+        for cond in conds:
+            cond_val = builder.accept(cond)
+            cont_block, rest_block = BasicBlock(), BasicBlock()
+            # If the condition is true we'll skip the continue.
+            builder.add_bool_branch(cond_val, rest_block, cont_block)
+            builder.activate_block(cont_block)
+            builder.nonlocal_control[-1].gen_continue(builder, cond.line)
+            builder.goto_and_activate(rest_block)
+
+        if remaining_loop_params:
+            # There's another nested level, so the body of this loop is another loop.
+            return handle_loop(remaining_loop_params)
+        else:
+            # We finally reached the actual body of the generator.
+            # Generate the IR for the inner loop body.
+            gen_inner_stmts()
+
+    handle_loop(loop_params)
+
+
+def is_range_ref(expr: RefExpr) -> bool:
+    return (
+        expr.fullname == "builtins.range"
+        or isinstance(expr.node, TypeAlias)
+        and expr.fullname == "six.moves.xrange"
+    )
+
+
+def make_for_loop_generator(
+    builder: IRBuilder,
+    index: Lvalue,
+    expr: Expression,
+    body_block: BasicBlock,
+    loop_exit: BasicBlock,
+    line: int,
+    is_async: bool = False,
+    nested: bool = False,
+) -> ForGenerator:
+    """Return helper object for generating a for loop over an iterable.
+
+    If "nested" is True, this is a nested iterator such as "e" in "enumerate(e)".
+    """
+
+    # Do an async loop if needed. async is always generic
+    if is_async:
+        expr_reg = builder.accept(expr)
+        async_obj = ForAsyncIterable(builder, index, body_block, loop_exit, line, nested)
+        item_type = builder._analyze_iterable_item_type(expr)
+        item_rtype = builder.type_to_rtype(item_type)
+        async_obj.init(expr_reg, item_rtype)
+        return async_obj
+
+    rtyp = builder.node_type(expr)
+    if is_sequence_rprimitive(rtyp):
+        # Special case "for x in <list>".
+        expr_reg = builder.accept(expr)
+        target_type = builder.get_sequence_type(expr)
+
+        for_list = ForSequence(builder, index, body_block, loop_exit, line, nested)
+        for_list.init(expr_reg, target_type, reverse=False)
+        return for_list
+
+    if is_dict_rprimitive(rtyp):
+        # Special case "for k in <dict>".
+        expr_reg = builder.accept(expr)
+        target_type = builder.get_dict_key_type(expr)
+
+        for_dict = ForDictionaryKeys(builder, index, body_block, loop_exit, line, nested)
+        for_dict.init(expr_reg, target_type)
+        return for_dict
+
+    if isinstance(expr, CallExpr) and isinstance(expr.callee, RefExpr):
+        if (
+            is_range_ref(expr.callee)
+            and (
+                len(expr.args) <= 2
+                or (len(expr.args) == 3 and builder.extract_int(expr.args[2]) is not None)
+            )
+            and set(expr.arg_kinds) == {ARG_POS}
+        ):
+            # Special case "for x in range(...)".
+            # We support the 3 arg form but only for int literals, since it doesn't
+            # seem worth the hassle of supporting dynamically determining which
+            # direction of comparison to do.
+            if len(expr.args) == 1:
+                start_reg: Value = Integer(0)
+                end_reg = builder.accept(expr.args[0])
+            else:
+                start_reg = builder.accept(expr.args[0])
+                end_reg = builder.accept(expr.args[1])
+            if len(expr.args) == 3:
+                step = builder.extract_int(expr.args[2])
+                assert step is not None
+                if step == 0:
+                    builder.error("range() step can't be zero", expr.args[2].line)
+            else:
+                step = 1
+
+            for_range = ForRange(builder, index, body_block, loop_exit, line, nested)
+            for_range.init(start_reg, end_reg, step)
+            return for_range
+
+        elif (
+            expr.callee.fullname == "builtins.enumerate"
+            and len(expr.args) == 1
+            and expr.arg_kinds == [ARG_POS]
+            and isinstance(index, TupleExpr)
+            and len(index.items) == 2
+        ):
+            # Special case "for i, x in enumerate(y)".
+            lvalue1 = index.items[0]
+            lvalue2 = index.items[1]
+            for_enumerate = ForEnumerate(builder, index, body_block, loop_exit, line, nested)
+            for_enumerate.init(lvalue1, lvalue2, expr.args[0])
+            return for_enumerate
+
+        elif (
+            expr.callee.fullname == "builtins.zip"
+            and len(expr.args) >= 2
+            and set(expr.arg_kinds) == {ARG_POS}
+            and isinstance(index, TupleExpr)
+            and len(index.items) == len(expr.args)
+        ):
+            # Special case "for x, y in zip(a, b)".
+            for_zip = ForZip(builder, index, body_block, loop_exit, line, nested)
+            for_zip.init(index.items, expr.args)
+            return for_zip
+
+        if (
+            expr.callee.fullname == "builtins.reversed"
+            and len(expr.args) == 1
+            and expr.arg_kinds == [ARG_POS]
+            and is_sequence_rprimitive(builder.node_type(expr.args[0]))
+        ):
+            # Special case "for x in reversed(<list>)".
+            expr_reg = builder.accept(expr.args[0])
+            target_type = builder.get_sequence_type(expr)
+
+            for_list = ForSequence(builder, index, body_block, loop_exit, line, nested)
+            for_list.init(expr_reg, target_type, reverse=True)
+            return for_list
+    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)
+        if is_dict_rprimitive(rtype) and expr.callee.name in ("keys", "values", "items"):
+            expr_reg = builder.accept(expr.callee.expr)
+            for_dict_type: type[ForGenerator] | None = None
+            if expr.callee.name == "keys":
+                target_type = builder.get_dict_key_type(expr.callee.expr)
+                for_dict_type = ForDictionaryKeys
+            elif expr.callee.name == "values":
+                target_type = builder.get_dict_value_type(expr.callee.expr)
+                for_dict_type = ForDictionaryValues
+            else:
+                target_type = builder.get_dict_item_type(expr.callee.expr)
+                for_dict_type = ForDictionaryItems
+            for_dict_gen = for_dict_type(builder, index, body_block, loop_exit, line, nested)
+            for_dict_gen.init(expr_reg, target_type)
+            return for_dict_gen
+
+    iterable_expr_reg: Value | None = None
+    if isinstance(expr, SetExpr):
+        # Special case "for x in <set literal>".
+        from mypyc.irbuild.expression import precompute_set_literal
+
+        set_literal = precompute_set_literal(builder, expr)
+        if set_literal is not None:
+            iterable_expr_reg = set_literal
+
+    # Default to a generic for loop.
+    if iterable_expr_reg is None:
+        iterable_expr_reg = builder.accept(expr)
+    for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested)
+    item_type = builder._analyze_iterable_item_type(expr)
+    item_rtype = builder.type_to_rtype(item_type)
+    for_obj.init(iterable_expr_reg, item_rtype)
+    return for_obj
+
+
+class ForGenerator:
+    """Abstract base class for generating for loops."""
+
+    def __init__(
+        self,
+        builder: IRBuilder,
+        index: Lvalue,
+        body_block: BasicBlock,
+        loop_exit: BasicBlock,
+        line: int,
+        nested: bool,
+    ) -> None:
+        self.builder = builder
+        self.index = index
+        self.body_block = body_block
+        self.line = line
+        # Some for loops need a cleanup block that we execute at exit. We
+        # create a cleanup block if needed. However, if we are generating a for
+        # loop for a nested iterator, such as "e" in "enumerate(e)", the
+        # outermost generator should generate the cleanup block -- we don't
+        # need to do it here.
+        if self.need_cleanup() and not nested:
+            # Create a new block to handle cleanup after loop exit.
+            self.loop_exit = BasicBlock()
+        else:
+            # Just use the existing loop exit block.
+            self.loop_exit = loop_exit
+
+    def need_cleanup(self) -> bool:
+        """If this returns true, we need post-loop cleanup."""
+        return False
+
+    def add_cleanup(self, exit_block: BasicBlock) -> None:
+        """Add post-loop cleanup, if needed."""
+        if self.need_cleanup():
+            self.builder.activate_block(self.loop_exit)
+            self.gen_cleanup()
+            self.builder.goto(exit_block)
+
+    def gen_condition(self) -> None:
+        """Generate check for loop exit (e.g. exhaustion of iteration)."""
+
+    def begin_body(self) -> None:
+        """Generate ops at the beginning of the body (if needed)."""
+
+    def gen_step(self) -> None:
+        """Generate stepping to the next item (if needed)."""
+
+    def gen_cleanup(self) -> None:
+        """Generate post-loop cleanup (if needed)."""
+
+    def load_len(self, expr: Value | AssignmentTarget) -> Value:
+        """A helper to get collection length, used by several subclasses."""
+        return self.builder.builder.builtin_len(self.builder.read(expr, self.line), self.line)
+
+
+class ForIterable(ForGenerator):
+    """Generate IR for a for loop over an arbitrary iterable (the general case)."""
+
+    def need_cleanup(self) -> bool:
+        # Create a new cleanup block for when the loop is finished.
+        return True
+
+    def init(self, expr_reg: Value, target_type: RType) -> None:
+        # Define targets to contain the expression, along with the iterator that will be used
+        # for the for-loop. If we are inside of a generator function, spill these into the
+        # environment class.
+        builder = self.builder
+        iter_reg = builder.call_c(iter_op, [expr_reg], self.line)
+        builder.maybe_spill(expr_reg)
+        self.iter_target = builder.maybe_spill(iter_reg)
+        self.target_type = target_type
+
+    def gen_condition(self) -> None:
+        # We call __next__ on the iterator and check to see if the return value
+        # is NULL, which signals either the end of the Iterable being traversed
+        # or an exception being raised. Note that Branch.IS_ERROR checks only
+        # for NULL (an exception does not necessarily have to be raised).
+        builder = self.builder
+        line = self.line
+        self.next_reg = builder.call_c(next_op, [builder.read(self.iter_target, line)], line)
+        builder.add(Branch(self.next_reg, self.loop_exit, self.body_block, Branch.IS_ERROR))
+
+    def begin_body(self) -> None:
+        # Assign the value obtained from __next__ to the
+        # lvalue so that it can be referenced by code in the body of the loop.
+        builder = self.builder
+        line = self.line
+        # We unbox here so that iterating with tuple unpacking generates a tuple based
+        # unpack instead of an iterator based one.
+        next_reg = builder.coerce(self.next_reg, self.target_type, line)
+        builder.assign(builder.get_assignment_target(self.index), next_reg, line)
+
+    def gen_step(self) -> None:
+        # Nothing to do here, since we get the next item as part of gen_condition().
+        pass
+
+    def gen_cleanup(self) -> None:
+        # We set the branch to go here if the conditional evaluates to true. If
+        # an exception was raised during the loop, then err_reg will be set to
+        # True. If no_err_occurred_op returns False, then the exception will be
+        # propagated using the ERR_FALSE flag.
+        self.builder.call_c(no_err_occurred_op, [], self.line)
+
+
+class ForAsyncIterable(ForGenerator):
+    """Generate IR for an async for loop."""
+
+    def init(self, expr_reg: Value, target_type: RType) -> None:
+        # Define targets to contain the expression, along with the
+        # iterator that will be used for the for-loop. We are inside
+        # of a generator function, so we will spill these into
+        # environment class.
+        builder = self.builder
+        iter_reg = builder.call_c(aiter_op, [expr_reg], self.line)
+        builder.maybe_spill(expr_reg)
+        self.iter_target = builder.maybe_spill(iter_reg)
+        self.target_type = target_type
+        self.stop_reg = Register(bool_rprimitive)
+
+    def gen_condition(self) -> None:
+        # This does the test and fetches the next value
+        # try:
+        #     TARGET = await type(iter).__anext__(iter)
+        #     stop = False
+        # except StopAsyncIteration:
+        #     stop = True
+        #
+        # What a pain.
+        # There are optimizations available here if we punch through some abstractions.
+
+        from mypyc.irbuild.statement import emit_await, transform_try_except
+
+        builder = self.builder
+        line = self.line
+
+        def except_match() -> Value:
+            addr = builder.add(LoadAddress(pointer_rprimitive, stop_async_iteration_op.src, line))
+            return builder.add(LoadMem(stop_async_iteration_op.type, addr))
+
+        def try_body() -> None:
+            awaitable = builder.call_c(anext_op, [builder.read(self.iter_target)], line)
+            self.next_reg = emit_await(builder, awaitable, line)
+            builder.assign(self.stop_reg, builder.false(), -1)
+
+        def except_body() -> None:
+            builder.assign(self.stop_reg, builder.true(), line)
+
+        transform_try_except(
+            builder, try_body, [((except_match, line), None, except_body)], None, line
+        )
+
+        builder.add(Branch(self.stop_reg, self.loop_exit, self.body_block, Branch.BOOL))
+
+    def begin_body(self) -> None:
+        # Assign the value obtained from await __anext__ to the
+        # lvalue so that it can be referenced by code in the body of the loop.
+        builder = self.builder
+        line = self.line
+        # We unbox here so that iterating with tuple unpacking generates a tuple based
+        # unpack instead of an iterator based one.
+        next_reg = builder.coerce(self.next_reg, self.target_type, line)
+        builder.assign(builder.get_assignment_target(self.index), next_reg, line)
+
+    def gen_step(self) -> None:
+        # Nothing to do here, since we get the next item as part of gen_condition().
+        pass
+
+
+def unsafe_index(builder: IRBuilder, target: Value, index: Value, line: int) -> Value:
+    """Emit a potentially unsafe index into a target."""
+    # This doesn't really fit nicely into any of our data-driven frameworks
+    # since we want to use __getitem__ if we don't have an unsafe version,
+    # so we just check manually.
+    if is_list_rprimitive(target.type):
+        return builder.call_c(list_get_item_unsafe_op, [target, index], line)
+    else:
+        return builder.gen_method_call(target, "__getitem__", [index], None, line)
+
+
+class ForSequence(ForGenerator):
+    """Generate optimized IR for a for loop over a sequence.
+
+    Supports iterating in both forward and reverse.
+    """
+
+    def init(self, expr_reg: Value, target_type: RType, reverse: bool) -> None:
+        builder = self.builder
+        self.reverse = reverse
+        # Define target to contain the expression, along with the index that will be used
+        # for the for-loop. If we are inside of a generator function, spill these into the
+        # environment class.
+        self.expr_target = builder.maybe_spill(expr_reg)
+        if not reverse:
+            index_reg: Value = Integer(0)
+        else:
+            index_reg = builder.binary_op(
+                self.load_len(self.expr_target), Integer(1), "-", self.line
+            )
+        self.index_target = builder.maybe_spill_assignable(index_reg)
+        self.target_type = target_type
+
+    def gen_condition(self) -> None:
+        builder = self.builder
+        line = self.line
+        # TODO: Don't reload the length each time when iterating an immutable sequence?
+        if self.reverse:
+            # If we are iterating in reverse order, we obviously need
+            # to check that the index is still positive. Somewhat less
+            # obviously we still need to check against the length,
+            # since it could shrink out from under us.
+            comparison = builder.binary_op(
+                builder.read(self.index_target, line), Integer(0), ">=", line
+            )
+            second_check = BasicBlock()
+            builder.add_bool_branch(comparison, second_check, self.loop_exit)
+            builder.activate_block(second_check)
+        # For compatibility with python semantics we recalculate the length
+        # at every iteration.
+        len_reg = self.load_len(self.expr_target)
+        comparison = builder.binary_op(builder.read(self.index_target, line), len_reg, "<", line)
+        builder.add_bool_branch(comparison, self.body_block, self.loop_exit)
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+        # Read the next list item.
+        value_box = unsafe_index(
+            builder,
+            builder.read(self.expr_target, line),
+            builder.read(self.index_target, line),
+            line,
+        )
+        assert value_box
+        # We coerce to the type of list elements here so that
+        # iterating with tuple unpacking generates a tuple based
+        # unpack instead of an iterator based one.
+        builder.assign(
+            builder.get_assignment_target(self.index),
+            builder.coerce(value_box, self.target_type, line),
+            line,
+        )
+
+    def gen_step(self) -> None:
+        # Step to the next item.
+        builder = self.builder
+        line = self.line
+        step = 1 if not self.reverse else -1
+        add = builder.int_op(
+            short_int_rprimitive,
+            builder.read(self.index_target, line),
+            Integer(step),
+            IntOp.ADD,
+            line,
+        )
+        builder.assign(self.index_target, add, line)
+
+
+class ForDictionaryCommon(ForGenerator):
+    """Generate optimized IR for a for loop over dictionary keys/values.
+
+    The logic is pretty straightforward, we use PyDict_Next() API wrapped in
+    a tuple, so that we can modify only a single register. The layout of the tuple:
+      * f0: are there more items (bool)
+      * f1: current offset (int)
+      * f2: next key (object)
+      * f3: next value (object)
+    For more info see https://docs.python.org/3/c-api/dict.html#c.PyDict_Next.
+
+    Note that for subclasses we fall back to generic PyObject_GetIter() logic,
+    since they may override some iteration methods in subtly incompatible manner.
+    The fallback logic is implemented in CPy.h via dynamic type check.
+    """
+
+    dict_next_op: ClassVar[CFunctionDescription]
+    dict_iter_op: ClassVar[CFunctionDescription]
+
+    def need_cleanup(self) -> bool:
+        # Technically, a dict subclass can raise an unrelated exception
+        # in __next__(), so we need this.
+        return True
+
+    def init(self, expr_reg: Value, target_type: RType) -> None:
+        builder = self.builder
+        self.target_type = target_type
+
+        # We add some variables to environment class, so they can be read across yield.
+        self.expr_target = builder.maybe_spill(expr_reg)
+        offset = Integer(0)
+        self.offset_target = builder.maybe_spill_assignable(offset)
+        self.size = builder.maybe_spill(self.load_len(self.expr_target))
+
+        # For dict class (not a subclass) this is the dictionary itself.
+        iter_reg = builder.call_c(self.dict_iter_op, [expr_reg], self.line)
+        self.iter_target = builder.maybe_spill(iter_reg)
+
+    def gen_condition(self) -> None:
+        """Get next key/value pair, set new offset, and check if we should continue."""
+        builder = self.builder
+        line = self.line
+        self.next_tuple = self.builder.call_c(
+            self.dict_next_op,
+            [builder.read(self.iter_target, line), builder.read(self.offset_target, line)],
+            line,
+        )
+
+        # Do this here instead of in gen_step() to minimize variables in environment.
+        new_offset = builder.add(TupleGet(self.next_tuple, 1, line))
+        builder.assign(self.offset_target, new_offset, line)
+
+        should_continue = builder.add(TupleGet(self.next_tuple, 0, line))
+        builder.add(Branch(should_continue, self.body_block, self.loop_exit, Branch.BOOL))
+
+    def gen_step(self) -> None:
+        """Check that dictionary didn't change size during iteration.
+
+        Raise RuntimeError if it is not the case to match CPython behavior.
+        """
+        builder = self.builder
+        line = self.line
+        # Technically, we don't need a new primitive for this, but it is simpler.
+        builder.call_c(
+            dict_check_size_op,
+            [builder.read(self.expr_target, line), builder.read(self.size, line)],
+            line,
+        )
+
+    def gen_cleanup(self) -> None:
+        # Same as for generic ForIterable.
+        self.builder.call_c(no_err_occurred_op, [], self.line)
+
+
+class ForDictionaryKeys(ForDictionaryCommon):
+    """Generate optimized IR for a for loop over dictionary keys."""
+
+    dict_next_op = dict_next_key_op
+    dict_iter_op = dict_key_iter_op
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+
+        # Key is stored at the third place in the tuple.
+        key = builder.add(TupleGet(self.next_tuple, 2, line))
+        builder.assign(
+            builder.get_assignment_target(self.index),
+            builder.coerce(key, self.target_type, line),
+            line,
+        )
+
+
+class ForDictionaryValues(ForDictionaryCommon):
+    """Generate optimized IR for a for loop over dictionary values."""
+
+    dict_next_op = dict_next_value_op
+    dict_iter_op = dict_value_iter_op
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+
+        # Value is stored at the third place in the tuple.
+        value = builder.add(TupleGet(self.next_tuple, 2, line))
+        builder.assign(
+            builder.get_assignment_target(self.index),
+            builder.coerce(value, self.target_type, line),
+            line,
+        )
+
+
+class ForDictionaryItems(ForDictionaryCommon):
+    """Generate optimized IR for a for loop over dictionary items."""
+
+    dict_next_op = dict_next_item_op
+    dict_iter_op = dict_item_iter_op
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+
+        key = builder.add(TupleGet(self.next_tuple, 2, line))
+        value = builder.add(TupleGet(self.next_tuple, 3, line))
+
+        # Coerce just in case e.g. key is itself a tuple to be unpacked.
+        assert isinstance(self.target_type, RTuple)
+        key = builder.coerce(key, self.target_type.types[0], line)
+        value = builder.coerce(value, self.target_type.types[1], line)
+
+        target = builder.get_assignment_target(self.index)
+        if isinstance(target, AssignmentTargetTuple):
+            # Simpler code for common case: for k, v in d.items().
+            if len(target.items) != 2:
+                builder.error("Expected a pair for dict item iteration", line)
+            builder.assign(target.items[0], key, line)
+            builder.assign(target.items[1], value, line)
+        else:
+            rvalue = builder.add(TupleSet([key, value], line))
+            builder.assign(target, rvalue, line)
+
+
+class ForRange(ForGenerator):
+    """Generate optimized IR for a for loop over an integer range."""
+
+    def init(self, start_reg: Value, end_reg: Value, step: int) -> None:
+        builder = self.builder
+        self.start_reg = start_reg
+        self.end_reg = end_reg
+        self.step = step
+        self.end_target = builder.maybe_spill(end_reg)
+        if is_short_int_rprimitive(start_reg.type) and is_short_int_rprimitive(end_reg.type):
+            index_type: RType = short_int_rprimitive
+        elif is_fixed_width_rtype(end_reg.type):
+            index_type = end_reg.type
+        else:
+            index_type = int_rprimitive
+        index_reg = Register(index_type)
+        builder.assign(index_reg, start_reg, -1)
+        self.index_reg = builder.maybe_spill_assignable(index_reg)
+        # Initialize loop index to 0. Assert that the index target is assignable.
+        self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index)
+        builder.assign(self.index_target, builder.read(self.index_reg, self.line), self.line)
+
+    def gen_condition(self) -> None:
+        builder = self.builder
+        line = self.line
+        # Add loop condition check.
+        cmp = "<" if self.step > 0 else ">"
+        comparison = builder.binary_op(
+            builder.read(self.index_reg, line), builder.read(self.end_target, line), cmp, line
+        )
+        builder.add_bool_branch(comparison, self.body_block, self.loop_exit)
+
+    def gen_step(self) -> None:
+        builder = self.builder
+        line = self.line
+
+        # Increment index register. If the range is known to fit in short ints, use
+        # short ints.
+        if is_short_int_rprimitive(self.start_reg.type) and is_short_int_rprimitive(
+            self.end_reg.type
+        ):
+            new_val = builder.int_op(
+                short_int_rprimitive,
+                builder.read(self.index_reg, line),
+                Integer(self.step),
+                IntOp.ADD,
+                line,
+            )
+
+        else:
+            new_val = builder.binary_op(
+                builder.read(self.index_reg, line), Integer(self.step), "+", line
+            )
+        builder.assign(self.index_reg, new_val, line)
+        builder.assign(self.index_target, new_val, line)
+
+
+class ForInfiniteCounter(ForGenerator):
+    """Generate optimized IR for a for loop counting from 0 to infinity."""
+
+    def init(self) -> None:
+        builder = self.builder
+        # Create a register to store the state of the loop index and
+        # initialize this register along with the loop index to 0.
+        zero = Integer(0)
+        self.index_reg = builder.maybe_spill_assignable(zero)
+        self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index)
+        builder.assign(self.index_target, zero, self.line)
+
+    def gen_step(self) -> None:
+        builder = self.builder
+        line = self.line
+        # We can safely assume that the integer is short, since we are not going to wrap
+        # around a 63-bit integer.
+        # NOTE: This would be questionable if short ints could be 32 bits.
+        new_val = builder.int_op(
+            short_int_rprimitive, builder.read(self.index_reg, line), Integer(1), IntOp.ADD, line
+        )
+        builder.assign(self.index_reg, new_val, line)
+        builder.assign(self.index_target, new_val, line)
+
+
+class ForEnumerate(ForGenerator):
+    """Generate optimized IR for a for loop of form "for i, x in enumerate(it)"."""
+
+    def need_cleanup(self) -> bool:
+        # The wrapped for loop might need cleanup. This might generate a
+        # redundant cleanup block, but that's okay.
+        return True
+
+    def init(self, index1: Lvalue, index2: Lvalue, expr: Expression) -> None:
+        # Count from 0 to infinity (for the index lvalue).
+        self.index_gen = ForInfiniteCounter(
+            self.builder, index1, self.body_block, self.loop_exit, self.line, nested=True
+        )
+        self.index_gen.init()
+        # Iterate over the actual iterable.
+        self.main_gen = make_for_loop_generator(
+            self.builder, index2, expr, self.body_block, self.loop_exit, self.line, nested=True
+        )
+
+    def gen_condition(self) -> None:
+        # No need for a check for the index generator, since it's unconditional.
+        self.main_gen.gen_condition()
+
+    def begin_body(self) -> None:
+        self.index_gen.begin_body()
+        self.main_gen.begin_body()
+
+    def gen_step(self) -> None:
+        self.index_gen.gen_step()
+        self.main_gen.gen_step()
+
+    def gen_cleanup(self) -> None:
+        self.index_gen.gen_cleanup()
+        self.main_gen.gen_cleanup()
+
+
+class ForZip(ForGenerator):
+    """Generate IR for a for loop of form `for x, ... in zip(a, ...)`."""
+
+    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, indexes: list[Lvalue], exprs: list[Expression]) -> None:
+        assert len(indexes) == len(exprs)
+        # Condition check will require multiple basic blocks, since there will be
+        # multiple conditions to check.
+        self.cond_blocks = [BasicBlock() for _ in range(len(indexes) - 1)] + [self.body_block]
+        self.gens: list[ForGenerator] = []
+        for index, expr, next_block in zip(indexes, exprs, self.cond_blocks):
+            gen = make_for_loop_generator(
+                self.builder, index, expr, next_block, self.loop_exit, self.line, nested=True
+            )
+            self.gens.append(gen)
+
+    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:
+        for gen in self.gens:
+            gen.begin_body()
+
+    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()