Fix register state dump keyerror for register_state_dump test

pythonbpf/allocation_pass.py

@@ -349,6 +349,9 @@ def _allocate_for_attribute(
             vmlinux_struct_name
         ):
             # Handle vmlinux struct field access
+            # 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}'"
@@ -370,16 +373,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/debuginfo/debug_info_generator.py

@@ -151,6 +151,21 @@ def create_struct_type(
             is_distinct=is_distinct,
         )
 
+    def create_union_type(
+        self, members: List[Any], size: int, is_distinct: bool
+    ) -> Any:
+        """Create an unnamed union type with the given members and size"""
+        return self.module.add_debug_info(
+            "DICompositeType",
+            {
+                "tag": dc.DW_TAG_union_type,
+                "file": self.module._file_metadata,
+                "size": size,
+                "elements": members,
+            },
+            is_distinct=is_distinct,
+        )
+
     def create_struct_type_with_name(
         self, name: str, members: List[Any], size: int, is_distinct: bool
     ) -> Any:

pythonbpf/vmlinux_parser/class_handler.py

@@ -43,6 +43,120 @@ def unwrap_pointer_type(type_obj: Any) -> Any:
     return current_type
 
 
+def _is_flattenable_scalar(member_type: Any) -> bool:
+    """
+    Report whether an anonymous member's member can be lifted into the parent.
+
+    Only plain ctypes scalars qualify. Pointers, arrays, function pointers and
+    nested vmlinux structs are deliberately excluded: they need the containing
+    type / array length bookkeeping that the top-level field walk does, and
+    half-registering them would produce silently wrong relocations.
+    """
+    if getattr(member_type, "__module__", None) != ctypes.__name__:
+        return False
+    if not isinstance(member_type, type):
+        return False
+    if issubclass(member_type, (ctypes._Pointer, ctypes.Array)):
+        return False
+    is_function_pointer = hasattr(member_type, "_restype_") and hasattr(
+        member_type, "_argtypes_"
+    )
+    return not is_function_pointer
+
+
+def flatten_anonymous_members(class_obj, dep_node) -> None:
+    """
+    Register the members of a struct's anonymous members as fields of the struct.
+
+    C lets an anonymous struct/union member's members be named directly on the
+    parent (`regs->cs` where `cs` lives inside an anonymous union), and ctypes
+    mirrors that by exposing them on the parent class. The top-level `_fields_`
+    walk only sees the anonymous member itself (`_0`), so those names would
+    otherwise be unknown to the compiler.
+
+    Each flattened member keeps the anonymous member registered as before, is
+    given the ABSOLUTE offset ctypes reports for it on the parent, and records
+    its CO-RE access path as (index of the anonymous member in the parent,
+    index of the member within the anonymous member).
+
+    Members that overlap are expected and correct: in `struct pt_regs`, `cs`
+    and `csx` are both at offset 136. Each keeps its own field.
+
+    This is a no-op for any struct without `_anonymous_`.
+    """
+    anonymous_names = getattr(class_obj, "_anonymous_", None)
+    if not anonymous_names:
+        return
+    declared_fields = getattr(class_obj, "_fields_", None)
+    if not declared_fields:
+        return
+
+    anonymous_names = set(anonymous_names)
+    for parent_index, declared_field in enumerate(declared_fields):
+        parent_name = declared_field[0]
+        if parent_name not in anonymous_names:
+            continue
+        parent_type = declared_field[1]
+        member_fields = getattr(parent_type, "_fields_", None)
+        if not member_fields:
+            logger.warning(
+                f"Anonymous member {dep_node.name}.{parent_name} has no _fields_, "
+                "its members will not be accessible"
+            )
+            continue
+
+        for member_index, member in enumerate(member_fields):
+            member_name = member[0]
+            member_type = member[1]
+            member_bitfield_size = member[2] if len(member) == 3 else None
+
+            if member_name in dep_node.fields:
+                # Never let a flattened member shadow a field we already have.
+                logger.warning(
+                    f"Anonymous member {dep_node.name}.{parent_name}.{member_name} "
+                    f"collides with an existing field of {dep_node.name}, keeping "
+                    "the existing field"
+                )
+                continue
+            if member_bitfield_size is not None:
+                logger.warning(
+                    f"Skipping bitfield {dep_node.name}.{parent_name}.{member_name}: "
+                    "bitfields inside anonymous members are not supported yet"
+                )
+                continue
+            if not _is_flattenable_scalar(member_type):
+                logger.warning(
+                    f"Skipping {dep_node.name}.{parent_name}.{member_name} of type "
+                    f"{member_type}: only scalar ctypes members of anonymous members "
+                    "are supported"
+                )
+                continue
+            if not hasattr(class_obj, member_name):
+                # ctypes did not lift the name onto the parent, so we have no
+                # authoritative absolute offset for it.
+                logger.warning(
+                    f"Skipping {dep_node.name}.{parent_name}.{member_name}: ctypes "
+                    "does not expose it on the parent struct"
+                )
+                continue
+
+            dep_node.add_field(
+                member_name,
+                member_type,
+                ready=False,
+                access_path=(parent_index, member_index),
+            )
+            dep_node.set_field_bitfield_size(member_name, None)
+            # set_field_ready reads the offset straight off the ctypes struct,
+            # which reports the ABSOLUTE offset for a lifted anonymous member.
+            dep_node.set_field_ready(member_name, is_ready=True)
+            logger.debug(
+                f"Flattened {dep_node.name}.{parent_name}.{member_name} at offset "
+                f"{dep_node.fields[member_name].offset} with access path "
+                f"{(parent_index, member_index)}"
+            )
+
+
 def process_vmlinux_class(
     node,
     llvm_module,
@@ -312,6 +426,11 @@ def process_vmlinux_post_ast(
                     raise ValueError(
                         f"{elem_name} with type {elem_type} from module {module_name} not supported in recursive resolver"
                     )
+
+            # Anonymous struct/union members expose their members on the parent
+            # in C and in ctypes, but they are invisible to the top-level
+            # _fields_ walk above. Register them too. No-op without _anonymous_.
+            flatten_anonymous_members(class_obj, new_dep_node)
     elif module_name == ctypes.__name__ or module_name is None:
         # Handle ctypes types - these don't need processing, just return
         logger.debug(f"Skipping ctypes type {current_symbol_name}")

pythonbpf/vmlinux_parser/dependency_node.py

@@ -17,6 +17,13 @@ class Field:
     offset: int
     value: Any = None
     ready: bool = False
+    # CO-RE access path from the enclosing struct down to this field, as a tuple
+    # of member indices. `None` means the field is a plain top-level member and
+    # its access path is just its own index in the struct. It is only populated
+    # for fields that live inside an anonymous member and were flattened into
+    # the parent, e.g. `struct pt_regs.cs` is member 0 of anonymous member 17,
+    # so its access_path is (17, 0).
+    access_path: Optional[tuple[int, ...]] = None
 
     def __hash__(self):
         """
@@ -154,6 +161,7 @@ def add_field(
         bitfield_size: Optional[int] = None,
         ready: bool = False,
         offset: int = 0,
+        access_path: Optional[tuple[int, ...]] = None,
     ) -> None:
         """Add a field to the node with an optional initial value and readiness state."""
         if self.depends_on is None:
@@ -168,6 +176,7 @@ def add_field(
             ctype_complex_type=ctype_complex_type,
             bitfield_size=bitfield_size,
             offset=offset,
+            access_path=access_path,
         )
         # Invalidate readiness cache
         self._ready_cache = None

pythonbpf/vmlinux_parser/ir_gen/debug_info_gen.py

@@ -43,10 +43,33 @@ def debug_info_generation(
     # Process all fields and create members for the struct
     members = []
 
-    sorted_fields = sorted(struct.fields.items(), key=lambda item: item[1].offset)
+    # Members lifted out of an anonymous member are not members of this struct in
+    # DWARF terms; they belong to the anonymous member's own type. Emitting them
+    # here would both duplicate offsets and shift every later member's index.
+    sorted_fields = sorted(
+        (item for item in struct.fields.items() if item[1].access_path is None),
+        key=lambda item: item[1].offset,
+    )
+
+    anonymous_names = set(getattr(struct.ctype_struct, "_anonymous_", None) or ())
 
     for field_name, field in sorted_fields:
         try:
+            if field_name in anonymous_names:
+                # An anonymous member has to reach BTF unnamed and with its own
+                # members, or a `$...:<n>:<m>` access string cannot be resolved
+                # against it at load time.
+                anonymous_member = _anonymous_member_debug_info(
+                    field, generator, generated_debug_info
+                )
+                if anonymous_member is not None:
+                    members.append(anonymous_member)
+                    continue
+                logger.warning(
+                    f"Could not describe anonymous member {struct.name}.{field_name}, "
+                    "falling back to an opaque member"
+                )
+
             # Get appropriate debug type for this field
             field_type = _get_field_debug_type(
                 field_name, field, generator, struct, generated_debug_info
@@ -75,6 +98,89 @@ def debug_info_generation(
     return struct_type
 
 
+def _lookup_generated_debug_info(
+    type_name: str, generated_debug_info: List[Tuple[DependencyNode, Any]]
+):
+    """Find already generated debug info for a vmlinux type by name."""
+    for existing_struct, debug_info in generated_debug_info:
+        if existing_struct.name == type_name:
+            return debug_info, existing_struct.__sizeof__() * 8
+    return None
+
+
+def _anonymous_member_debug_info(
+    field,
+    generator: DebugInfoGenerator,
+    generated_debug_info: List[Tuple[DependencyNode, Any]],
+):
+    """
+    Describe an anonymous struct/union member the way a C compiler would.
+
+    The member itself is emitted WITHOUT a name, so it lands in BTF as `(anon)`,
+    and its type is a real composite carrying its own members in declaration
+    order. Both matter: libbpf walks a CO-RE access string by member index into
+    the local type and then matches by name in the target type, so an opaque or
+    named stand-in makes any access through the anonymous member unresolvable.
+
+    Returns None if the member cannot be described faithfully, in which case the
+    caller keeps the previous behaviour.
+    """
+    anonymous_type = field.type
+    declared = getattr(anonymous_type, "_fields_", None)
+    if not declared or not isinstance(anonymous_type, type):
+        return None
+
+    inner_members = []
+    for declared_member in declared:
+        if len(declared_member) != 2:
+            # Bitfield members need BTF bitfield encoding we do not emit yet.
+            return None
+        member_name, member_type = declared_member
+        try:
+            member_offset_bits = getattr(anonymous_type, member_name).offset * 8
+        except AttributeError:
+            return None
+
+        if getattr(member_type, "__module__", None) == "vmlinux":
+            member_type_name = getattr(member_type, "__name__", None)
+            member_debug_type = (
+                _lookup_generated_debug_info(member_type_name, generated_debug_info)
+                if member_type_name
+                else None
+            )
+            if member_debug_type is None:
+                # Keep the member so the indices stay right, but leave its type
+                # opaque, exactly as nested structs are handled elsewhere.
+                member_debug_type = (
+                    generator.create_struct_type([], 0, is_distinct=True),
+                    0,
+                )
+        else:
+            member_debug_type = _get_basic_debug_type(member_type, generator)
+        if not isinstance(member_debug_type, tuple) or len(member_debug_type) != 2:
+            return None
+
+        inner_members.append(
+            generator.create_struct_member_vmlinux(
+                member_name, member_debug_type, member_offset_bits
+            )
+        )
+
+    size_bits = ctypes.sizeof(anonymous_type) * 8
+    if issubclass(anonymous_type, ctypes.Union):
+        composite = generator.create_union_type(
+            inner_members, size_bits, is_distinct=True
+        )
+    else:
+        composite = generator.create_struct_type(
+            inner_members, size_bits, is_distinct=True
+        )
+
+    return generator.create_struct_member_vmlinux(
+        "", (composite, size_bits), field.offset * 8
+    )
+
+
 def _get_field_debug_type(
     field_name: str,
     field,