Fix all issues with vmlinux struct usage in syntax

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
+
+            # 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
+            # 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