Core: Augmented assignment to a map-lookup local rebinds it, as v = v + 1 does

`v = m.lookup(k); v += 1` was rejected with "augmented assignment needs
an integer target, got i64*": the Name branch took the local's slot type,
the pointer, as the target type. `v = v + 1` already compiles, and to a
rebinding: the value is read through the pointer, null-checked, and the
local is pointed at its spare `v_tmp` slot holding the result; the map is
not written. `v += 1` now does the same. The read is get_typed_operand's,
so the operation takes the map's declared value type and sign (u64 >>= is
an lshr); the write is a rebind_map_value_local helper, now shared with
plain assignment so the two cannot drift apart. A struct-valued lookup
local is refused with a pointer to its fields.

Author: varun-r-mallya

docs/user-guide/maps.md

@@ -60,6 +60,10 @@ def trace_open(ctx: c_void_p) -> c_int64:
     return 0
 ```
 
+An integer value behaves like a Python `int`: `value += 1`, like
+`value = value + 1`, rebinds the local to the new number and leaves the map
+unchanged. To store it, call `update`.
+
 ##### update(key, value, flags=None)
 
 Update or insert a key-value pair.

pythonbpf/assign_pass.py

@@ -87,6 +87,24 @@ def _copy_string_to_char_array(
     emit_probe_read_kernel_str_call(builder, dst_i8_ptr, array_size, src_ptr)
 
 
+def is_map_value_local(local_sym_tab, var_name):
+    """Whether `var_name` holds a map-lookup result: a pointer to the value,
+    allocated with a spare `<name>_tmp` slot for values computed from it."""
+    return (
+        isinstance(local_sym_tab[var_name].ir_type, ir.PointerType)
+        and f"{var_name}_tmp" in local_sym_tab
+    )
+
+
+def rebind_map_value_local(builder, local_sym_tab, var_name, val):
+    """Bind a map-lookup local to a computed value, the way Python rebinds a
+    name: the value goes into the local's `<name>_tmp` slot and the local is
+    pointed at it. The map itself is not written."""
+    tmp = local_sym_tab[f"{var_name}_tmp"].var
+    builder.store(val, tmp)
+    builder.store(tmp, local_sym_tab[var_name].var)
+
+
 def _is_char_array(ir_type):
     """Check if type is [N x i8]."""
     return (
@@ -223,12 +241,9 @@ def handle_variable_assignment(
                 return False
         elif isinstance(val_type, ir.IntType) and isinstance(var_type, ir.PointerType):
             # NOTE: This is assignment to a PTR_TO_MAP_VALUE_OR_NULL
-            logger.info(
-                f"Creating temporary variable for pointer assignment to {var_name}"
-            )
-            var_ptr_tmp = local_sym_tab[f"{var_name}_tmp"].var
-            builder.store(val, var_ptr_tmp)
-            val = var_ptr_tmp
+            rebind_map_value_local(builder, local_sym_tab, var_name, val)
+            logger.info(f"Rebound map-lookup local {var_name} to a computed value")
+            return True
         else:
             logger.error(
                 f"Type mismatch for variable {var_name}: {val_type} vs {var_type}"

pythonbpf/functions/functions_pass.py

@@ -22,6 +22,8 @@
 from pythonbpf.assign_pass import (
     handle_variable_assignment,
     handle_struct_field_assignment,
+    is_map_value_local,
+    rebind_map_value_local,
 )
 from pythonbpf.allocation_pass import (
     handle_assign_allocation,
@@ -207,6 +209,9 @@ def handle_aug_assign(func, compilation_context, builder, stmt, local_sym_tab):
     read, and the operator table is apply_binop, the same one binary-op
     evaluation uses.
     """
+    # A map-lookup local reads through its pointer and is rebound, not
+    # written through, exactly as `v = v + 1` does it.
+    map_value_local = None
     if isinstance(stmt.target, ast.Name):
         name = stmt.target.id
         # One table: a declared global is a local_sym_tab entry whose slot is
@@ -220,6 +225,13 @@ def handle_aug_assign(func, compilation_context, builder, stmt, local_sym_tab):
                 raise SyntaxError(
                     f"cannot assign to '{name}': it is the context parameter"
                 )
+            if is_map_value_local(local_sym_tab, name):
+                if local_sym_tab[name].metadata in compilation_context.structs_sym_tab:
+                    raise SyntaxError(
+                        f"augmented assignment to '{name}', a struct map value; "
+                        f"update one of its fields instead ({name}.field += ...)"
+                    )
+                map_value_local = name
         elif name in compilation_context.bpf_globals:
             # `x += v` binds x as a local and reads it unbound: UnboundLocalError.
             raise SyntaxError(
@@ -263,13 +275,20 @@ def handle_aug_assign(func, compilation_context, builder, stmt, local_sym_tab):
             f"Unsupported augmented-assignment target: {ast.dump(stmt.target)}"
         )
 
-    if not isinstance(slot_type, ir.IntType):
+    # Python evaluates the target's current value before the right-hand side.
+    if map_value_local is not None:
+        # The null-checked read every other use of the local gets; its
+        # descriptor is the map's declared value type, which is also the
+        # type the result is stored back as.
+        current, slot_type = get_typed_operand(
+            func, compilation_context, stmt.target, builder, local_sym_tab
+        )
+    elif isinstance(slot_type, ir.IntType):
+        current = builder.load(slot)
+    else:
         raise SyntaxError(
             f"augmented assignment needs an integer target, got {slot_type}"
         )
-
-    # Python evaluates the target's current value before the right-hand side.
-    current = builder.load(slot)
     rhs, rhs_ty = get_typed_operand(
         func, compilation_context, stmt.value, builder, local_sym_tab
     )
@@ -287,7 +306,11 @@ def handle_aug_assign(func, compilation_context, builder, stmt, local_sym_tab):
         apply_binop(builder, stmt.op, current, rhs, signedness(result_ty)),
         result_ty,
     )
-    builder.store(convert(builder, result, result_ty, slot_type), slot)
+    result = convert(builder, result, result_ty, slot_type)
+    if map_value_local is not None:
+        rebind_map_value_local(builder, local_sym_tab, map_value_local, result)
+    else:
+        builder.store(result, slot)
 
 
 def handle_cond(func, compilation_context, builder, cond, local_sym_tab):

tests/passing_tests/assign/augassign_map_value.py

@@ -0,0 +1,32 @@
+# Augmented assignment to a map-lookup local: `v += 1` is `v = v + 1`. It reads
+# the value through the pointer (null-checked) and rebinds v to the result; the
+# map itself is not written.
+from pythonbpf import bpf, map, section, bpfglobal, compile
+from pythonbpf.maps import HashMap
+from ctypes import c_void_p, c_int64, c_uint64
+
+
+@bpf
+@map
+def m() -> HashMap:
+    return HashMap(key=c_int64, value=c_uint64, max_entries=4)
+
+
+@bpf
+@section("tracepoint/raw_syscalls/sys_enter")
+def prog(ctx: c_void_p) -> c_int64:
+    v = m.lookup(0)
+    if v:
+        v += 1
+        v >>= 1  # c_uint64 value: a logical shift
+        return v
+    return c_int64(0)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

tests/test_signedness_ir.py

@@ -81,6 +81,12 @@
         [r"trunc i64 .* to i32", r"zext i32 .* to i64"],
         [r"sext i32 .* to i64"],
     ),
+    # `v += 1; v >>= 1` on a c_uint64 map value: read through the pointer,
+    # operate unsigned, rebind v to the lookup local's spare slot.
+    "assign/augassign_map_value.py": (
+        [r"deref_1_not_null", r"\blshr i64", r"store i64\* %\"v_tmp\", i64\*\* %\"v\""],
+        [r"\bashr i64"],
+    ),
     # `return p` on a map lookup dereferences through a null check and returns
     # the i64, never the pointer.
     "return/map_value.py": (