Core: Lower for-over-range and while loops, with break and continue

`for i in range(...)` steps a hidden induction counter, allocated in the
entry block next to the loop variable, and copies it into the variable each
iteration, so rebinding the variable in the body leaves the trip count alone
as in Python (and the loop stays visibly bounded for the verifier). The
counter is 64-bit, signed unless the bounds make C's arithmetic unsigned.
Bounds are evaluated once, before the loop; the step must be a nonzero
integer literal, because its sign picks the loop test.

`while` re-evaluates its test at the top of each iteration. `break` and
`continue` branch through a loop-target stack on the compilation context;
outside a loop they are a SyntaxError, as in Python. A loop's else-branch
runs on normal exit only. Anything but range() as the iterable is rejected
with NotImplementedError.

The allocation pass now descends into loop bodies and counts helper temps in
loop headers. Checked against clang -O2 on tests/c-form/loops.bpf.c: the
constant-bound cases fold to the same ret, and a helper-bounded loop gets the
same rotated loop.

Author: varun-r-mallya

pythonbpf/allocation_pass.py

@@ -8,6 +8,7 @@
 from .expr import VmlinuxHandlerRegistry
 from pythonbpf.type_deducer import ctypes_to_ir, is_ctypes, IntTy, signedness
 from pythonbpf.expr.type_inference import infer_int_type
+from pythonbpf.expr.operators import usual_arithmetic_conversions
 from pythonbpf.maps import BPFMapType
 
 logger = logging.getLogger(__name__)
@@ -89,6 +90,110 @@ def handle_ann_assign_allocation(compilation_context, builder, stmt, local_sym_t
     )
 
 
+def handle_for_allocation(compilation_context, builder, stmt, local_sym_tab):
+    """Handle memory allocation for `for <name> in range(...)`.
+
+    Two slots, both in the entry block so a loop never grows the stack: the
+    loop variable, and a hidden induction counter the loop actually steps.
+    Keeping them apart is what makes rebinding the loop variable in the body
+    leave the trip count alone, as in Python -- and a counter the body cannot
+    touch is what keeps the loop visibly bounded for the verifier.
+    """
+    start, stop, step = parse_range(stmt)
+
+    # range() yields Python ints; like any undeclared local they are 64-bit,
+    # signed unless the bounds make C's arithmetic unsigned.
+    bound_types = [
+        infer_int_type(bound, local_sym_tab, compilation_context)
+        for bound in (start, stop)
+        if bound is not None
+    ]
+    signed = True
+    if all(ty is not None for ty in bound_types):
+        common = bound_types[0]
+        for ty in bound_types[1:]:
+            common = usual_arithmetic_conversions(common, ty)
+        signed = signedness(common)
+    if not signed and step < 0:
+        raise SyntaxError(
+            f"range() on line {stmt.lineno} counts down over unsigned bounds; "
+            f"the counter would wrap instead of stopping"
+        )
+    loop_ty = IntTy(64, signed)
+
+    counter = range_counter_name(stmt)
+    local_sym_tab[counter] = LocalSymbol(
+        _allocate_with_type(builder, counter, loop_ty), loop_ty
+    )
+    _bind_name(
+        compilation_context,
+        stmt,
+        stmt.target,
+        local_sym_tab,
+        lambda var_name: local_sym_tab.__setitem__(
+            var_name,
+            LocalSymbol(_allocate_with_type(builder, var_name, loop_ty), loop_ty),
+        ),
+    )
+
+
+def range_counter_name(stmt):
+    """Symbol-table name of a range loop's hidden induction counter, unique per
+    loop so nested loops each get their own."""
+    return f"__range_idx_{stmt.lineno}_{stmt.col_offset}"
+
+
+def parse_range(stmt):
+    """Split `for <name> in range(...)` into (start, stop, step): start is an
+    expression or None (meaning 0), stop an expression, step a nonzero int.
+
+    step has to be known at compile time, because its sign decides whether the
+    loop runs while the counter is below stop or above it.
+    """
+    it = stmt.iter
+    if not (
+        isinstance(it, ast.Call)
+        and isinstance(it.func, ast.Name)
+        and it.func.id == "range"
+    ):
+        raise NotImplementedError(
+            f"for loop on line {stmt.lineno}: only range(...) can be iterated "
+            f"so far, got {ast.unparse(it)}"
+        )
+    if not isinstance(stmt.target, ast.Name):
+        raise NotImplementedError(
+            f"for loop on line {stmt.lineno}: only a plain name can be the loop "
+            f"variable, got {ast.unparse(stmt.target)}"
+        )
+    if it.keywords or not 1 <= len(it.args) <= 3:
+        raise SyntaxError(
+            f"range() on line {stmt.lineno} takes 1 to 3 positional arguments"
+        )
+
+    if len(it.args) == 1:
+        return None, it.args[0], 1
+    start, stop = it.args[0], it.args[1]
+    if len(it.args) == 2:
+        return start, stop, 1
+
+    step_node = it.args[2]
+    negate = isinstance(step_node, ast.UnaryOp) and isinstance(step_node.op, ast.USub)
+    literal = step_node.operand if negate else step_node
+    if not (
+        isinstance(literal, ast.Constant)
+        and isinstance(literal.value, int)
+        and not isinstance(literal.value, bool)
+    ):
+        raise SyntaxError(
+            f"range() step on line {stmt.lineno} must be an integer literal, "
+            f"got {ast.unparse(step_node)}"
+        )
+    step = -literal.value if negate else literal.value
+    if step == 0:
+        raise ValueError(f"range() arg 3 must not be zero (line {stmt.lineno})")
+    return start, stop, step
+
+
 def annotation_to_ir(annotation, lineno):
     """IR type for a ctypes annotation, written `c_int32` or `ctypes.c_int32`."""
     if isinstance(annotation, ast.Name):

pythonbpf/context.py

@@ -82,7 +82,12 @@ def __init__(self, module: ir.Module):
         # Current function context (optional, if needed globally during function processing)
         self.current_func = None
 
+        # Enclosing loops of the statement being lowered, innermost last, as
+        # (continue target, break target) blocks.
+        self.loop_stack: list[tuple[ir.Block, ir.Block]] = []
+
     def reset(self):
         """Reset state between functions if necessary, though new context per compile is preferred."""
         self.scratch_pool.reset()
         self.current_func = None
+        self.loop_stack = []

pythonbpf/functions/functions_pass.py

@@ -26,6 +26,9 @@
 from pythonbpf.allocation_pass import (
     handle_assign_allocation,
     handle_ann_assign_allocation,
+    handle_for_allocation,
+    parse_range,
+    range_counter_name,
     allocate_temp_pool,
     create_targets_and_rvals,
     LocalSymbol,
@@ -82,10 +85,11 @@ def count_temps_in_call(call_node, local_sym_tab):
     return count
 
 
-def handle_if_allocation(
+def handle_block_allocation(
     compilation_context, builder, stmt, func, ret_type, local_sym_tab
 ):
-    """Recursively handle allocations in if/else branches."""
+    """Recursively handle allocations in the body and else-branch of an
+    if, for or while statement."""
     if stmt.body:
         allocate_mem(
             compilation_context,
@@ -117,15 +121,23 @@ def merge_type_counts(count_dict):
     def update_max_temps_for_stmt(stmt):
         nonlocal max_temps_needed
 
-        if isinstance(stmt, ast.If):
+        if isinstance(stmt, (ast.If, ast.For, ast.While)):
+            # A loop header is evaluated like a statement of its own: range()
+            # bounds once before the loop, a while test on every iteration.
+            header = {ast.For: "iter", ast.While: "test"}.get(type(stmt))
+            if header is not None:
+                count_temps_in_tree(getattr(stmt, header))
             for s in stmt.body:
                 update_max_temps_for_stmt(s)
             for s in stmt.orelse:
                 update_max_temps_for_stmt(s)
             return
 
+        count_temps_in_tree(stmt)
+
+    def count_temps_in_tree(tree):
         stmt_temps = {}
-        for node in ast.walk(stmt):
+        for node in ast.walk(tree):
             if isinstance(node, ast.Call):
                 call_temps = count_temps_in_call(node, local_sym_tab)
                 for typ, cnt in call_temps.items():
@@ -136,8 +148,10 @@ def update_max_temps_for_stmt(stmt):
         update_max_temps_for_stmt(stmt)
 
         # Handle allocations
-        if isinstance(stmt, ast.If):
-            handle_if_allocation(
+        if isinstance(stmt, ast.For):
+            handle_for_allocation(compilation_context, builder, stmt, local_sym_tab)
+        if isinstance(stmt, (ast.If, ast.For, ast.While)):
+            handle_block_allocation(
                 compilation_context,
                 builder,
                 stmt,
@@ -353,6 +367,158 @@ def handle_if(func, compilation_context, builder, stmt, local_sym_tab, ret_type)
     builder.position_at_end(merge_block)
 
 
+def _lower_loop(
+    func,
+    compilation_context,
+    builder,
+    stmt,
+    local_sym_tab,
+    ret_type,
+    body_block,
+    continue_block,
+    end_block,
+    else_block,
+):
+    """What for and while share once their header is emitted: the body, with
+    `continue` and `break` bound to this loop, falling through to
+    continue_block; then the else-branch, which runs only when the loop ends
+    without a break, so it sits between the exit test and end_block."""
+    builder.position_at_end(body_block)
+    compilation_context.loop_stack.append((continue_block, end_block))
+    try:
+        process_block(
+            func, compilation_context, builder, stmt.body, local_sym_tab, ret_type
+        )
+    finally:
+        compilation_context.loop_stack.pop()
+    if not builder.block.is_terminated:
+        builder.branch(continue_block)
+
+    if else_block is not None:
+        # Outside this loop's scope: a break here leaves the enclosing loop.
+        builder.position_at_end(else_block)
+        process_block(
+            func, compilation_context, builder, stmt.orelse, local_sym_tab, ret_type
+        )
+        if not builder.block.is_terminated:
+            builder.branch(end_block)
+
+    builder.position_at_end(end_block)
+
+
+def handle_while(func, compilation_context, builder, stmt, local_sym_tab, ret_type):
+    """Handle `while test: body [else: orelse]`. The test is re-evaluated at
+    the top of every iteration, and is where `continue` goes."""
+    cond_block = func.append_basic_block(name="while.cond")
+    body_block = func.append_basic_block(name="while.body")
+    else_block = func.append_basic_block(name="while.else") if stmt.orelse else None
+    end_block = func.append_basic_block(name="while.end")
+
+    builder.branch(cond_block)
+    builder.position_at_end(cond_block)
+    cond = handle_cond(func, compilation_context, builder, stmt.test, local_sym_tab)
+    builder.cbranch(cond, body_block, else_block or end_block)
+
+    _lower_loop(
+        func,
+        compilation_context,
+        builder,
+        stmt,
+        local_sym_tab,
+        ret_type,
+        body_block,
+        cond_block,
+        end_block,
+        else_block,
+    )
+
+
+def handle_for(func, compilation_context, builder, stmt, local_sym_tab, ret_type):
+    """Handle `for name in range(...): body [else: orelse]`.
+
+    The allocation pass made a hidden induction counter (typed from the
+    bounds) next to the loop variable. The bounds are evaluated once, before
+    the loop, as Python does; each iteration copies the counter into the loop
+    variable, and `continue` goes to the step, not straight back to the test.
+    """
+    start, stop, step = parse_range(stmt)
+    counter = local_sym_tab[range_counter_name(stmt)]
+    loop_ty = counter.ir_type
+
+    def bound(expr):
+        val, ty = get_typed_operand(
+            func, compilation_context, expr, builder, local_sym_tab
+        )
+        if val is None or not isinstance(ty, ir.IntType):
+            raise SyntaxError(
+                f"range() bound on line {stmt.lineno} must be an integer: "
+                f"{ast.unparse(expr)}"
+            )
+        return convert(builder, val, ty, loop_ty)
+
+    start_val = ir.Constant(loop_ty, 0) if start is None else bound(start)
+    stop_val = bound(stop)
+    builder.store(start_val, counter.var)
+
+    target = local_sym_tab[stmt.target.id]
+    if target.var is None:
+        raise SyntaxError(
+            f"cannot use '{stmt.target.id}' as a loop variable: it is the "
+            f"context parameter"
+        )
+
+    cond_block = func.append_basic_block(name="for.cond")
+    body_block = func.append_basic_block(name="for.body")
+    inc_block = func.append_basic_block(name="for.inc")
+    else_block = func.append_basic_block(name="for.else") if stmt.orelse else None
+    end_block = func.append_basic_block(name="for.end")
+
+    builder.branch(cond_block)
+    builder.position_at_end(cond_block)
+    idx = builder.load(counter.var)
+    # Counting up runs while below stop, counting down while above it.
+    predicate = "<" if step > 0 else ">"
+    compare = builder.icmp_signed if signedness(loop_ty) else builder.icmp_unsigned
+    builder.cbranch(
+        compare(predicate, idx, stop_val), body_block, else_block or end_block
+    )
+
+    # The loop variable is bound to the counter's value, per iteration.
+    builder.position_at_end(body_block)
+    builder.store(
+        convert(builder, builder.load(counter.var), loop_ty, target.ir_type),
+        target.var,
+    )
+
+    builder.position_at_end(inc_block)
+    next_idx = builder.add(builder.load(counter.var), ir.Constant(loop_ty, step))
+    builder.store(next_idx, counter.var)
+    builder.branch(cond_block)
+
+    _lower_loop(
+        func,
+        compilation_context,
+        builder,
+        stmt,
+        local_sym_tab,
+        ret_type,
+        body_block,
+        inc_block,
+        end_block,
+        else_block,
+    )
+
+
+def handle_loop_jump(compilation_context, builder, stmt):
+    """Handle `break` and `continue`: branch to the innermost loop's exit or
+    next-iteration block."""
+    keyword = "break" if isinstance(stmt, ast.Break) else "continue"
+    if not compilation_context.loop_stack:
+        raise SyntaxError(f"'{keyword}' outside loop (line {stmt.lineno})")
+    continue_block, break_block = compilation_context.loop_stack[-1]
+    builder.branch(break_block if keyword == "break" else continue_block)
+
+
 def handle_return(
     func, builder, stmt, local_sym_tab, ret_type, compilation_context=None
 ):
@@ -423,6 +589,12 @@ def process_stmt(
         logger.debug(f"global declaration of {', '.join(stmt.names)} already bound")
     elif isinstance(stmt, ast.If):
         handle_if(func, compilation_context, builder, stmt, local_sym_tab, ret_type)
+    elif isinstance(stmt, ast.While):
+        handle_while(func, compilation_context, builder, stmt, local_sym_tab, ret_type)
+    elif isinstance(stmt, ast.For):
+        handle_for(func, compilation_context, builder, stmt, local_sym_tab, ret_type)
+    elif isinstance(stmt, (ast.Break, ast.Continue)):
+        handle_loop_jump(compilation_context, builder, stmt)
     elif isinstance(stmt, ast.Return):
         did_return = handle_return(
             func, builder, stmt, local_sym_tab, ret_type, compilation_context
@@ -438,10 +610,10 @@ def process_stmt(
 
 
 def process_block(func, compilation_context, builder, stmts, local_sym_tab, ret_type):
-    """Process a nested statement list, such as an if-branch, in the
-    enclosing function's return type. Stops at the first statement that ends
-    the block (a return), because whatever follows it in the same list can
-    never run, and would otherwise be emitted after a terminator."""
+    """Process a nested statement list (an if-branch or loop body). Stops at
+    the first statement that ends the block -- break, continue or return --
+    because whatever follows it in the same list can never run, and would
+    otherwise be emitted after a terminator."""
     for s in stmts:
         if builder.block.is_terminated:
             break