Core: Write struct fields through a map-lookup pointer, null-checked

`s = m.lookup(k); s.count = 5` (and `s.count += 1`) compiled without an
error to IR llc rejects: "invalid getelementptr indices" on
`getelementptr inbounds i64*, i64** %s, i32 0, i32 0`. Both the
assignment and augmented-assignment paths indexed into the local's own
slot, which for a lookup or cast local holds a pointer to the struct, not
the struct. Reads already got this right in access_struct_field.

with_struct_field_ptr now resolves a field for a write the same way:
directly for a struct local; for a pointer local, by loading the pointer,
casting it to the struct type, and emitting the store only on the
non-null path (emit_if_not_null, the write-side counterpart of
_null_checked_operation). The write lands in the map entry itself, so
`s.count += 1` is an in-place update with no update() call, and the
maps/structs guide examples that did update(k, s) after a field
assignment now just assign.

Author: varun-r-mallya

docs/user-guide/maps.md

@@ -353,8 +353,8 @@ def track_stats(ctx: c_void_p) -> c_int64:
     stats = process_stats.lookup(process_id)
 
     if stats:
-        stats.count = stats.count + 1
-        process_stats.update(process_id, stats)
+        # A field assignment writes into the map entry itself
+        stats.count += 1
     else:
         new_stats = Stats()
         new_stats.count = 1

docs/user-guide/structs.md

@@ -161,9 +161,8 @@ def track_syscalls(ctx: c_void_p) -> c_int64:
     s = stats.lookup(process_id)
 
     if s:
-        # Update existing stats
-        s.syscall_count = s.syscall_count + 1
-        stats.update(process_id, s)
+        # A field assignment writes into the map entry itself
+        s.syscall_count += 1
     else:
         # Create new stats
         new_stats = ProcessStats()

pythonbpf/assign_pass.py

@@ -3,7 +3,7 @@
 from inspect import isclass
 
 from llvmlite import ir
-from pythonbpf.expr import eval_expr, convert
+from pythonbpf.expr import eval_expr, convert, with_struct_field_ptr
 from pythonbpf.helper import emit_probe_read_kernel_str_call
 from pythonbpf.type_deducer import ctypes_to_ir
 from pythonbpf.vmlinux_parser.dependency_node import Field
@@ -30,8 +30,7 @@ def handle_struct_field_assignment(
         logger.error(f"Field '{field_name}' not found in struct '{struct_type}'")
         return
 
-    # Get field pointer and evaluate value
-    field_ptr = struct_info.gep(builder, local_sym_tab[var_name].var, field_name)
+    # Python evaluates the value before the target.
     field_type = struct_info.field_type(field_name)
     val_result = eval_expr(func, compilation_context, builder, rval, local_sym_tab)
 
@@ -43,24 +42,29 @@ def handle_struct_field_assignment(
 
     # Special case: i8* string to [N x i8] char array
     if _is_char_array(field_type) and _is_i8_ptr(val_type):
-        _copy_string_to_char_array(
-            func,
-            builder,
-            val,
-            field_ptr,
-            field_type,
-            local_sym_tab,
-        )
-        logger.info(f"Copied string to char array {var_name}.{field_name}")
-        return
 
-    # Same implicit widening/truncation as assignment to a local: expressions
-    # evaluate in i64, but a field may be narrower.
-    if isinstance(val_type, ir.IntType) and isinstance(field_type, ir.IntType):
-        val = convert(builder, val, val_type, field_type)
+        def store(builder, field_ptr):
+            _copy_string_to_char_array(
+                func,
+                builder,
+                val,
+                field_ptr,
+                field_type,
+                local_sym_tab,
+            )
+
+    else:
+        # Same implicit widening/truncation as assignment to a local:
+        # expressions evaluate in i64, but a field may be narrower.
+        if isinstance(val_type, ir.IntType) and isinstance(field_type, ir.IntType):
+            val = convert(builder, val, val_type, field_type)
+
+        def store(builder, field_ptr):
+            builder.store(val, field_ptr)
 
-    # Regular assignment
-    builder.store(val, field_ptr)
+    with_struct_field_ptr(
+        func, builder, local_sym_tab[var_name], struct_info, field_name, store
+    )
     logger.info(f"Assigned to struct field {var_name}.{field_name}")
 
 

pythonbpf/expr/__init__.py

@@ -7,7 +7,7 @@
     to_promoted,
 )
 from .operators import usual_arithmetic_conversions
-from .ir_ops import deref_to_depth, access_struct_field
+from .ir_ops import deref_to_depth, access_struct_field, with_struct_field_ptr
 from .operators import apply_binop
 from .call_registry import CallHandlerRegistry
 from .vmlinux_registry import VmlinuxHandlerRegistry
@@ -25,6 +25,7 @@
     "deref_to_depth",
     "apply_binop",
     "access_struct_field",
+    "with_struct_field_ptr",
     "CallHandlerRegistry",
     "VmlinuxHandlerRegistry",
 ]

pythonbpf/expr/ir_ops.py

@@ -62,6 +62,42 @@ def _null_checked_operation(func, builder, ptr, operation, result_type, name_pre
     return phi
 
 
+def emit_if_not_null(func, builder, ptr, emit, name_prefix):
+    """Emit `emit(builder)` only on the path where `ptr` is non-null. The
+    write-side counterpart of _null_checked_operation: a store has no value to
+    merge, so the null path simply skips it."""
+    not_null_block = func.append_basic_block(name=f"{name_prefix}_not_null")
+    merge_block = func.append_basic_block(name=f"{name_prefix}_merge")
+
+    is_not_null = builder.icmp_signed("!=", ptr, ir.Constant(ptr.type, None))
+    builder.cbranch(is_not_null, not_null_block, merge_block)
+
+    builder.position_at_end(not_null_block)
+    emit(builder)
+    builder.branch(merge_block)
+
+    builder.position_at_end(merge_block)
+
+
+def with_struct_field_ptr(func, builder, symbol, struct_info, field_name, emit):
+    """Call `emit(builder, field_ptr)` with a pointer to `field_name` of the
+    struct `symbol` holds. A struct local holds the struct itself; a
+    map-lookup or cast local holds a pointer to it, which may be null, and
+    `emit` then runs only when it is not, as the read of the field yields 0
+    there instead (see access_struct_field)."""
+    if not isinstance(symbol.ir_type, ir.PointerType):
+        emit(builder, struct_info.gep(builder, symbol.var, field_name))
+        return
+
+    struct_ptr = builder.load(symbol.var)
+
+    def emit_through_ptr(builder):
+        typed_ptr = builder.bitcast(struct_ptr, struct_info.ir_type.as_pointer())
+        emit(builder, struct_info.gep(builder, typed_ptr, field_name))
+
+    emit_if_not_null(func, builder, struct_ptr, emit_through_ptr, f"field_{field_name}")
+
+
 def access_struct_field(
     builder, var_ptr, var_type, var_metadata, field_name, structs_sym_tab, func=None
 ):

pythonbpf/functions/functions_pass.py

@@ -17,6 +17,8 @@
     to_promoted,
     canonicalise,
     usual_arithmetic_conversions,
+    access_struct_field,
+    with_struct_field_ptr,
     VmlinuxHandlerRegistry,
 )
 from pythonbpf.assign_pass import (
@@ -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.
     """
+    # Set for a field reached through a pointer (a map-lookup or cast local):
+    # it is read and written through null checks rather than as a plain slot.
+    field_via_pointer = 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
@@ -250,8 +255,12 @@ def handle_aug_assign(func, compilation_context, builder, stmt, local_sym_tab):
         struct_info = structs_sym_tab[metadata]
         if field_name not in struct_info.fields:
             raise SyntaxError(f"Field '{field_name}' not found in struct '{metadata}'")
-        slot = struct_info.gep(builder, local_sym_tab[var_name].var, field_name)
+        symbol = local_sym_tab[var_name]
         slot_type = struct_info.field_type(field_name)
+        if isinstance(symbol.ir_type, ir.PointerType):
+            field_via_pointer = (symbol, struct_info, field_name)
+        else:
+            slot = struct_info.gep(builder, symbol.var, field_name)
     elif isinstance(stmt.target, ast.Attribute):
         raise SyntaxError(
             f"augmented assignment to a nested struct field "
@@ -269,7 +278,19 @@ def handle_aug_assign(func, compilation_context, builder, stmt, local_sym_tab):
         )
 
     # Python evaluates the target's current value before the right-hand side.
-    current = builder.load(slot)
+    if field_via_pointer is not None:
+        symbol, struct_info, field_name = field_via_pointer
+        current, _ = access_struct_field(
+            builder,
+            symbol.var,
+            symbol.ir_type,
+            symbol.metadata,
+            field_name,
+            compilation_context.structs_sym_tab,
+            func,
+        )
+    else:
+        current = builder.load(slot)
     rhs, rhs_ty = get_typed_operand(
         func, compilation_context, stmt.value, builder, local_sym_tab
     )
@@ -287,7 +308,16 @@ 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 field_via_pointer is not None:
+        with_struct_field_ptr(
+            func,
+            builder,
+            *field_via_pointer,
+            lambda builder, field_ptr: builder.store(result, field_ptr),
+        )
+    else:
+        builder.store(result, slot)
 
 
 def handle_cond(func, compilation_context, builder, cond, local_sym_tab):

tests/passing_tests/assign/struct_map_value_field.py

@@ -0,0 +1,41 @@
+# Assignment to a field of a struct map value writes through the lookup
+# pointer into the map, null-checked: `stats.count += 1` updates the entry in
+# place, with no update() call. Also a sub-64-bit field and a plain store.
+from pythonbpf import bpf, map, section, bpfglobal, compile, struct
+from pythonbpf.maps import HashMap
+from ctypes import c_void_p, c_int64, c_uint64, c_uint32
+
+
+@bpf
+@struct
+class stats_t:
+    count: c_uint64
+    flags: c_uint32
+
+
+@bpf
+@map
+def stats() -> HashMap:
+    return HashMap(key=c_uint32, value=stats_t, max_entries=16)
+
+
+@bpf
+@section("tracepoint/raw_syscalls/sys_enter")
+def prog(ctx: c_void_p) -> c_int64:
+    k = c_uint32(0)
+    s = stats.lookup(k)
+    if s:
+        s.count += 1
+        s.flags = 3
+        s.flags <<= 1
+        return s.count
+    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"],
     ),
+    # A field of a struct map value is written through the lookup pointer,
+    # behind a null check, never by indexing into the local's own slot.
+    "assign/struct_map_value_field.py": (
+        [r"field_count_not_null", r"field_flags_not_null", r"store i32 .*, i32\* %"],
+        [r"getelementptr inbounds i64\*, i64\*\*"],
+    ),
     # `return p` on a map lookup dereferences through a null check and returns
     # the i64, never the pointer.
     "return/map_value.py": (