|
11 | 11 | eval_expr, |
12 | 12 | handle_expr, |
13 | 13 | convert_to_bool, |
| 14 | + get_operand_value, |
| 15 | + apply_binop, |
14 | 16 | VmlinuxHandlerRegistry, |
15 | 17 | ) |
16 | 18 | from pythonbpf.assign_pass import ( |
@@ -190,30 +192,81 @@ def handle_assign(func, compilation_context, builder, stmt, local_sym_tab): |
190 | 192 |
|
191 | 193 |
|
192 | 194 | 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. |
198 | 205 | """ |
199 | 206 | 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) |
205 | 241 | else: |
206 | 242 | raise SyntaxError( |
207 | 243 | f"Unsupported augmented-assignment target: {ast.dump(stmt.target)}" |
208 | 244 | ) |
209 | 245 |
|
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 |
213 | 255 | ) |
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) |
217 | 270 |
|
218 | 271 |
|
219 | 272 | def handle_cond(func, compilation_context, builder, cond, local_sym_tab): |
|
0 commit comments