Merge pull request #74 from pythonbpf/fix-vmlinux-ir-gen

Fix some issues with vmlinux struct usage in syntax
TODO: THIS DOES NOT FIX the XDP example. 
The issue with the XDP example here is that the range is not being tracked due to multiple loads from stack which is making the verifier lose track of the range on that value.

.github/workflows/format.yml

@@ -12,7 +12,7 @@ jobs:
     name: Format
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v5
+    - uses: actions/checkout@v6
     - uses: actions/setup-python@v6
       with:
         python-version: "3.x"

.github/workflows/python-publish.yml

@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
 
       - uses: actions/setup-python@v6
         with:

BCC-Examples/disksnoop.ipynb

@@ -0,0 +1,122 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c3520e58-e50f-4bc1-8f9d-a6fecbf6e9f0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from vmlinux import struct_request, struct_pt_regs\n",
+    "from pythonbpf import bpf, section, bpfglobal, map, BPF\n",
+    "from pythonbpf.helper import ktime\n",
+    "from pythonbpf.maps import HashMap\n",
+    "from ctypes import c_int64, c_uint64, c_int32\n",
+    "\n",
+    "REQ_WRITE = 1\n",
+    "\n",
+    "\n",
+    "@bpf\n",
+    "@map\n",
+    "def start() -> HashMap:\n",
+    "    return HashMap(key=c_uint64, value=c_uint64, max_entries=10240)\n",
+    "\n",
+    "\n",
+    "@bpf\n",
+    "@section(\"kprobe/blk_mq_end_request\")\n",
+    "def trace_completion(ctx: struct_pt_regs) -> c_int64:\n",
+    "    # Get request pointer from first argument\n",
+    "    req_ptr = ctx.di\n",
+    "    req = struct_request(ctx.di)\n",
+    "    # Print: data_len, cmd_flags, latency_us\n",
+    "    data_len = req.__data_len\n",
+    "    cmd_flags = req.cmd_flags\n",
+    "    # Lookup start timestamp\n",
+    "    req_tsp = start.lookup(req_ptr)\n",
+    "    if req_tsp:\n",
+    "        # Calculate delta in nanoseconds\n",
+    "        delta = ktime() - req_tsp\n",
+    "\n",
+    "        # Convert to microseconds for printing\n",
+    "        delta_us = delta // 1000\n",
+    "\n",
+    "        print(f\"{data_len} {cmd_flags:x} {delta_us}\\n\")\n",
+    "\n",
+    "        # Delete the entry\n",
+    "        start.delete(req_ptr)\n",
+    "\n",
+    "    return c_int64(0)\n",
+    "\n",
+    "\n",
+    "@bpf\n",
+    "@section(\"kprobe/blk_mq_start_request\")\n",
+    "def trace_start(ctx1: struct_pt_regs) -> c_int32:\n",
+    "    req = ctx1.di\n",
+    "    ts = ktime()\n",
+    "    start.update(req, ts)\n",
+    "    return c_int32(0)\n",
+    "\n",
+    "\n",
+    "@bpf\n",
+    "@bpfglobal\n",
+    "def LICENSE() -> str:\n",
+    "    return \"GPL\"\n",
+    "\n",
+    "\n",
+    "b = BPF()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "97040f73-98e0-4993-94c6-125d1b42d931",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "b.load()\n",
+    "b.attach_all()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b1bd4f51-fa25-42e1-877c-e48a2605189f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pythonbpf import trace_pipe"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "96b4b59b-b0db-4952-9534-7a714f685089",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trace_pipe()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

BCC-Examples/disksnoop.py

@@ -0,0 +1,61 @@
+from vmlinux import struct_request, struct_pt_regs
+from pythonbpf import bpf, section, bpfglobal, compile_to_ir, compile, map
+from pythonbpf.helper import ktime
+from pythonbpf.maps import HashMap
+import logging
+from ctypes import c_int64, c_uint64, c_int32
+
+# Constants
+REQ_WRITE = 1  # from include/linux/blk_types.h
+
+
+@bpf
+@map
+def start() -> HashMap:
+    return HashMap(key=c_uint64, value=c_uint64, max_entries=10240)
+
+
+@bpf
+@section("kprobe/blk_mq_end_request")
+def trace_completion(ctx: struct_pt_regs) -> c_int64:
+    # Get request pointer from first argument
+    req_ptr = ctx.di
+    req = struct_request(ctx.di)
+    # Print: data_len, cmd_flags, latency_us
+    data_len = req.__data_len
+    cmd_flags = req.cmd_flags
+    # Lookup start timestamp
+    req_tsp = start.lookup(req_ptr)
+    if req_tsp:
+        # Calculate delta in nanoseconds
+        delta = ktime() - req_tsp
+
+        # Convert to microseconds for printing
+        delta_us = delta // 1000
+
+        print(f"{data_len} {cmd_flags:x} {delta_us}\n")
+
+        # Delete the entry
+        start.delete(req_ptr)
+
+    return c_int64(0)
+
+
+@bpf
+@section("kprobe/blk_mq_start_request")
+def trace_start(ctx1: struct_pt_regs) -> c_int32:
+    req = ctx1.di
+    ts = ktime()
+    start.update(req, ts)
+    return c_int32(0)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+if __name__ == "__main__":
+    compile_to_ir("disksnoop.py", "disksnoop.ll", loglevel=logging.INFO)
+    compile()

pythonbpf/allocation_pass.py

@@ -114,9 +114,22 @@ def _allocate_for_call(
         # Struct constructors
         elif call_type in structs_sym_tab:
             struct_info = structs_sym_tab[call_type]
-            var = builder.alloca(struct_info.ir_type, name=var_name)
-            local_sym_tab[var_name] = LocalSymbol(var, struct_info.ir_type, call_type)
-            logger.info(f"Pre-allocated {var_name} for struct {call_type}")
+            if len(rval.args) == 0:
+                # Zero-arg constructor: allocate the struct itself
+                var = builder.alloca(struct_info.ir_type, name=var_name)
+                local_sym_tab[var_name] = LocalSymbol(
+                    var, struct_info.ir_type, call_type
+                )
+                logger.info(f"Pre-allocated {var_name} for struct {call_type}")
+            else:
+                # Pointer cast: allocate as pointer to struct
+                ptr_type = ir.PointerType(struct_info.ir_type)
+                var = builder.alloca(ptr_type, name=var_name)
+                var.align = 8
+                local_sym_tab[var_name] = LocalSymbol(var, ptr_type, call_type)
+                logger.info(
+                    f"Pre-allocated {var_name} for struct pointer cast to {call_type}"
+                )
 
         elif VmlinuxHandlerRegistry.is_vmlinux_struct(call_type):
             # When calling struct_name(pointer), we're doing a cast, not construction
@@ -371,6 +384,7 @@ def _allocate_for_attribute(builder, var_name, rval, local_sym_tab, structs_sym_
                         f"Could not determine size for ctypes field {field_name}: {e}"
                     )
                     actual_ir_type = ir.IntType(64)
+                    field_size_bits = 64
 
             # Check if it's a nested vmlinux struct or complex type
             elif field.type.__module__ == "vmlinux":
@@ -379,23 +393,34 @@ def _allocate_for_attribute(builder, var_name, rval, local_sym_tab, structs_sym_
                     field.ctype_complex_type, ctypes._Pointer
                 ):
                     actual_ir_type = ir.IntType(64)  # Pointer is always 64-bit
+                    field_size_bits = 64
                 # For embedded structs, this is more complex - might need different handling
                 else:
                     logger.warning(
                         f"Field {field_name} is a nested vmlinux struct, using i64 for now"
                     )
                     actual_ir_type = ir.IntType(64)
+                    field_size_bits = 64
             else:
                 logger.warning(
                     f"Unknown field type module {field.type.__module__} for {field_name}"
                 )
                 actual_ir_type = ir.IntType(64)
+                field_size_bits = 64
 
-            # Allocate with the actual IR type
+            # Pre-allocate the tmp storage used by load_struct_field (so we don't alloca inside handler)
+            tmp_name = f"{struct_var}_{field_name}_tmp"
+            tmp_ir_type = ir.IntType(field_size_bits)
+            tmp_var = builder.alloca(tmp_ir_type, name=tmp_name)
+            tmp_var.align = tmp_ir_type.width // 8
+            local_sym_tab[tmp_name] = LocalSymbol(tmp_var, tmp_ir_type)
+            logger.info(
+                f"Pre-allocated temp {tmp_name} (i{field_size_bits}) for vmlinux field read {vmlinux_struct_name}.{field_name}"
+            )
+
+            # Allocate with the actual IR type for the destination var
             var = _allocate_with_type(builder, var_name, actual_ir_type)
-            local_sym_tab[var_name] = LocalSymbol(
-                var, actual_ir_type, field
-            )  # <-- Store Field metadata
+            local_sym_tab[var_name] = LocalSymbol(var, actual_ir_type, field)
 
             logger.info(
                 f"Pre-allocated {var_name} as {actual_ir_type} from vmlinux struct {vmlinux_struct_name}.{field_name}"

pythonbpf/assign_pass.py

@@ -174,6 +174,23 @@ def handle_variable_assignment(
                     f"Type mismatch: vmlinux struct pointer requires i64, got {var_type}"
                 )
                 return False
+        # Handle user-defined struct pointer casts
+        # val_type is a string (struct name), var_type is a pointer to the struct
+        if isinstance(val_type, str) and val_type in structs_sym_tab:
+            struct_info = structs_sym_tab[val_type]
+            expected_ptr_type = ir.PointerType(struct_info.ir_type)
+
+            # Check if var_type matches the expected pointer type
+            if isinstance(var_type, ir.PointerType) and var_type == expected_ptr_type:
+                # val is already the correct pointer type from inttoptr/bitcast
+                builder.store(val, var_ptr)
+                logger.info(f"Assigned user-defined struct pointer cast to {var_name}")
+                return True
+            else:
+                logger.error(
+                    f"Type mismatch: user-defined struct pointer cast requires pointer type, got {var_type}"
+                )
+                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

pythonbpf/expr/__init__.py

@@ -1,6 +1,6 @@
 from .expr_pass import eval_expr, handle_expr, get_operand_value
 from .type_normalization import convert_to_bool, get_base_type_and_depth
-from .ir_ops import deref_to_depth
+from .ir_ops import deref_to_depth, access_struct_field
 from .call_registry import CallHandlerRegistry
 from .vmlinux_registry import VmlinuxHandlerRegistry
 
@@ -10,6 +10,7 @@ __all__ = [
     "convert_to_bool",
     "get_base_type_and_depth",
     "deref_to_depth",
+    "access_struct_field",
     "get_operand_value",
     "CallHandlerRegistry",
     "VmlinuxHandlerRegistry",

pythonbpf/expr/expr_pass.py

@@ -6,11 +6,11 @@ from typing import Dict
 
 from pythonbpf.type_deducer import ctypes_to_ir, is_ctypes
 from .call_registry import CallHandlerRegistry
+from .ir_ops import deref_to_depth, access_struct_field
 from .type_normalization import (
     convert_to_bool,
     handle_comparator,
     get_base_type_and_depth,
-    deref_to_depth,
 )
 from .vmlinux_registry import VmlinuxHandlerRegistry
 from ..vmlinux_parser.dependency_node import Field
@@ -77,89 +77,6 @@ def _handle_attribute_expr(
             logger.info(
                 f"Variable type: {var_type}, Variable ptr: {var_ptr}, Variable Metadata: {var_metadata}"
             )
-            # Check if this is a pointer to a struct (from map lookup)
-            if (
-                isinstance(var_type, ir.PointerType)
-                and var_metadata
-                and isinstance(var_metadata, str)
-            ):
-                if var_metadata in structs_sym_tab:
-                    logger.info(
-                        f"Handling pointer to struct {var_metadata} from map lookup"
-                    )
-
-                    if func is None:
-                        raise ValueError(
-                            f"func parameter required for null-safe pointer access to {var_name}.{attr_name}"
-                        )
-
-                    # Load the pointer value (ptr<struct>)
-                    struct_ptr = builder.load(var_ptr)
-
-                    # Create blocks for null check
-                    null_check_block = builder.block
-                    not_null_block = func.append_basic_block(
-                        name=f"{var_name}_not_null"
-                    )
-                    merge_block = func.append_basic_block(name=f"{var_name}_merge")
-
-                    # Check if pointer is null
-                    null_ptr = ir.Constant(struct_ptr.type, None)
-                    is_not_null = builder.icmp_signed("!=", struct_ptr, null_ptr)
-                    logger.info(f"Inserted null check for pointer {var_name}")
-
-                    builder.cbranch(is_not_null, not_null_block, merge_block)
-
-                    # Not-null block: Access the field
-                    builder.position_at_end(not_null_block)
-
-                    # Get struct metadata
-                    metadata = structs_sym_tab[var_metadata]
-                    struct_ptr = builder.bitcast(
-                        struct_ptr, metadata.ir_type.as_pointer()
-                    )
-
-                    if attr_name not in metadata.fields:
-                        raise ValueError(
-                            f"Field '{attr_name}' not found in struct '{var_metadata}'"
-                        )
-
-                    # GEP to field
-                    field_gep = metadata.gep(builder, struct_ptr, attr_name)
-
-                    # Load field value
-                    field_val = builder.load(field_gep)
-                    field_type = metadata.field_type(attr_name)
-
-                    logger.info(
-                        f"Loaded field {attr_name} from struct pointer, type: {field_type}"
-                    )
-
-                    # Branch to merge
-                    not_null_after_load = builder.block
-                    builder.branch(merge_block)
-
-                    # Merge block: PHI node for the result
-                    builder.position_at_end(merge_block)
-                    phi = builder.phi(field_type, name=f"{var_name}_{attr_name}")
-
-                    # If null, return zero/default value
-                    if isinstance(field_type, ir.IntType):
-                        zero_value = ir.Constant(field_type, 0)
-                    elif isinstance(field_type, ir.PointerType):
-                        zero_value = ir.Constant(field_type, None)
-                    elif isinstance(field_type, ir.ArrayType):
-                        # For arrays, we can't easily create a zero constant
-                        # This case is tricky - for now, just use undef
-                        zero_value = ir.Constant(field_type, ir.Undefined)
-                    else:
-                        zero_value = ir.Constant(field_type, ir.Undefined)
-
-                    phi.add_incoming(zero_value, null_check_block)
-                    phi.add_incoming(field_val, not_null_after_load)
-
-                    logger.info(f"Created PHI node for {var_name}.{attr_name}")
-                    return phi, field_type
             if (
                 hasattr(var_metadata, "__module__")
                 and var_metadata.__module__ == "vmlinux"
@@ -180,13 +97,23 @@ def _handle_attribute_expr(
                 )
                 return None
 
-            # Regular user-defined struct
-            metadata = structs_sym_tab.get(var_metadata)
-            if metadata and attr_name in metadata.fields:
-                gep = metadata.gep(builder, var_ptr, attr_name)
-                val = builder.load(gep)
-                field_type = metadata.field_type(attr_name)
-                return val, field_type
+            if var_metadata in structs_sym_tab:
+                return access_struct_field(
+                    builder,
+                    var_ptr,
+                    var_type,
+                    var_metadata,
+                    expr.attr,
+                    structs_sym_tab,
+                    func,
+                )
+            else:
+                logger.error(f"Struct metadata for '{var_name}' not found")
+        else:
+            logger.error(f"Undefined variable '{var_name}' for attribute access")
+    else:
+        logger.error("Unsupported attribute base expression type")
+
     return None
 
 
@@ -241,7 +168,11 @@ def get_operand_value(
             var_type = var.type
             base_type, depth = get_base_type_and_depth(var_type)
             logger.info(f"var is {var}, base_type is {base_type}, depth is {depth}")
-            val = deref_to_depth(func, builder, var, depth)
+            if depth == 1:
+                val = builder.load(var)
+                return val
+            else:
+                val = deref_to_depth(func, builder, var, depth)
             return val
         else:
             # Check if it's a vmlinux enum/constant
@@ -618,7 +549,7 @@ def _handle_boolean_op(
 
 
 # ============================================================================
-# VMLinux casting
+# Struct casting (including vmlinux struct casting)
 # ============================================================================
 
 
@@ -666,17 +597,85 @@ def _handle_vmlinux_cast(
     # Cast the integer/value to a pointer to the struct
     # If arg_val is an integer type, we need to inttoptr it
     ptr_type = ir.PointerType()
-    # TODO: add a integer check here later
-    if ctypes_to_ir(arg_type.type.__name__):
-        # Cast integer to pointer
-        casted_ptr = builder.inttoptr(arg_val, ptr_type)
+    # TODO: add a field value type check here
+    # print(arg_type)
+    if isinstance(arg_type, Field):
+        if ctypes_to_ir(arg_type.type.__name__):
+            # Cast integer to pointer
+            casted_ptr = builder.inttoptr(arg_val, ptr_type)
+        else:
+            logger.error(f"Unsupported type for vmlinux cast: {arg_type}")
+            return None
     else:
-        logger.error(f"Unsupported type for vmlinux cast: {arg_type}")
-        return None
+        casted_ptr = builder.inttoptr(arg_val, ptr_type)
 
     return casted_ptr, vmlinux_struct_type
 
 
+def _handle_user_defined_struct_cast(
+    func,
+    module,
+    builder,
+    expr,
+    local_sym_tab,
+    map_sym_tab,
+    structs_sym_tab,
+):
+    """Handle user-defined struct cast expressions like iphdr(nh).
+
+    This casts a pointer/integer value to a pointer to the user-defined struct,
+    similar to how vmlinux struct casts work but for user-defined @struct types.
+    """
+    if len(expr.args) != 1:
+        logger.info("User-defined struct cast takes exactly one argument")
+        return None
+
+    # Get the struct name
+    struct_name = expr.func.id
+
+    if struct_name not in structs_sym_tab:
+        logger.error(f"Struct {struct_name} not found in structs_sym_tab")
+        return None
+
+    struct_info = structs_sym_tab[struct_name]
+
+    # Evaluate the argument (e.g.,
+    # an address/pointer value)
+    arg_result = eval_expr(
+        func,
+        module,
+        builder,
+        expr.args[0],
+        local_sym_tab,
+        map_sym_tab,
+        structs_sym_tab,
+    )
+
+    if arg_result is None:
+        logger.info("Failed to evaluate argument to user-defined struct cast")
+        return None
+
+    arg_val, arg_type = arg_result
+
+    # Cast the integer/pointer value to a pointer to the struct type
+    # The struct pointer type is a pointer to the struct's IR type
+    struct_ptr_type = ir.PointerType(struct_info.ir_type)
+
+    # If arg_val is an integer type (like i64), convert to pointer using inttoptr
+    if isinstance(arg_val.type, ir.IntType):
+        casted_ptr = builder.inttoptr(arg_val, struct_ptr_type)
+        logger.info(f"Cast integer to pointer for struct {struct_name}")
+    elif isinstance(arg_val.type, ir.PointerType):
+        # If already a pointer, bitcast to the struct pointer type
+        casted_ptr = builder.bitcast(arg_val, struct_ptr_type)
+        logger.info(f"Bitcast pointer to struct pointer for {struct_name}")
+    else:
+        logger.error(f"Unsupported type for user-defined struct cast: {arg_val.type}")
+        return None
+
+    return casted_ptr, struct_name
+
+
 # ============================================================================
 # Expression Dispatcher
 # ============================================================================
@@ -722,6 +721,16 @@ def eval_expr(
                 map_sym_tab,
                 structs_sym_tab,
             )
+        if isinstance(expr.func, ast.Name) and (expr.func.id in structs_sym_tab):
+            return _handle_user_defined_struct_cast(
+                func,
+                module,
+                builder,
+                expr,
+                local_sym_tab,
+                map_sym_tab,
+                structs_sym_tab,
+            )
 
         result = CallHandlerRegistry.handle_call(
             expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab

pythonbpf/expr/ir_ops.py

@@ -17,41 +17,100 @@ def deref_to_depth(func, builder, val, target_depth):
 
         # dereference with null check
         pointee_type = cur_type.pointee
-        null_check_block = builder.block
-        not_null_block = func.append_basic_block(name=f"deref_not_null_{depth}")
-        merge_block = func.append_basic_block(name=f"deref_merge_{depth}")
 
-        null_ptr = ir.Constant(cur_type, None)
-        is_not_null = builder.icmp_signed("!=", cur_val, null_ptr)
-        logger.debug(f"Inserted null check for pointer at depth {depth}")
+        def load_op(builder, ptr):
+            return builder.load(ptr)
 
-        builder.cbranch(is_not_null, not_null_block, merge_block)
-
-        builder.position_at_end(not_null_block)
-        dereferenced_val = builder.load(cur_val)
-        logger.debug(f"Dereferenced to depth {depth - 1}, type: {pointee_type}")
-        builder.branch(merge_block)
-
-        builder.position_at_end(merge_block)
-        phi = builder.phi(pointee_type, name=f"deref_result_{depth}")
-
-        zero_value = (
-            ir.Constant(pointee_type, 0)
-            if isinstance(pointee_type, ir.IntType)
-            else ir.Constant(pointee_type, None)
+        cur_val = _null_checked_operation(
+            func, builder, cur_val, load_op, pointee_type, f"deref_{depth}"
         )
-        phi.add_incoming(zero_value, null_check_block)
-
-        phi.add_incoming(dereferenced_val, not_null_block)
-
-        # Continue with phi result
-        cur_val = phi
         cur_type = pointee_type
+        logger.debug(f"Dereferenced to depth {depth}, type: {pointee_type}")
     return cur_val
 
 
-def deref_struct_ptr(
-    func, builder, struct_ptr, struct_metadata, field_name, structs_sym_tab
+def _null_checked_operation(func, builder, ptr, operation, result_type, name_prefix):
+    """
+    Generic null-checked operation on a pointer.
+    """
+    curr_block = builder.block
+    not_null_block = func.append_basic_block(name=f"{name_prefix}_not_null")
+    merge_block = func.append_basic_block(name=f"{name_prefix}_merge")
+
+    null_ptr = ir.Constant(ptr.type, None)
+    is_not_null = builder.icmp_signed("!=", ptr, null_ptr)
+    builder.cbranch(is_not_null, not_null_block, merge_block)
+
+    builder.position_at_end(not_null_block)
+    result = operation(builder, ptr)
+    not_null_after = builder.block
+    builder.branch(merge_block)
+
+    builder.position_at_end(merge_block)
+    phi = builder.phi(result_type, name=f"{name_prefix}_result")
+
+    if isinstance(result_type, ir.IntType):
+        null_val = ir.Constant(result_type, 0)
+    elif isinstance(result_type, ir.PointerType):
+        null_val = ir.Constant(result_type, None)
+    else:
+        null_val = ir.Constant(result_type, ir.Undefined)
+
+    phi.add_incoming(null_val, curr_block)
+    phi.add_incoming(result, not_null_after)
+
+    return phi
+
+
+def access_struct_field(
+    builder, var_ptr, var_type, var_metadata, field_name, structs_sym_tab, func=None
 ):
-    """Dereference a pointer to a struct type."""
-    return deref_to_depth(func, builder, struct_ptr, 1)
+    """
+    Access a struct field - automatically returns value or pointer based on field type.
+    """
+    metadata = (
+        structs_sym_tab.get(var_metadata)
+        if isinstance(var_metadata, str)
+        else var_metadata
+    )
+    if not metadata or field_name not in metadata.fields:
+        raise ValueError(f"Field '{field_name}' not found in struct")
+
+    field_type = metadata.field_type(field_name)
+    is_ptr_to_struct = isinstance(var_type, ir.PointerType) and isinstance(
+        var_metadata, str
+    )
+
+    # Get struct pointer
+    struct_ptr = builder.load(var_ptr) if is_ptr_to_struct else var_ptr
+
+    should_load = not isinstance(field_type, ir.ArrayType)
+
+    def field_access_op(builder, ptr):
+        typed_ptr = builder.bitcast(ptr, metadata.ir_type.as_pointer())
+        field_ptr = metadata.gep(builder, typed_ptr, field_name)
+        return builder.load(field_ptr) if should_load else field_ptr
+
+    # Handle null check for pointer-to-struct
+    if is_ptr_to_struct:
+        if func is None:
+            raise ValueError("func required for null-safe struct pointer access")
+
+        if should_load:
+            result_type = field_type
+        else:
+            result_type = field_type.as_pointer()
+
+        result = _null_checked_operation(
+            func,
+            builder,
+            struct_ptr,
+            field_access_op,
+            result_type,
+            f"field_{field_name}",
+        )
+        return result, field_type
+
+    field_ptr = metadata.gep(builder, struct_ptr, field_name)
+    result = builder.load(field_ptr) if should_load else field_ptr
+    return result, field_type

pythonbpf/helper/helper_utils.py

@@ -5,6 +5,7 @@ from llvmlite import ir
 from pythonbpf.expr import (
     get_operand_value,
     eval_expr,
+    access_struct_field,
 )
 
 logger = logging.getLogger(__name__)
@@ -135,7 +136,7 @@ def get_or_create_ptr_from_arg(
             and field_type.element.width == 8
         ):
             ptr, sz = get_char_array_ptr_and_size(
-                arg, builder, local_sym_tab, struct_sym_tab
+                arg, builder, local_sym_tab, struct_sym_tab, func
             )
             if not ptr:
                 raise ValueError("Failed to get char array pointer from struct field")
@@ -266,7 +267,9 @@ def get_buffer_ptr_and_size(buf_arg, builder, local_sym_tab, struct_sym_tab):
         )
 
 
-def get_char_array_ptr_and_size(buf_arg, builder, local_sym_tab, struct_sym_tab):
+def get_char_array_ptr_and_size(
+    buf_arg, builder, local_sym_tab, struct_sym_tab, func=None
+):
     """Get pointer to char array and its size."""
 
     # Struct field: obj.field
@@ -277,11 +280,11 @@ def get_char_array_ptr_and_size(buf_arg, builder, local_sym_tab, struct_sym_tab)
         if not (local_sym_tab and var_name in local_sym_tab):
             raise ValueError(f"Variable '{var_name}' not found")
 
-        struct_type = local_sym_tab[var_name].metadata
-        if not (struct_sym_tab and struct_type in struct_sym_tab):
-            raise ValueError(f"Struct type '{struct_type}' not found")
+        struct_ptr, struct_type, struct_metadata = local_sym_tab[var_name]
+        if not (struct_sym_tab and struct_metadata in struct_sym_tab):
+            raise ValueError(f"Struct type '{struct_metadata}' not found")
 
-        struct_info = struct_sym_tab[struct_type]
+        struct_info = struct_sym_tab[struct_metadata]
         if field_name not in struct_info.fields:
             raise ValueError(f"Field '{field_name}' not found")
 
@@ -292,8 +295,24 @@ def get_char_array_ptr_and_size(buf_arg, builder, local_sym_tab, struct_sym_tab)
             )
             return None, 0
 
-        struct_ptr = local_sym_tab[var_name].var
-        field_ptr = struct_info.gep(builder, struct_ptr, field_name)
+        # Check if char array
+        if not (
+            isinstance(field_type, ir.ArrayType)
+            and isinstance(field_type.element, ir.IntType)
+            and field_type.element.width == 8
+        ):
+            logger.warning("Field is not a char array")
+            return None, 0
+
+        field_ptr, _ = access_struct_field(
+            builder,
+            struct_ptr,
+            struct_type,
+            struct_metadata,
+            field_name,
+            struct_sym_tab,
+            func,
+        )
 
         # GEP to first element: [N x i8]* -> i8*
         buf_ptr = builder.gep(

pythonbpf/helper/printk_formatter.py

@@ -222,7 +222,7 @@ def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_ta
     # Special case: struct field char array needs pointer to first element
     if isinstance(expr, ast.Attribute):
         char_array_ptr, _ = get_char_array_ptr_and_size(
-            expr, builder, local_sym_tab, struct_sym_tab
+            expr, builder, local_sym_tab, struct_sym_tab, func
         )
         if char_array_ptr:
             return char_array_ptr

pythonbpf/type_deducer.py

@@ -18,6 +18,10 @@ mapping = {
     "c_longlong": ir.IntType(64),
     "c_uint": ir.IntType(32),
     "c_int": ir.IntType(32),
+    "c_ushort": ir.IntType(16),
+    "c_short": ir.IntType(16),
+    "c_ubyte": ir.IntType(8),
+    "c_byte": ir.IntType(8),
     # Not so sure about this one
     "str": ir.PointerType(ir.IntType(8)),
 }

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -77,7 +77,7 @@ class VmlinuxHandler:
         return None
 
     def get_vmlinux_enum_value(self, name):
-        """Handle vmlinux enum constants by returning LLVM IR constants"""
+        """Handle vmlinux.enum constants by returning LLVM IR constants"""
         if self.is_vmlinux_enum(name):
             value = self.vmlinux_symtab[name].value
             logger.info(f"The value of vmlinux enum {name} = {value}")
@@ -119,9 +119,18 @@ class VmlinuxHandler:
                 # Load the struct pointer from the local variable
                 struct_ptr = builder.load(var_info.var)
 
+                # Determine the preallocated tmp name that assignment pass should have created
+                tmp_name = f"{struct_var_name}_{field_name}_tmp"
+
                 # Use bpf_probe_read_kernel for non-context struct field access
                 field_value = self.load_struct_field(
-                    builder, struct_ptr, globvar_ir, field_data, struct_name
+                    builder,
+                    struct_ptr,
+                    globvar_ir,
+                    field_data,
+                    struct_name,
+                    local_sym_tab,
+                    tmp_name,
                 )
                 # Return field value and field type
                 return field_value, field_data
@@ -130,7 +139,13 @@ class VmlinuxHandler:
 
     @staticmethod
     def load_struct_field(
-        builder, struct_ptr_int, offset_global, field_data, struct_name=None
+        builder,
+        struct_ptr_int,
+        offset_global,
+        field_data,
+        struct_name=None,
+        local_sym_tab=None,
+        tmp_name: str | None = None,
     ):
         """
         Generate LLVM IR to load a field from a regular (non-context) struct using bpf_probe_read_kernel.
@@ -141,6 +156,8 @@ class VmlinuxHandler:
             offset_global: Global variable containing the field offset (i64)
             field_data: contains data about the field
             struct_name: Name of the struct being accessed (optional)
+            local_sym_tab: symbol table (optional) - used to locate preallocated tmp storage
+            tmp_name: name of the preallocated temporary storage to use (preferred)
         Returns:
             The loaded value
         """
@@ -203,9 +220,18 @@ class VmlinuxHandler:
                 else:
                     logger.warning("Complex vmlinux field type, using default 64 bits")
 
-        # Allocate local storage for the field value
-        local_storage = builder.alloca(ir.IntType(int_width))
-        local_storage_i8_ptr = builder.bitcast(local_storage, i8_ptr_type)
+        # Use preallocated temporary storage if provided by allocation pass
+
+        local_storage_i8_ptr = None
+        if tmp_name and local_sym_tab and tmp_name in local_sym_tab:
+            # Expect the tmp to be an alloca created during allocation pass
+            tmp_alloca = local_sym_tab[tmp_name].var
+            local_storage_i8_ptr = builder.bitcast(tmp_alloca, i8_ptr_type)
+        else:
+            # Fallback: allocate inline (not ideal, but preserves behavior)
+            local_storage = builder.alloca(ir.IntType(int_width))
+            local_storage_i8_ptr = builder.bitcast(local_storage, i8_ptr_type)
+            logger.warning(f"Temp storage '{tmp_name}' not found. Allocating inline")
 
         # Use bpf_probe_read_kernel to safely read the field
         # This generates:
@@ -219,7 +245,9 @@ class VmlinuxHandler:
         )
 
         # Load the value from local storage
-        value = builder.load(local_storage)
+        value = builder.load(
+            builder.bitcast(local_storage_i8_ptr, ir.PointerType(ir.IntType(int_width)))
+        )
 
         # Zero-extend i32 to i64 if needed
         if needs_zext:

tests/c-form/disksnoop.bpf.c

@@ -0,0 +1,66 @@
+// disksnoop.bpf.c
+// eBPF program (compile with: clang -O2 -g -target bpf -c disksnoop.bpf.c -o disksnoop.bpf.o)
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_core_read.h>
+
+char LICENSE[] SEC("license") = "GPL";
+
+struct {
+    __uint(type, BPF_MAP_TYPE_HASH);
+    __type(key, __u64);
+    __type(value, __u64);
+    __uint(max_entries, 10240);
+} start_map SEC(".maps");
+
+/* kprobe: record start timestamp keyed by request pointer */
+SEC("kprobe/blk_mq_start_request")
+int trace_start(struct pt_regs *ctx)
+{
+    /* request * is first arg */
+    __u64 reqp = (__u64)(ctx->di);
+    __u64 ts = bpf_ktime_get_ns();
+
+    bpf_map_update_elem(&start_map, &reqp, &ts, BPF_ANY);
+
+//  /* optional debug:
+    bpf_printk("start: req=%llu ts=%llu\n", reqp, ts);
+//  */
+    return 0;
+}
+
+/* completion: compute latency and print data_len, cmd_flags, latency_us */
+SEC("kprobe/blk_mq_end_request")
+int trace_completion(struct pt_regs *ctx)
+{
+    __u64 reqp = (__u64)(ctx->di);
+    __u64 *tsp;
+    __u64 now_ns;
+    __u64 delta_ns;
+    __u64 delta_us = 0;
+    bpf_printk("%lld", reqp);
+    tsp = bpf_map_lookup_elem(&start_map, &reqp);
+    if (!tsp)
+        return 0;
+
+    now_ns = bpf_ktime_get_ns();
+    delta_ns = now_ns - *tsp;
+    delta_us = delta_ns / 1000;
+
+    /* read request fields using CO-RE; needs vmlinux.h/BTF */
+    __u32 data_len = 0;
+    __u32 cmd_flags = 0;
+
+    /* __data_len is usually a 32/64-bit; use CORE read to be safe */
+    data_len = ( __u32 ) BPF_CORE_READ((struct request *)reqp, __data_len);
+    cmd_flags = ( __u32 ) BPF_CORE_READ((struct request *)reqp, cmd_flags);
+
+    /* print: "<bytes> <flags_hex> <latency_us>" */
+    bpf_printk("%u %x %llu\n", data_len, cmd_flags, delta_us);
+
+    /* remove from map */
+    bpf_map_delete_elem(&start_map, &reqp);
+
+    return 0;
+}

tests/c-form/xdp_test.bpf.c

@@ -0,0 +1,36 @@
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <linux/bpf.h>
+#include <linux/if_ether.h>
+#include <linux/ip.h>
+
+struct fake_iphdr {
+  unsigned short useless;
+  unsigned short tot_len;
+  unsigned short id;
+  unsigned short frag_off;
+  unsigned char ttl;
+  unsigned char protocol;
+  unsigned short check;
+  unsigned int saddr;
+  unsigned int daddr;
+};
+
+SEC("xdp")
+int xdp_prog(struct xdp_md *ctx) {
+  unsigned long data = ctx->data;
+  unsigned long data_end = ctx->data_end;
+
+  if (data + sizeof(struct ethhdr) + sizeof(struct fake_iphdr) <= data_end) {
+    struct fake_iphdr *iph = (void *)data + sizeof(struct ethhdr);
+
+    bpf_printk("%d", iph->saddr);
+
+    return XDP_PASS;
+  } else {
+    return XDP_ABORTED;
+  }
+  struct task_struct * a = btf_bpf_get_current_task_btf();
+}
+
+char _license[] SEC("license") = "GPL";

tests/failing_tests/xdp/xdp_test_1.py

@@ -0,0 +1,46 @@
+from vmlinux import XDP_PASS, XDP_ABORTED
+from vmlinux import (
+    struct_xdp_md,
+)
+from pythonbpf import bpf, section, bpfglobal, compile, compile_to_ir, struct
+from ctypes import c_int64, c_ubyte, c_ushort, c_uint32, c_void_p
+
+
+@bpf
+@struct
+class iphdr:
+    useless: c_ushort
+    tot_len: c_ushort
+    id: c_ushort
+    frag_off: c_ushort
+    ttl: c_ubyte
+    protocol: c_ubyte
+    check: c_ushort
+    saddr: c_uint32
+    daddr: c_uint32
+
+
+@bpf
+@section("xdp")
+def ip_detector(ctx: struct_xdp_md) -> c_int64:
+    data = c_void_p(ctx.data)
+    data_end = c_void_p(ctx.data_end)
+    if data + 34 < data_end:
+        hdr = data + 14
+        iph = iphdr(hdr)
+        addr = iph.saddr
+        print(f"ipaddress: {addr}")
+    else:
+        return c_int64(XDP_ABORTED)
+
+    return c_int64(XDP_PASS)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile_to_ir("xdp_test_1.py", "xdp_test_1.ll")
+compile()

tests/passing_tests/assign/ptr_to_char_array.py

@@ -1,4 +1,4 @@
-from pythonbpf import bpf, struct, section, bpfglobal
+from pythonbpf import bpf, struct, section, bpfglobal, compile
 from pythonbpf.helper import comm
 
 from ctypes import c_void_p, c_int64
@@ -26,3 +26,6 @@ def hello(ctx: c_void_p) -> c_int64:
 @bpfglobal
 def LICENSE() -> str:
     return "GPL"
+
+
+compile()

tests/passing_tests/hash_map_struct.py

@@ -1,6 +1,6 @@
 from pythonbpf import bpf, section, struct, bpfglobal, compile, map
 from pythonbpf.maps import HashMap
-from pythonbpf.helper import pid
+from pythonbpf.helper import pid, comm
 from ctypes import c_void_p, c_int64
 
 
@@ -9,6 +9,7 @@ from ctypes import c_void_p, c_int64
 class val_type:
     counter: c_int64
     shizzle: c_int64
+    comm: str(16)
 
 
 @bpf
@@ -22,6 +23,7 @@ def last() -> HashMap:
 def hello_world(ctx: c_void_p) -> c_int64:
     obj = val_type()
     obj.counter, obj.shizzle = 42, 96
+    comm(obj.comm)
     t = last.lookup(obj)
     if t:
         print(f"Found existing entry: counter={obj.counter}, pid={t}")

tests/passing_tests/struct_pylib.py

@@ -0,0 +1,93 @@
+"""
+Test struct values in HashMap.
+
+This example stores a struct in a HashMap and reads it back,
+testing the new set_value_struct() functionality in pylibbpf.
+"""
+
+from pythonbpf import bpf, map, struct, section, bpfglobal, BPF
+from pythonbpf.helper import ktime, smp_processor_id, pid, comm
+from pythonbpf.maps import HashMap
+from ctypes import c_void_p, c_int64, c_uint32, c_uint64
+import time
+import os
+
+
+@bpf
+@struct
+class task_info:
+    pid: c_uint64
+    timestamp: c_uint64
+    comm: str(16)
+
+
+@bpf
+@map
+def cpu_tasks() -> HashMap:
+    return HashMap(key=c_uint32, value=task_info, max_entries=256)
+
+
+@bpf
+@section("tracepoint/sched/sched_switch")
+def trace_sched_switch(ctx: c_void_p) -> c_int64:
+    cpu = smp_processor_id()
+
+    # Create task info struct
+    info = task_info()
+    info.pid = pid()
+    info.timestamp = ktime()
+    comm(info.comm)
+
+    # Store in map
+    cpu_tasks.update(cpu, info)
+
+    return 0  # type: ignore
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+# Compile and load
+b = BPF()
+b.load()
+b.attach_all()
+
+print("Testing HashMap with Struct Values")
+
+cpu_map = b["cpu_tasks"]
+cpu_map.set_value_struct("task_info")  # Enable struct deserialization
+
+print("Listening for context switches.. .\n")
+
+num_cpus = os.cpu_count() or 16
+
+try:
+    while True:
+        time.sleep(1)
+
+        print(f"--- Snapshot at {time.strftime('%H:%M:%S')} ---")
+
+        for cpu in range(num_cpus):
+            try:
+                info = cpu_map.lookup(cpu)
+
+                if info:
+                    comm_str = (
+                        bytes(info.comm).decode("utf-8", errors="ignore").rstrip("\x00")
+                    )
+                    ts_sec = info.timestamp / 1e9
+
+                    print(
+                        f"  CPU {cpu}: PID={info.pid}, comm={comm_str}, ts={ts_sec:.3f}s"
+                    )
+            except KeyError:
+                # No data for this CPU yet
+                pass
+
+        print()
+
+except KeyboardInterrupt:
+    print("\nStopped")