Core: Widen any sub-64-bit context field instead of special-casing xdp_md

Reading a context field narrower than a register was hardcoded to
struct_xdp_md's i32 fields in three places that had to agree with each other:
the destination alloca in allocation_pass, the zext in load_ctx_field, and
the store in assign_pass. struct pt_regs' cs and ss are 2 bytes, so a second
special case would have been needed.

Replace all three with the general rule: a context field is loaded at its
natural width and zero-extended to i64, so its destination is i64.
allocation_pass gains the same context discriminator load_ctx_field uses
(a context argument has no alloca of its own), so non-context field reads,
which go through load_struct_field and bpf_probe_read_kernel, keep their
natural width and are untouched.

struct_xdp_md's generated IR is byte-identical, as is every other vmlinux/xdp
test program's; the only diff in the whole corpus remains the four new globals
from the previous commit.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: r41k0u

pythonbpf/allocation_pass.py

@@ -343,6 +343,9 @@ def _allocate_for_attribute(
         if VmlinuxHandlerRegistry.is_vmlinux_struct(struct_type.__name__):
             # Handle vmlinux struct field access
             vmlinux_struct_name = struct_type.__name__
+            # Same discriminator handle_vmlinux_struct_field uses: a context
+            # argument has no alloca of its own.
+            is_context_field = local_sym_tab[struct_var].var is None
             if not VmlinuxHandlerRegistry.has_field(vmlinux_struct_name, field_name):
                 logger.error(
                     f"Field '{field_name}' not found in vmlinux struct '{vmlinux_struct_name}'"
@@ -364,16 +367,15 @@ def _allocate_for_attribute(
                     field_size_bits = field_size_bytes * 8
 
                     if field_size_bits in [8, 16, 32, 64]:
-                        # Special case: struct_xdp_md i32 fields should allocate as i64
-                        # because load_ctx_field will zero-extend them to i64
-                        if (
-                            vmlinux_struct_name == "struct_xdp_md"
-                            and field_size_bits == 32
-                        ):
+                        # Sub-register-width context fields allocate as i64,
+                        # because load_ctx_field zero-extends them to i64.
+                        # Non-context fields go through load_struct_field, which
+                        # keeps them at their natural width.
+                        if is_context_field and field_size_bits < 64:
                             actual_ir_type = ir.IntType(64)
                             logger.info(
-                                f"Allocating {var_name} as i64 for i32 field from struct_xdp_md.{field_name} "
-                                "(will be zero-extended during load)"
+                                f"Allocating {var_name} as i64 for i{field_size_bits} field from "
+                                f"{vmlinux_struct_name}.{field_name} (will be zero-extended during load)"
                             )
                         else:
                             actual_ir_type = ir.IntType(field_size_bits)

pythonbpf/assign_pass.py

@@ -185,22 +185,24 @@ def handle_variable_assignment(
                 return False
         if isinstance(val_type, Field):
             logger.info("Handling assignment to struct field")
-            # Special handling for struct_xdp_md i32 fields that are zero-extended to i64
-            # The load_ctx_field already extended them, so val is i64 but val_type.type shows c_uint
+            field_ir_type = ctypes_to_ir(val_type.type.__name__)
+            # Sub-register-width context fields are zero-extended to i64 by
+            # load_ctx_field, so val is already i64 even though the field type
+            # says otherwise (c_uint for xdp_md, c_ushort for pt_regs.cs/ss).
             if (
-                hasattr(val_type, "type")
-                and val_type.type.__name__ == "c_uint"
+                isinstance(field_ir_type, ir.IntType)
+                and field_ir_type.width < 64
                 and isinstance(var_type, ir.IntType)
                 and var_type.width == 64
             ):
-                # This is the struct_xdp_md case - value is already i64
                 builder.store(val, var_ptr)
                 logger.info(
-                    f"Assigned zero-extended struct_xdp_md i32 field to {var_name} (i64)"
+                    f"Assigned zero-extended i{field_ir_type.width} context field "
+                    f"to {var_name} (i64)"
                 )
                 return True
             # TODO: handling only ctype struct fields for now. Handle other stuff too later.
-            elif var_type == ctypes_to_ir(val_type.type.__name__):
+            elif var_type == field_ir_type:
                 builder.store(val, var_ptr)
                 logger.info(f"Assigned ctype struct field to {var_name}")
                 return True

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -315,7 +315,7 @@ def load_ctx_field(builder, ctx_arg, offset_global, field_data, struct_name=None
 
         # Determine the appropriate IR type based on field information
         int_width = 64  # Default to 64-bit
-        needs_zext = False  # Track if we need zero-extension for xdp_md
+        needs_zext = False  # Track if we need zero-extension to a full register
 
         if field_data is not None:
             # Try to determine the size from field metadata
@@ -328,12 +328,14 @@ def load_ctx_field(builder, ctx_arg, offset_global, field_data, struct_name=None
                         int_width = field_size_bits
                         logger.info(f"Determined field size: {int_width} bits")
 
-                        # Special handling for struct_xdp_md i32 fields
-                        # Load as i32 but extend to i64 before storing
-                        if struct_name == "struct_xdp_md" and int_width == 32:
+                        # Context fields are loaded at their natural width and
+                        # widened to a full 64-bit register, so that everything
+                        # downstream sees one uniform integer type.
+                        if int_width < 64:
                             needs_zext = True
                             logger.info(
-                                "struct_xdp_md i32 field detected, will zero-extend to i64"
+                                f"i{int_width} field {struct_name} detected, "
+                                "will zero-extend to i64"
                             )
                     else:
                         logger.warning(
@@ -363,10 +365,10 @@ def load_ctx_field(builder, ctx_arg, offset_global, field_data, struct_name=None
         # Load and return the value
         value = builder.load(typed_ptr)
 
-        # Zero-extend i32 to i64 for struct_xdp_md fields
+        # Widen sub-register-width context fields to i64
         if needs_zext:
             value = builder.zext(value, ir.IntType(64))
-            logger.info("Zero-extended i32 value to i64 for struct_xdp_md field")
+            logger.info(f"Zero-extended i{int_width} context field value to i64")
 
         return value