Skip to content

loops #106

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
varun-r-mallya merged 8 commits into from
Sep 25, 2026
Merged

loops #106

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
6220180
Core: Lower annotated assignment (x: c_int32 = 0)
varun-r-mallya Sep 24, 2026
cd3cb1b
Tests: Annotated assignment passes now
varun-r-mallya Sep 24, 2026
5ab6281
Core: Lower if-branches in the function's return type
varun-r-mallya Sep 24, 2026
f1555a6
Core: Lower for-over-range and while loops, with break and continue
varun-r-mallya Sep 24, 2026
bae742b
Tests: Loops pass now, plus the cases the lowering promises
varun-r-mallya Sep 24, 2026
3a6277e
Tests: Indent the loop C reference with spaces
varun-r-mallya Sep 24, 2026
1225e94
Core: Fix loop edge cases the lowering got wrong
varun-r-mallya Sep 25, 2026
e9e067e
Tests: Run loop programs on the host, plus adversarial loop cases
varun-r-mallya Sep 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
269 changes: 222 additions & 47 deletions pythonbpf/allocation_pass.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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):
Expand Down
5 changes: 5 additions & 0 deletions pythonbpf/context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Loading
Loading