Fix enum constant assignments to locals

pythonbpf/allocation_pass.py

@@ -324,12 +324,19 @@ def _allocate_for_name(builder, var_name, rval, local_sym_tab, compilation_conte
     """Allocate memory for variable-to-variable assignment (b = a)."""
     source_var = rval.id
 
-    # Local first, then a BPF global: the copy takes the source's type either
-    # way (a c_uint32 global gives a c_uint32 local).
+    # Same resolution order as every other read of a bare name: local, then
+    # BPF global, then vmlinux enum constant. A variable source gives the copy
+    # its type; an enum constant is an immediate with no storage, so the copy
+    # gets the i64 slot a literal gets.
     if source_var in local_sym_tab:
         source_symbol = local_sym_tab[source_var]
     elif source_var in compilation_context.bpf_globals:
         source_symbol = compilation_context.bpf_globals[source_var]
+    elif VmlinuxHandlerRegistry.handle_name(source_var) is not None:
+        var = _allocate_with_type(builder, var_name, ir.IntType(64))
+        local_sym_tab[var_name] = LocalSymbol(var, ir.IntType(64))
+        logger.info(f"Pre-allocated {var_name} from enum constant {source_var}")
+        return
     else:
         logger.error(f"Source variable '{source_var}' not found in symbol table")
         return

tests/failing_tests/vmlinux/assignment_handling.py → tests/passing_tests/vmlinux/assignment_handling.py

@@ -7,8 +7,8 @@
 @bpf
 @section("kprobe/blk_mq_start_request")
 def example(ctx: c_void_p) -> c_int64:
-    d = XDP_PASS  # This gives an error, but
-    e = XDP_PASS + 0  # this does not
+    d = XDP_PASS  # bare enum constant: used to be rejected at allocation
+    e = XDP_PASS + 0  # through a binop: always worked; the two must agree
     print(f"test1 {e} test2 {d}")
     return c_int64(0)
 

tests/passing_tests/vmlinux/enum_to_local.py

@@ -0,0 +1,26 @@
+# A local initialised from a vmlinux enum constant, the C shape
+# `int act = XDP_PASS;`. Reads of a bare name resolve local, then BPF global,
+# then vmlinux enum everywhere else; the allocation pass used to stop one step
+# short for this shape and reject it, while `act = XDP_PASS + 0` compiled.
+from ctypes import c_int64
+from pythonbpf import bpf, section, bpfglobal, compile
+from vmlinux import struct_xdp_md
+from vmlinux import XDP_PASS, XDP_DROP
+
+
+@bpf
+@section("xdp")
+def prog(ctx: struct_xdp_md) -> c_int64:
+    act = XDP_PASS
+    if ctx.data_end - ctx.data > 1500:
+        act = XDP_DROP
+    return c_int64(act)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

tests/test_config.toml

@@ -29,7 +29,6 @@
 
 "failing_tests/vmlinux/args_test.py" = {reason = "struct_trace_event_raw_sys_enter args field access not supported", level = "ir"}
 
-"failing_tests/vmlinux/assignment_handling.py" = {reason = "Assigning vmlinux enum value (XDP_PASS) to a local variable not yet supported", level = "ir"}
 
 "failing_tests/xdp_pass.py" = {reason = "XDP program using vmlinux structs (struct_xdp_md) and complex map/struct interaction not yet supported", level = "ir"}