Skip to content

Commit aa321db

Browse files
r41k0uclaude
andcommitted
Core: Lower augmented assignment directly instead of desugaring
The first version synthesised an Assign(BinOp(...)) node and fed it back through handle_assign. That is a standard compiler move, but it is not how this compiler works: every pass walks the tree the user wrote, and a node invented mid-codegen is invisible to the allocation pass that already ran, carries no source location for diagnostics, and shows up in logs as a statement nobody typed. handle_aug_assign now resolves the target slot (declared global, local, or struct field), loads it, evaluates the right-hand side, applies the operator, and stores back. Semantic agreement with "x = x op v" no longer rests on shared AST but on shared value-level helpers: the RHS goes through get_operand_value like any other read, and the operator table is extracted from binary-op evaluation into apply_binop so it exists exactly once. Two behavioural refinements fall out: augmented assignment to a global name without a global statement is now the same loud SyntaxError plain assignment gives (previously it slipped past the allocation pass, which only walks Assign), and struct-field targets compute their GEP once rather than twice. Target read precedes RHS evaluation, matching Python. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01L1PX8EuP9C3o3veWGA84RF
1 parent 274bac9 commit aa321db

4 files changed

Lines changed: 117 additions & 44 deletions

File tree

‎pythonbpf/expr/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from .expr_pass import eval_expr, handle_expr, get_operand_value
22
from .type_normalization import convert_to_bool, get_base_type_and_depth
3-
from .ir_ops import deref_to_depth, access_struct_field
3+
from .ir_ops import deref_to_depth, access_struct_field, apply_binop
44
from .call_registry import CallHandlerRegistry
55
from .vmlinux_registry import VmlinuxHandlerRegistry
66

@@ -10,6 +10,7 @@
1010
"convert_to_bool",
1111
"get_base_type_and_depth",
1212
"deref_to_depth",
13+
"apply_binop",
1314
"access_struct_field",
1415
"get_operand_value",
1516
"CallHandlerRegistry",

‎pythonbpf/expr/expr_pass.py‎

Lines changed: 18 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
from pythonbpf.type_deducer import ctypes_to_ir, is_ctypes
88
from .call_registry import CallHandlerRegistry
9-
from .ir_ops import deref_to_depth, access_struct_field
9+
from .ir_ops import deref_to_depth, access_struct_field, apply_binop
1010
from .type_normalization import (
1111
convert_to_bool,
1212
handle_comparator,
@@ -31,8 +31,7 @@ def _handle_name_expr(
3131
val = builder.load(var)
3232
return val, local_sym_tab[expr.id].ir_type
3333
elif expr.id in compilation_context.bpf_globals:
34-
# A @bpfglobal: read straight off the global symbol, exactly the
35-
# `load i64, ptr @counter` form clang emits (tests/c-form/global_vars).
34+
# A @bpfglobal
3635
sym = compilation_context.bpf_globals[expr.id]
3736
val = builder.load(sym.var)
3837
return val, sym.ir_type
@@ -86,7 +85,9 @@ def _handle_attribute_expr(
8685
var_ptr, var_type, var_metadata = local_sym_tab[var_name]
8786
logger.info(f"Loading attribute {attr_name} from variable {var_name}")
8887
logger.info(
89-
f"Variable type: {var_type}, Variable ptr: {var_ptr}, Variable Metadata: {var_metadata}"
88+
f"Variable type: {var_type}, Variable ptr: {
89+
var_ptr
90+
}, Variable Metadata: {var_metadata}"
9091
)
9192
if (
9293
hasattr(var_metadata, "__module__")
@@ -104,7 +105,9 @@ def _handle_attribute_expr(
104105

105106
elif isinstance(var_metadata, Field):
106107
logger.error(
107-
f"Cannot access field '{attr_name}' on already-loaded field value '{var_name}'"
108+
f"Cannot access field '{attr_name}' on already-loaded field value '{
109+
var_name
110+
}'"
108111
)
109112
return None
110113

@@ -234,25 +237,7 @@ def _handle_binary_op_impl(func, compilation_context, rval, builder, local_sym_t
234237
right = builder.sext(right, ir.IntType(64))
235238

236239
# Map AST operation nodes to LLVM IR builder methods
237-
op_map = {
238-
ast.Add: builder.add,
239-
ast.Sub: builder.sub,
240-
ast.Mult: builder.mul,
241-
ast.Div: builder.sdiv,
242-
ast.Mod: builder.srem,
243-
ast.LShift: builder.shl,
244-
ast.RShift: builder.lshr,
245-
ast.BitOr: builder.or_,
246-
ast.BitXor: builder.xor,
247-
ast.BitAnd: builder.and_,
248-
ast.FloorDiv: builder.udiv,
249-
}
250-
251-
if type(op) in op_map:
252-
result = op_map[type(op)](left, right)
253-
return result
254-
else:
255-
raise SyntaxError("Unsupported binary operation")
240+
return apply_binop(builder, op, left, right)
256241

257242

258243
def _handle_binary_op(
@@ -315,7 +300,9 @@ def _handle_ctypes_call(
315300
# Get the IR type from the value itself
316301
actual_ir_type = value.type
317302
logger.info(
318-
f"Converting vmlinux field {val_type.name} (IR type: {actual_ir_type}) to {call_type}"
303+
f"Converting vmlinux field {val_type.name} (IR type: {actual_ir_type}) to {
304+
call_type
305+
}"
319306
)
320307
else:
321308
actual_ir_type = val_type
@@ -328,7 +315,9 @@ def _handle_ctypes_call(
328315
if actual_ir_type.width < expected_type.width:
329316
value = builder.sext(value, expected_type)
330317
logger.info(
331-
f"Sign-extended from i{actual_ir_type.width} to i{expected_type.width}"
318+
f"Sign-extended from i{actual_ir_type.width} to i{
319+
expected_type.width
320+
}"
332321
)
333322
elif actual_ir_type.width > expected_type.width:
334323
value = builder.trunc(value, expected_type)
@@ -340,7 +329,9 @@ def _handle_ctypes_call(
340329
pass
341330
else:
342331
raise ValueError(
343-
f"Type mismatch: expected {expected_type}, got {actual_ir_type} (original type: {val_type})"
332+
f"Type mismatch: expected {expected_type}, got {
333+
actual_ir_type
334+
} (original type: {val_type})"
344335
)
345336

346337
return value, expected_type

‎pythonbpf/expr/ir_ops.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,37 @@
1+
import ast
12
import logging
23
from llvmlite import ir
34

45
logger = logging.getLogger(__name__)
56

67

8+
BINOP_METHODS = {
9+
ast.Add: "add",
10+
ast.Sub: "sub",
11+
ast.Mult: "mul",
12+
ast.Div: "sdiv",
13+
ast.Mod: "srem",
14+
ast.LShift: "shl",
15+
ast.RShift: "lshr",
16+
ast.BitOr: "or_",
17+
ast.BitXor: "xor",
18+
ast.BitAnd: "and_",
19+
ast.FloorDiv: "udiv",
20+
}
21+
22+
23+
def apply_binop(builder, op, left, right):
24+
"""Emit the LLVM instruction for a Python binary operator.
25+
26+
Shared by binary-op evaluation and augmented assignment so the operator
27+
table exists exactly once.
28+
"""
29+
method = BINOP_METHODS.get(type(op))
30+
if method is None:
31+
raise SyntaxError("Unsupported binary operation")
32+
return getattr(builder, method)(left, right)
33+
34+
735
def deref_to_depth(func, builder, val, target_depth):
836
"""Dereference a pointer to a certain depth."""
937

‎pythonbpf/functions/functions_pass.py‎

Lines changed: 69 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
eval_expr,
1212
handle_expr,
1313
convert_to_bool,
14+
get_operand_value,
15+
apply_binop,
1416
VmlinuxHandlerRegistry,
1517
)
1618
from pythonbpf.assign_pass import (
@@ -190,30 +192,81 @@ def handle_assign(func, compilation_context, builder, stmt, local_sym_tab):
190192

191193

192194
def handle_aug_assign(func, compilation_context, builder, stmt, local_sym_tab):
193-
"""Handle `x += v` by desugaring to `x = x op v` and reusing handle_assign.
194-
195-
That is the statement's Python semantics for the targets we support, and it
196-
means globals come along for free: `counter += 1` under `global counter`
197-
becomes load/add/store on @counter.
195+
"""Handle `x += v` and friends by direct lowering: resolve the target's
196+
slot, load it, apply the operator, store back.
197+
198+
No desugaring into synthetic Assign/BinOp nodes: every pass in this
199+
compiler walks the tree the user wrote, and nodes invented mid-codegen are
200+
invisible to the passes that already ran and carry no source locations.
201+
Semantic agreement with `x = x op v` comes from sharing the value-level
202+
helpers instead — the RHS goes through get_operand_value like any other
203+
read, and the operator table is apply_binop, the same one binary-op
204+
evaluation uses.
198205
"""
199206
if isinstance(stmt.target, ast.Name):
200-
load_target = ast.Name(id=stmt.target.id, ctx=ast.Load())
201-
elif isinstance(stmt.target, ast.Attribute):
202-
load_target = ast.Attribute(
203-
value=stmt.target.value, attr=stmt.target.attr, ctx=ast.Load()
204-
)
207+
name = stmt.target.id
208+
if name in compilation_context.current_func_globals:
209+
sym = compilation_context.bpf_globals[name]
210+
slot, slot_type = sym.var, sym.ir_type
211+
elif name in compilation_context.bpf_globals:
212+
raise SyntaxError(
213+
f"augmented assignment to '{name}' shadows the BPF global of "
214+
f"the same name — add 'global {name}' to write to it"
215+
)
216+
elif name in local_sym_tab:
217+
slot = local_sym_tab[name].var
218+
slot_type = local_sym_tab[name].ir_type
219+
else:
220+
raise SyntaxError(f"augmented assignment to undefined variable '{name}'")
221+
elif isinstance(stmt.target, ast.Attribute) and isinstance(
222+
stmt.target.value, ast.Name
223+
):
224+
var_name, field_name = stmt.target.value.id, stmt.target.attr
225+
if var_name not in local_sym_tab:
226+
raise SyntaxError(
227+
f"augmented assignment to field of undefined variable '{var_name}'"
228+
)
229+
metadata = local_sym_tab[var_name].metadata
230+
structs_sym_tab = compilation_context.structs_sym_tab
231+
if metadata not in structs_sym_tab:
232+
raise SyntaxError(
233+
f"augmented assignment to '{var_name}.{field_name}' is only "
234+
f"supported for struct fields"
235+
)
236+
struct_info = structs_sym_tab[metadata]
237+
if field_name not in struct_info.fields:
238+
raise SyntaxError(f"Field '{field_name}' not found in struct '{metadata}'")
239+
slot = struct_info.gep(builder, local_sym_tab[var_name].var, field_name)
240+
slot_type = struct_info.field_type(field_name)
205241
else:
206242
raise SyntaxError(
207243
f"Unsupported augmented-assignment target: {ast.dump(stmt.target)}"
208244
)
209245

210-
desugared = ast.Assign(
211-
targets=[stmt.target],
212-
value=ast.BinOp(left=load_target, op=stmt.op, right=stmt.value),
246+
if not isinstance(slot_type, ir.IntType):
247+
raise SyntaxError(
248+
f"augmented assignment needs an integer target, got {slot_type}"
249+
)
250+
251+
# Python evaluates the target's current value before the right-hand side.
252+
current = builder.load(slot)
253+
rhs = get_operand_value(
254+
func, compilation_context, stmt.value, builder, local_sym_tab
213255
)
214-
ast.copy_location(desugared, stmt)
215-
ast.fix_missing_locations(desugared)
216-
handle_assign(func, compilation_context, builder, desugared, local_sym_tab)
256+
if rhs is None:
257+
raise SyntaxError(
258+
f"Failed to evaluate augmented-assignment value: {ast.dump(stmt.value)}"
259+
)
260+
# Same width discipline as binary-op evaluation: compute in i64, narrow
261+
# back to the slot's width on the way out.
262+
if current.type.width < 64:
263+
current = builder.sext(current, ir.IntType(64))
264+
if isinstance(rhs.type, ir.IntType) and rhs.type.width < 64:
265+
rhs = builder.sext(rhs, ir.IntType(64))
266+
result = apply_binop(builder, stmt.op, current, rhs)
267+
if result.type.width > slot_type.width:
268+
result = builder.trunc(result, slot_type)
269+
builder.store(result, slot)
217270

218271

219272
def handle_cond(func, compilation_context, builder, cond, local_sym_tab):

0 commit comments

Comments
 (0)