Skip to content

Commit 107f035

Browse files
Merge #106 (loops) for bpfsnake
# Conflicts: # tests/test_config.toml
2 parents cb1952d + 89b7eb6 commit 107f035

20 files changed

Lines changed: 719 additions & 89 deletions

‎pythonbpf/allocation_pass.py‎

Lines changed: 214 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from .expr import VmlinuxHandlerRegistry
99
from pythonbpf.type_deducer import ctypes_to_ir, is_ctypes, IntTy, signedness
1010
from pythonbpf.expr.type_inference import infer_int_type
11+
from pythonbpf.expr.operators import usual_arithmetic_conversions
1112
from pythonbpf.maps import BPFMapType
1213

1314
logger = logging.getLogger(__name__)
@@ -49,57 +50,223 @@ def handle_assign_allocation(compilation_context, builder, stmt, local_sym_tab):
4950
)
5051
continue
5152

52-
var_name = target.id
53+
_bind_name(
54+
compilation_context,
55+
stmt,
56+
target,
57+
local_sym_tab,
58+
lambda var_name: _allocate_for_value(
59+
builder, var_name, rval, local_sym_tab, compilation_context
60+
),
61+
)
5362

54-
# Already bound in this scope: a parameter, an earlier assignment, or a
55-
# `global` declaration (whose slot is the GlobalVariable). No slot needed.
56-
if var_name in local_sym_tab:
57-
logger.debug(f"'{var_name}' already bound, no allocation needed")
58-
continue
5963

60-
# Not declared `global`, yet named like one: Python creates a local
61-
# that shadows the global for the whole function body, and leaves the
62-
# global untouched. Do the same.
63-
shadows_global = var_name in compilation_context.bpf_globals
64-
if shadows_global:
65-
logger.info(
66-
f"'{var_name}' is assigned without a 'global' declaration, so it "
67-
f"is a local shadowing the @bpfglobal of the same name"
68-
)
64+
def handle_ann_assign_allocation(compilation_context, builder, stmt, local_sym_tab):
65+
"""Handle memory allocation for annotated assignment (`x: c_int32 = 0`).
6966
70-
# Determine type and allocate based on rval
71-
if isinstance(rval, ast.Call):
72-
_allocate_for_call(
73-
builder, var_name, rval, local_sym_tab, compilation_context
74-
)
75-
elif isinstance(rval, ast.Constant):
76-
_allocate_for_constant(builder, var_name, rval, local_sym_tab)
77-
elif isinstance(rval, ast.BinOp):
78-
_allocate_for_binop(
79-
builder, var_name, rval, local_sym_tab, compilation_context
80-
)
81-
elif isinstance(rval, ast.Name):
82-
# Variable-to-variable assignment (b = a)
83-
_allocate_for_name(
84-
builder, var_name, rval, local_sym_tab, compilation_context
85-
)
86-
elif isinstance(rval, ast.Attribute):
87-
# Struct field-to-variable assignment (a = dat.fld)
88-
_allocate_for_attribute(
89-
builder, var_name, rval, local_sym_tab, compilation_context
90-
)
91-
else:
92-
logger.warning(
93-
f"Unsupported assignment value type for {var_name}: {type(rval).__name__}"
94-
)
67+
The annotation, not the value, types the slot: `x: c_int32 = 0` is an i32
68+
even though the literal alone would give an i64. With no value
69+
(`x: c_int64`) the slot is still made, and the name stays unbound until
70+
something assigns it, as in Python.
71+
"""
72+
logger.info(f"Handling annotated assignment for allocation: {ast.dump(stmt)}")
9573

96-
if shadows_global and var_name in local_sym_tab:
97-
# Where the binding ends, so that a read above it is reported the
98-
# way Python reports it. end_lineno, not lineno, so a read on a
99-
# continuation line of a multi-line binding counts as above it too.
100-
local_sym_tab[var_name].shadows_global_from = (
101-
getattr(stmt, "end_lineno", None) or target.lineno
102-
)
74+
if not isinstance(stmt.target, ast.Name):
75+
raise SyntaxError(
76+
f"annotated assignment on line {stmt.lineno} must target a plain "
77+
f"name, got {type(stmt.target).__name__}"
78+
)
79+
80+
ir_type = annotation_to_ir(stmt.annotation, stmt.lineno)
81+
_bind_name(
82+
compilation_context,
83+
stmt,
84+
stmt.target,
85+
local_sym_tab,
86+
lambda var_name: local_sym_tab.__setitem__(
87+
var_name,
88+
LocalSymbol(_allocate_with_type(builder, var_name, ir_type), ir_type),
89+
),
90+
)
91+
92+
93+
def handle_for_allocation(compilation_context, builder, stmt, local_sym_tab):
94+
"""Handle memory allocation for `for <name> in range(...)`.
95+
96+
Two slots, both in the entry block so a loop never grows the stack: the
97+
loop variable, and a hidden induction counter the loop actually steps.
98+
Keeping them apart is what makes rebinding the loop variable in the body
99+
leave the trip count alone, as in Python -- and a counter the body cannot
100+
touch is what keeps the loop visibly bounded for the verifier.
101+
"""
102+
start, stop, step = parse_range(stmt)
103+
104+
# range() yields Python ints; like any undeclared local they are 64-bit,
105+
# signed unless the bounds make C's arithmetic unsigned.
106+
bound_types = [
107+
infer_int_type(bound, local_sym_tab, compilation_context)
108+
for bound in (start, stop)
109+
if bound is not None
110+
]
111+
signed = True
112+
if all(ty is not None for ty in bound_types):
113+
common = bound_types[0]
114+
for ty in bound_types[1:]:
115+
common = usual_arithmetic_conversions(common, ty)
116+
signed = signedness(common)
117+
if not signed and step < 0:
118+
raise SyntaxError(
119+
f"range() on line {stmt.lineno} counts down over unsigned bounds; "
120+
f"the counter would wrap instead of stopping"
121+
)
122+
loop_ty = IntTy(64, signed)
123+
124+
counter = range_counter_name(stmt)
125+
local_sym_tab[counter] = LocalSymbol(
126+
_allocate_with_type(builder, counter, loop_ty), loop_ty
127+
)
128+
_bind_name(
129+
compilation_context,
130+
stmt,
131+
stmt.target,
132+
local_sym_tab,
133+
lambda var_name: local_sym_tab.__setitem__(
134+
var_name,
135+
LocalSymbol(_allocate_with_type(builder, var_name, loop_ty), loop_ty),
136+
),
137+
)
138+
139+
140+
def range_counter_name(stmt):
141+
"""Symbol-table name of a range loop's hidden induction counter, unique per
142+
loop so nested loops each get their own."""
143+
return f"__range_idx_{stmt.lineno}_{stmt.col_offset}"
144+
145+
146+
def parse_range(stmt):
147+
"""Split `for <name> in range(...)` into (start, stop, step): start is an
148+
expression or None (meaning 0), stop an expression, step a nonzero int.
149+
150+
step has to be known at compile time, because its sign decides whether the
151+
loop runs while the counter is below stop or above it.
152+
"""
153+
it = stmt.iter
154+
if not (
155+
isinstance(it, ast.Call)
156+
and isinstance(it.func, ast.Name)
157+
and it.func.id == "range"
158+
):
159+
raise NotImplementedError(
160+
f"for loop on line {stmt.lineno}: only range(...) can be iterated "
161+
f"so far, got {ast.unparse(it)}"
162+
)
163+
if not isinstance(stmt.target, ast.Name):
164+
raise NotImplementedError(
165+
f"for loop on line {stmt.lineno}: only a plain name can be the loop "
166+
f"variable, got {ast.unparse(stmt.target)}"
167+
)
168+
if it.keywords or not 1 <= len(it.args) <= 3:
169+
raise SyntaxError(
170+
f"range() on line {stmt.lineno} takes 1 to 3 positional arguments"
171+
)
172+
173+
if len(it.args) == 1:
174+
return None, it.args[0], 1
175+
start, stop = it.args[0], it.args[1]
176+
if len(it.args) == 2:
177+
return start, stop, 1
178+
179+
step_node = it.args[2]
180+
negate = isinstance(step_node, ast.UnaryOp) and isinstance(step_node.op, ast.USub)
181+
literal = step_node.operand if negate else step_node
182+
if not (
183+
isinstance(literal, ast.Constant)
184+
and isinstance(literal.value, int)
185+
and not isinstance(literal.value, bool)
186+
):
187+
raise SyntaxError(
188+
f"range() step on line {stmt.lineno} must be an integer literal, "
189+
f"got {ast.unparse(step_node)}"
190+
)
191+
step = -literal.value if negate else literal.value
192+
if step == 0:
193+
raise ValueError(f"range() arg 3 must not be zero (line {stmt.lineno})")
194+
return start, stop, step
195+
196+
197+
def annotation_to_ir(annotation, lineno):
198+
"""IR type for a ctypes annotation, written `c_int32` or `ctypes.c_int32`."""
199+
if isinstance(annotation, ast.Name):
200+
name = annotation.id
201+
elif isinstance(annotation, ast.Attribute):
202+
name = annotation.attr
203+
else:
204+
name = None
205+
if name is None or not is_ctypes(name):
206+
raise SyntaxError(
207+
f"unsupported annotation on line {lineno}: {ast.unparse(annotation)} "
208+
f"(annotate locals with a ctypes integer type such as c_int64)"
209+
)
210+
return ctypes_to_ir(name)
211+
212+
213+
def _bind_name(compilation_context, stmt, target, local_sym_tab, allocate):
214+
"""What every statement that binds a bare name shares, around the
215+
statement-specific `allocate(var_name)` that makes the slot.
216+
217+
A name already bound needs no new slot. A name that is also a @bpfglobal
218+
but was not declared `global` becomes a local shadowing it, and records
219+
where its binding ends so a read above it is reported as Python would.
220+
"""
221+
var_name = target.id
222+
223+
# Already bound in this scope: a parameter, an earlier assignment, or a
224+
# `global` declaration (whose slot is the GlobalVariable). No slot needed.
225+
if var_name in local_sym_tab:
226+
logger.debug(f"'{var_name}' already bound, no allocation needed")
227+
return
228+
229+
# Not declared `global`, yet named like one: Python creates a local
230+
# that shadows the global for the whole function body, and leaves the
231+
# global untouched. Do the same.
232+
shadows_global = var_name in compilation_context.bpf_globals
233+
if shadows_global:
234+
logger.info(
235+
f"'{var_name}' is assigned without a 'global' declaration, so it "
236+
f"is a local shadowing the @bpfglobal of the same name"
237+
)
238+
239+
allocate(var_name)
240+
241+
if shadows_global and var_name in local_sym_tab:
242+
# Where the binding ends, so that a read above it is reported the
243+
# way Python reports it. end_lineno, not lineno, so a read on a
244+
# continuation line of a multi-line binding counts as above it too.
245+
local_sym_tab[var_name].shadows_global_from = (
246+
getattr(stmt, "end_lineno", None) or target.lineno
247+
)
248+
249+
250+
def _allocate_for_value(builder, var_name, rval, local_sym_tab, compilation_context):
251+
"""Allocate a slot for `var_name = rval`, typed from the value."""
252+
if isinstance(rval, ast.Call):
253+
_allocate_for_call(builder, var_name, rval, local_sym_tab, compilation_context)
254+
elif isinstance(rval, ast.Constant):
255+
_allocate_for_constant(builder, var_name, rval, local_sym_tab)
256+
elif isinstance(rval, ast.BinOp):
257+
_allocate_for_binop(builder, var_name, rval, local_sym_tab, compilation_context)
258+
elif isinstance(rval, ast.Name):
259+
# Variable-to-variable assignment (b = a)
260+
_allocate_for_name(builder, var_name, rval, local_sym_tab, compilation_context)
261+
elif isinstance(rval, ast.Attribute):
262+
# Struct field-to-variable assignment (a = dat.fld)
263+
_allocate_for_attribute(
264+
builder, var_name, rval, local_sym_tab, compilation_context
265+
)
266+
else:
267+
logger.warning(
268+
f"Unsupported assignment value type for {var_name}: {type(rval).__name__}"
269+
)
103270

104271

105272
def _allocate_for_call(builder, var_name, rval, local_sym_tab, compilation_context):

‎pythonbpf/context.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,12 @@ def __init__(self, module: ir.Module):
8282
# Current function context (optional, if needed globally during function processing)
8383
self.current_func = None
8484

85+
# Enclosing loops of the statement being lowered, innermost last, as
86+
# (continue target, break target) blocks.
87+
self.loop_stack: list[tuple[ir.Block, ir.Block]] = []
88+
8589
def reset(self):
8690
"""Reset state between functions if necessary, though new context per compile is preferred."""
8791
self.scratch_pool.reset()
8892
self.current_func = None
93+
self.loop_stack = []

0 commit comments

Comments
 (0)