loops (#106)

* Core: Lower annotated assignment (x: c_int32 = 0)

The annotation types the slot, so `x: c_int32 = 0` is an i32 where the
literal alone would give an i64; the store goes through the same
handle_variable_assignment plain assignment uses. A bare `x: T` makes the
slot and binds nothing.

The already-bound and global-shadowing bookkeeping that Assign allocation did
inline moves into _bind_name, which both statements now call.

* Tests: Annotated assignment passes now

* Core: Lower if-branches in the function's return type

handle_if processed its branches without passing ret_type down, so a
return inside an if fell back to process_stmt's i64 default: a -> c_int32
function got `ret i64 7`, which llc rejects. The branches now go through
process_block, which threads ret_type and stops at a statement that ends
the block instead of emitting after its terminator.

* 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.

* Tests: Loops pass now, plus the cases the lowering promises

- passing_tests/loops/: the seven range/while/break/continue/nested tests
  move over from failing_tests/, joined by rebinding the loop variable, a
  negative step, loop else-branches, and a helper-bounded loop over a
  @bpfglobal that survives -O2 as a real loop for the verifier.
- failing_tests/loops/: break outside a loop, and a non-literal range step.
- for_map_items stays xfail: map iteration needs its own design.
- c-form/loops.bpf.c: the clang reference the lowering was checked against.

* Tests: Indent the loop C reference with spaces

The remove-tabs pre-commit hook rejects tab indentation outside docs/ and
Makefiles, which failed the Format job.

* Core: Fix loop edge cases the lowering got wrong

- `pass` is accepted as a statement, the natural body of a loop run for its
  side effects.
- Falling off the end of a function returns 0 in the function's own return
  type. After `while True` that block is unreachable, but a c_int32 function
  still emitted `ret i64 0`, which llc rejected.
- A range() step other than 1 could wrap the counter past INT64_MAX or
  UINT64_MAX back inside the range and loop forever. The step block now ends
  the loop when the distance left to stop, exact as an unsigned number, is
  no longer than the step. That is a normal exit, so the else-branch runs.
  Steps of 1 cannot overshoot and keep the plain loop, so the helper-bounded
  loop the verifier walks is unchanged.
- A step of 2**64 or more no longer truncates to 0: it is a compile error,
  since it does not fit the counter. Steps from 2**63 up are emitted by
  their bits.
- The counter's signedness folds each bound into a signed 64-bit type rather
  than joining the bounds alone, so only a c_uint64 bound makes it unsigned:
  range(-1, n) over a c_uint32 n runs 4 times, as in C with an __s64 counter.

* Tests: Run loop programs on the host, plus adversarial loop cases

Levels 1 and 2 only prove a loop compiles; one can still run the wrong
number of times, or forever. test_loops_run.py JIT-compiles each
helper-free passing_tests/loops case on the host (framework/host_jit.py,
in a subprocess with a timeout) and checks its return value against what
CPython returns for the same body.

- passing_tests/loops/: pass bodies, `while True` exited by a return in a
  c_int32 function, step overflow past INT64_MAX, INT64_MIN and UINT64_MAX,
  a c_uint32 stop with a negative start, and loop variables that must stay
  one local (a nested same-name loop, a read after the loop).
- failing_tests/loops/range_step_too_wide.py: a 2**64 step is a compile error.
- reuse_loop_var_signedness.py is xfail in test_loops_run.py: every local
  gets one slot typed by its first binding, so a second loop's `i` inherits
  the first loop's unsigned type. `x = n_u64; x = -3` does the same, so the
  fix is general, not loop-specific.

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, byte_size
 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__)
@@ -49,57 +50,231 @@ def handle_assign_allocation(compilation_context, builder, stmt, local_sym_tab):
             )
             continue
 
-        var_name = target.id
+        _bind_name(
+            compilation_context,
+            stmt,
+            target,
+            local_sym_tab,
+            lambda var_name: _allocate_for_value(
+                builder, var_name, rval, local_sym_tab, compilation_context
+            ),
+        )
 
-        # Already bound in this scope: a parameter, an earlier assignment, or a
-        # `global` declaration (whose slot is the GlobalVariable). No slot needed.
-        if var_name in local_sym_tab:
-            logger.debug(f"'{var_name}' already bound, no allocation needed")
-            continue
 
-        # Not declared `global`, yet named like one: Python creates a local
-        # that shadows the global for the whole function body, and leaves the
-        # global untouched. Do the same.
-        shadows_global = var_name in compilation_context.bpf_globals
-        if shadows_global:
-            logger.info(
-                f"'{var_name}' is assigned without a 'global' declaration, so it "
-                f"is a local shadowing the @bpfglobal of the same name"
-            )
+def handle_ann_assign_allocation(compilation_context, builder, stmt, local_sym_tab):
+    """Handle memory allocation for annotated assignment (`x: c_int32 = 0`).
 
-        # Determine type and allocate based on rval
-        if isinstance(rval, ast.Call):
-            _allocate_for_call(
-                builder, var_name, rval, local_sym_tab, compilation_context
-            )
-        elif isinstance(rval, ast.Constant):
-            _allocate_for_constant(builder, var_name, rval, local_sym_tab)
-        elif isinstance(rval, ast.BinOp):
-            _allocate_for_binop(
-                builder, var_name, rval, local_sym_tab, compilation_context
-            )
-        elif isinstance(rval, ast.Name):
-            # Variable-to-variable assignment (b = a)
-            _allocate_for_name(
-                builder, var_name, rval, local_sym_tab, compilation_context
-            )
-        elif isinstance(rval, ast.Attribute):
-            # Struct field-to-variable assignment (a = dat.fld)
-            _allocate_for_attribute(
-                builder, var_name, rval, local_sym_tab, compilation_context
-            )
-        else:
-            logger.warning(
-                f"Unsupported assignment value type for {var_name}: {type(rval).__name__}"
-            )
+    The annotation, not the value, types the slot: `x: c_int32 = 0` is an i32
+    even though the literal alone would give an i64. With no value
+    (`x: c_int64`) the slot is still made, and the name stays unbound until
+    something assigns it, as in Python.
+    """
+    logger.info(f"Handling annotated assignment for allocation: {ast.dump(stmt)}")
 
-        if shadows_global and var_name in local_sym_tab:
-            # Where the binding ends, so that a read above it is reported the
-            # way Python reports it. end_lineno, not lineno, so a read on a
-            # continuation line of a multi-line binding counts as above it too.
-            local_sym_tab[var_name].shadows_global_from = (
-                getattr(stmt, "end_lineno", None) or target.lineno
-            )
+    if not isinstance(stmt.target, ast.Name):
+        raise SyntaxError(
+            f"annotated assignment on line {stmt.lineno} must target a plain "
+            f"name, got {type(stmt.target).__name__}"
+        )
+
+    ir_type = annotation_to_ir(stmt.annotation, stmt.lineno)
+    _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, ir_type), ir_type),
+        ),
+    )
+
+
+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. That arithmetic
+    # is the counter's against each bound, so it starts from a signed 64-bit
+    # counter, and only an unsigned 64-bit bound turns it unsigned: every
+    # narrower unsigned bound fits in the signed counter, as in C.
+    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 = IntTy(64, True)
+        for ty in bound_types:
+            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})")
+    if abs(step) >= 1 << 64:
+        raise ValueError(
+            f"range() step on line {stmt.lineno} does not fit the 64-bit "
+            f"loop counter: {ast.unparse(step_node)}"
+        )
+    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):
+        name = annotation.id
+    elif isinstance(annotation, ast.Attribute):
+        name = annotation.attr
+    else:
+        name = None
+    if name is None or not is_ctypes(name):
+        raise SyntaxError(
+            f"unsupported annotation on line {lineno}: {ast.unparse(annotation)} "
+            f"(annotate locals with a ctypes integer type such as c_int64)"
+        )
+    return ctypes_to_ir(name)
+
+
+def _bind_name(compilation_context, stmt, target, local_sym_tab, allocate):
+    """What every statement that binds a bare name shares, around the
+    statement-specific `allocate(var_name)` that makes the slot.
+
+    A name already bound needs no new slot. A name that is also a @bpfglobal
+    but was not declared `global` becomes a local shadowing it, and records
+    where its binding ends so a read above it is reported as Python would.
+    """
+    var_name = target.id
+
+    # Already bound in this scope: a parameter, an earlier assignment, or a
+    # `global` declaration (whose slot is the GlobalVariable). No slot needed.
+    if var_name in local_sym_tab:
+        logger.debug(f"'{var_name}' already bound, no allocation needed")
+        return
+
+    # Not declared `global`, yet named like one: Python creates a local
+    # that shadows the global for the whole function body, and leaves the
+    # global untouched. Do the same.
+    shadows_global = var_name in compilation_context.bpf_globals
+    if shadows_global:
+        logger.info(
+            f"'{var_name}' is assigned without a 'global' declaration, so it "
+            f"is a local shadowing the @bpfglobal of the same name"
+        )
+
+    allocate(var_name)
+
+    if shadows_global and var_name in local_sym_tab:
+        # Where the binding ends, so that a read above it is reported the
+        # way Python reports it. end_lineno, not lineno, so a read on a
+        # continuation line of a multi-line binding counts as above it too.
+        local_sym_tab[var_name].shadows_global_from = (
+            getattr(stmt, "end_lineno", None) or target.lineno
+        )
+
+
+def _allocate_for_value(builder, var_name, rval, local_sym_tab, compilation_context):
+    """Allocate a slot for `var_name = rval`, typed from the value."""
+    if isinstance(rval, ast.Call):
+        _allocate_for_call(builder, var_name, rval, local_sym_tab, compilation_context)
+    elif isinstance(rval, ast.Constant):
+        _allocate_for_constant(builder, var_name, rval, local_sym_tab)
+    elif isinstance(rval, ast.BinOp):
+        _allocate_for_binop(builder, var_name, rval, local_sym_tab, compilation_context)
+    elif isinstance(rval, ast.Name):
+        # Variable-to-variable assignment (b = a)
+        _allocate_for_name(builder, var_name, rval, local_sym_tab, compilation_context)
+    elif isinstance(rval, ast.Attribute):
+        # Struct field-to-variable assignment (a = dat.fld)
+        _allocate_for_attribute(
+            builder, var_name, rval, local_sym_tab, compilation_context
+        )
+    else:
+        logger.warning(
+            f"Unsupported assignment value type for {var_name}: {type(rval).__name__}"
+        )
 
 
 def _allocate_for_call(builder, var_name, rval, local_sym_tab, compilation_context):

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 = []