Initial plan

Add support for userspace+kernelspace stack trace example using blazesym

BCC-Examples/sync_perf_output.py

@@ -68,8 +68,6 @@ def callback(cpu, event):
 
 perf = b["events"].open_perf_buffer(callback, struct_name="data_t")
 print("Starting to poll... (Ctrl+C to stop)")
-print("Try running: fork() or clone() system calls to trigger events")
-
 try:
     while True:
         b["events"].poll(1000)

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 ]
 readme = "README.md"
 license = {text = "Apache-2.0"}
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 
 dependencies = [
   "llvmlite",

pythonbpf/allocation_pass.py

@@ -7,6 +7,7 @@ from pythonbpf.helper import HelperHandlerRegistry
 from pythonbpf.vmlinux_parser.dependency_node import Field
 from .expr import VmlinuxHandlerRegistry
 from pythonbpf.type_deducer import ctypes_to_ir
+from pythonbpf.maps import BPFMapType
 
 logger = logging.getLogger(__name__)
 
@@ -25,7 +26,9 @@ def create_targets_and_rvals(stmt):
     return stmt.targets, [stmt.value]
 
 
-def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
+def handle_assign_allocation(
+    builder, stmt, local_sym_tab, map_sym_tab, structs_sym_tab
+):
     """Handle memory allocation for assignment statements."""
 
     logger.info(f"Handling assignment for allocation: {ast.dump(stmt)}")
@@ -55,7 +58,9 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
 
         # Determine type and allocate based on rval
         if isinstance(rval, ast.Call):
-            _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab)
+            _allocate_for_call(
+                builder, var_name, rval, local_sym_tab, map_sym_tab, structs_sym_tab
+            )
         elif isinstance(rval, ast.Constant):
             _allocate_for_constant(builder, var_name, rval, local_sym_tab)
         elif isinstance(rval, ast.BinOp):
@@ -74,7 +79,9 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
             )
 
 
-def _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab):
+def _allocate_for_call(
+    builder, var_name, rval, local_sym_tab, map_sym_tab, structs_sym_tab
+):
     """Allocate memory for variable assigned from a call."""
 
     if isinstance(rval.func, ast.Name):
@@ -116,15 +123,74 @@ def _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab):
 
     elif isinstance(rval.func, ast.Attribute):
         # Map method calls - need double allocation for ptr handling
-        _allocate_for_map_method(builder, var_name, local_sym_tab)
+        _allocate_for_map_method(
+            builder, var_name, rval, local_sym_tab, map_sym_tab, structs_sym_tab
+        )
 
     else:
         logger.warning(f"Unsupported call function type for {var_name}")
 
 
-def _allocate_for_map_method(builder, var_name, local_sym_tab):
+def _allocate_for_map_method(
+    builder, var_name, rval, local_sym_tab, map_sym_tab, structs_sym_tab
+):
     """Allocate memory for variable assigned from map method (double alloc)."""
 
+    map_name = rval.func.value.id
+    method_name = rval.func.attr
+
+    # NOTE: We will have to special case HashMap.lookup which returns a pointer to value type
+    # The value type can be a struct as well, so we need to handle that properly
+    # This special casing is not ideal, as over time other map methods may need similar handling
+    # But for now, we will just handle lookup specifically
+    if map_name not in map_sym_tab:
+        logger.error(f"Map '{map_name}' not found for allocation")
+        return
+
+    if method_name != "lookup":
+        # Fallback allocation for other map methods
+        _allocate_for_map_method_fallback(builder, var_name, local_sym_tab)
+        return
+
+    map_params = map_sym_tab[map_name].params
+    if map_params["type"] != BPFMapType.HASH:
+        logger.warning(
+            "Map method lookup used on non-hash map, using fallback allocation"
+        )
+        _allocate_for_map_method_fallback(builder, var_name, local_sym_tab)
+        return
+
+    value_type = map_params["value"]
+    # Determine IR type for value
+    if isinstance(value_type, str) and value_type in structs_sym_tab:
+        struct_info = structs_sym_tab[value_type]
+        value_ir_type = struct_info.ir_type
+    else:
+        value_ir_type = ctypes_to_ir(value_type)
+
+    if value_ir_type is None:
+        logger.warning(
+            f"Could not determine IR type for map value '{value_type}', using fallback allocation"
+        )
+        _allocate_for_map_method_fallback(builder, var_name, local_sym_tab)
+        return
+
+    # Main variable (pointer to pointer)
+    ir_type = ir.PointerType(ir.IntType(64))
+    var = builder.alloca(ir_type, name=var_name)
+    local_sym_tab[var_name] = LocalSymbol(var, ir_type)
+    # Temporary variable for computed values
+    tmp_ir_type = value_ir_type
+    var_tmp = builder.alloca(tmp_ir_type, name=f"{var_name}_tmp")
+    local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, tmp_ir_type)
+    logger.info(
+        f"Pre-allocated {var_name} and {var_name}_tmp for map method lookup of type {value_ir_type}"
+    )
+
+
+def _allocate_for_map_method_fallback(builder, var_name, local_sym_tab):
+    """Fallback allocation for map method variable (i64* and i64**)."""
+
     # Main variable (pointer to pointer)
     ir_type = ir.PointerType(ir.IntType(64))
     var = builder.alloca(ir_type, name=var_name)
@@ -135,7 +201,9 @@ def _allocate_for_map_method(builder, var_name, local_sym_tab):
     var_tmp = builder.alloca(tmp_ir_type, name=f"{var_name}_tmp")
     local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, tmp_ir_type)
 
-    logger.info(f"Pre-allocated {var_name} and {var_name}_tmp for map method")
+    logger.info(
+        f"Pre-allocated {var_name} and {var_name}_tmp for map method (fallback)"
+    )
 
 
 def _allocate_for_constant(builder, var_name, rval, local_sym_tab):
@@ -177,17 +245,33 @@ def _allocate_for_binop(builder, var_name, local_sym_tab):
     logger.info(f"Pre-allocated {var_name} for binop result")
 
 
+def _get_type_name(ir_type):
+    """Get a string representation of an IR type."""
+    if isinstance(ir_type, ir.IntType):
+        return f"i{ir_type.width}"
+    elif isinstance(ir_type, ir.PointerType):
+        return "ptr"
+    elif isinstance(ir_type, ir.ArrayType):
+        return f"[{ir_type.count}x{_get_type_name(ir_type.element)}]"
+    else:
+        return str(ir_type).replace(" ", "")
+
+
 def allocate_temp_pool(builder, max_temps, local_sym_tab):
     """Allocate the temporary scratch space pool for helper arguments."""
-    if max_temps == 0:
+    if not max_temps:
+        logger.info("No temp pool allocation needed")
         return
 
-    logger.info(f"Allocating temp pool of {max_temps} variables")
-    for i in range(max_temps):
-        temp_name = f"__helper_temp_{i}"
-        temp_var = builder.alloca(ir.IntType(64), name=temp_name)
-        temp_var.align = 8
-        local_sym_tab[temp_name] = LocalSymbol(temp_var, ir.IntType(64))
+    for tmp_type, cnt in max_temps.items():
+        type_name = _get_type_name(tmp_type)
+        logger.info(f"Allocating temp pool of {cnt} variables of type {type_name}")
+        for i in range(cnt):
+            temp_name = f"__helper_temp_{type_name}_{i}"
+            temp_var = builder.alloca(tmp_type, name=temp_name)
+            temp_var.align = _get_alignment(tmp_type)
+            local_sym_tab[temp_name] = LocalSymbol(temp_var, tmp_type)
+            logger.debug(f"Allocated temp variable: {temp_name}")
 
 
 def _allocate_for_name(builder, var_name, rval, local_sym_tab):

pythonbpf/codegen.py

@@ -86,7 +86,7 @@ def processor(source_code, filename, module):
     license_processing(tree, module)
     globals_processing(tree, module)
     structs_sym_tab = structs_proc(tree, module, bpf_chunks)
-    map_sym_tab = maps_proc(tree, module, bpf_chunks)
+    map_sym_tab = maps_proc(tree, module, bpf_chunks, structs_sym_tab)
     func_proc(tree, module, bpf_chunks, map_sym_tab, structs_sym_tab)
 
     globals_list_creation(tree, module)
@@ -218,13 +218,11 @@ def compile(loglevel=logging.WARNING) -> bool:
 def BPF(loglevel=logging.WARNING) -> BpfObject:
     caller_frame = inspect.stack()[1]
     src = inspect.getsource(caller_frame.frame)
-    with tempfile.NamedTemporaryFile(
-        mode="w+", delete=True, suffix=".py"
-    ) as f, tempfile.NamedTemporaryFile(
-        mode="w+", delete=True, suffix=".ll"
-    ) as inter, tempfile.NamedTemporaryFile(
-        mode="w+", delete=False, suffix=".o"
-    ) as obj_file:
+    with (
+        tempfile.NamedTemporaryFile(mode="w+", delete=True, suffix=".py") as f,
+        tempfile.NamedTemporaryFile(mode="w+", delete=True, suffix=".ll") as inter,
+        tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".o") as obj_file,
+    ):
         f.write(src)
         f.flush()
         source = f.name

pythonbpf/debuginfo/debug_info_generator.py

@@ -49,6 +49,10 @@ class DebugInfoGenerator:
             )
         return self._type_cache[key]
 
+    def get_uint8_type(self) -> Any:
+        """Get debug info for signed 8-bit integer"""
+        return self.get_basic_type("char", 8, dc.DW_ATE_unsigned)
+
     def get_int32_type(self) -> Any:
         """Get debug info for signed 32-bit integer"""
         return self.get_basic_type("int", 32, dc.DW_ATE_signed)

pythonbpf/functions/functions_pass.py

@@ -39,7 +39,7 @@ logger = logging.getLogger(__name__)
 def count_temps_in_call(call_node, local_sym_tab):
     """Count the number of temporary variables needed for a function call."""
 
-    count = 0
+    count = {}
     is_helper = False
 
     # NOTE: We exclude print calls for now
@@ -49,21 +49,28 @@ def count_temps_in_call(call_node, local_sym_tab):
             and call_node.func.id != "print"
         ):
             is_helper = True
+            func_name = call_node.func.id
     elif isinstance(call_node.func, ast.Attribute):
         if HelperHandlerRegistry.has_handler(call_node.func.attr):
             is_helper = True
+            func_name = call_node.func.attr
 
     if not is_helper:
-        return 0
+        return {}  # No temps needed
 
-    for arg in call_node.args:
+    for arg_idx in range(len(call_node.args)):
         # NOTE: Count all non-name arguments
         # For struct fields, if it is being passed as an argument,
         # The struct object should already exist in the local_sym_tab
-        if not isinstance(arg, ast.Name) and not (
+        arg = call_node.args[arg_idx]
+        if isinstance(arg, ast.Name) or (
             isinstance(arg, ast.Attribute) and arg.value.id in local_sym_tab
         ):
-            count += 1
+            continue
+        param_type = HelperHandlerRegistry.get_param_type(func_name, arg_idx)
+        if isinstance(param_type, ir.PointerType):
+            pointee_type = param_type.pointee
+            count[pointee_type] = count.get(pointee_type, 0) + 1
 
     return count
 
@@ -99,11 +106,15 @@ def handle_if_allocation(
 def allocate_mem(
     module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
 ):
-    max_temps_needed = 0
+    max_temps_needed = {}
+
+    def merge_type_counts(count_dict):
+        nonlocal max_temps_needed
+        for typ, cnt in count_dict.items():
+            max_temps_needed[typ] = max(max_temps_needed.get(typ, 0), cnt)
 
     def update_max_temps_for_stmt(stmt):
         nonlocal max_temps_needed
-        temps_needed = 0
 
         if isinstance(stmt, ast.If):
             for s in stmt.body:
@@ -112,10 +123,13 @@ def allocate_mem(
                 update_max_temps_for_stmt(s)
             return
 
+        stmt_temps = {}
         for node in ast.walk(stmt):
             if isinstance(node, ast.Call):
-                temps_needed += count_temps_in_call(node, local_sym_tab)
-        max_temps_needed = max(max_temps_needed, temps_needed)
+                call_temps = count_temps_in_call(node, local_sym_tab)
+                for typ, cnt in call_temps.items():
+                    stmt_temps[typ] = stmt_temps.get(typ, 0) + cnt
+        merge_type_counts(stmt_temps)
 
     for stmt in body:
         update_max_temps_for_stmt(stmt)
@@ -133,7 +147,9 @@ def allocate_mem(
                 structs_sym_tab,
             )
         elif isinstance(stmt, ast.Assign):
-            handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab)
+            handle_assign_allocation(
+                builder, stmt, local_sym_tab, map_sym_tab, structs_sym_tab
+            )
 
     allocate_temp_pool(builder, max_temps_needed, local_sym_tab)
 

pythonbpf/helper/__init__.py

@@ -1,7 +1,21 @@
 from .helper_registry import HelperHandlerRegistry
 from .helper_utils import reset_scratch_pool
 from .bpf_helper_handler import handle_helper_call, emit_probe_read_kernel_str_call
-from .helpers import ktime, pid, deref, comm, probe_read_str, XDP_DROP, XDP_PASS
+from .helpers import (
+    ktime,
+    pid,
+    deref,
+    comm,
+    probe_read_str,
+    random,
+    probe_read,
+    smp_processor_id,
+    uid,
+    skb_store_bytes,
+    get_stack,
+    XDP_DROP,
+    XDP_PASS,
+)
 
 
 # Register the helper handler with expr module
@@ -65,6 +79,12 @@ __all__ = [
     "deref",
     "comm",
     "probe_read_str",
+    "random",
+    "probe_read",
+    "smp_processor_id",
+    "uid",
+    "skb_store_bytes",
+    "get_stack",
     "XDP_DROP",
     "XDP_PASS",
 ]

pythonbpf/helper/bpf_helper_handler.py

@@ -8,30 +8,43 @@ from .helper_utils import (
     get_flags_val,
     get_data_ptr_and_size,
     get_buffer_ptr_and_size,
-    get_char_array_ptr_and_size,
     get_ptr_from_arg,
+    get_int_value_from_arg,
 )
 from .printk_formatter import simple_string_print, handle_fstring_print
-
-from logging import Logger
+from pythonbpf.maps import BPFMapType
 import logging
 
-logger: Logger = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 
 class BPFHelperID(Enum):
     BPF_MAP_LOOKUP_ELEM = 1
     BPF_MAP_UPDATE_ELEM = 2
     BPF_MAP_DELETE_ELEM = 3
+    BPF_PROBE_READ = 4
     BPF_KTIME_GET_NS = 5
     BPF_PRINTK = 6
+    BPF_GET_PRANDOM_U32 = 7
+    BPF_GET_SMP_PROCESSOR_ID = 8
+    BPF_SKB_STORE_BYTES = 9
     BPF_GET_CURRENT_PID_TGID = 14
+    BPF_GET_CURRENT_UID_GID = 15
     BPF_GET_CURRENT_COMM = 16
     BPF_PERF_EVENT_OUTPUT = 25
+    BPF_GET_STACK = 67
     BPF_PROBE_READ_KERNEL_STR = 115
+    BPF_RINGBUF_OUTPUT = 130
+    BPF_RINGBUF_RESERVE = 131
+    BPF_RINGBUF_SUBMIT = 132
+    BPF_RINGBUF_DISCARD = 133
 
 
-@HelperHandlerRegistry.register("ktime")
+@HelperHandlerRegistry.register(
+    "ktime",
+    param_types=[],
+    return_type=ir.IntType(64),
+)
 def bpf_ktime_get_ns_emitter(
     call,
     map_ptr,
@@ -54,7 +67,11 @@ def bpf_ktime_get_ns_emitter(
     return result, ir.IntType(64)
 
 
-@HelperHandlerRegistry.register("lookup")
+@HelperHandlerRegistry.register(
+    "lookup",
+    param_types=[ir.PointerType(ir.IntType(64))],
+    return_type=ir.PointerType(ir.IntType(64)),
+)
 def bpf_map_lookup_elem_emitter(
     call,
     map_ptr,
@@ -96,6 +113,7 @@ def bpf_map_lookup_elem_emitter(
     return result, ir.PointerType()
 
 
+# NOTE: This has special handling so we won't reflect the signature here.
 @HelperHandlerRegistry.register("print")
 def bpf_printk_emitter(
     call,
@@ -144,7 +162,15 @@ def bpf_printk_emitter(
     return True
 
 
-@HelperHandlerRegistry.register("update")
+@HelperHandlerRegistry.register(
+    "update",
+    param_types=[
+        ir.PointerType(ir.IntType(64)),
+        ir.PointerType(ir.IntType(64)),
+        ir.IntType(64),
+    ],
+    return_type=ir.PointerType(ir.IntType(64)),
+)
 def bpf_map_update_elem_emitter(
     call,
     map_ptr,
@@ -199,7 +225,11 @@ def bpf_map_update_elem_emitter(
     return result, None
 
 
-@HelperHandlerRegistry.register("delete")
+@HelperHandlerRegistry.register(
+    "delete",
+    param_types=[ir.PointerType(ir.IntType(64))],
+    return_type=ir.PointerType(ir.IntType(64)),
+)
 def bpf_map_delete_elem_emitter(
     call,
     map_ptr,
@@ -239,7 +269,11 @@ def bpf_map_delete_elem_emitter(
     return result, None
 
 
-@HelperHandlerRegistry.register("comm")
+@HelperHandlerRegistry.register(
+    "comm",
+    param_types=[ir.PointerType(ir.IntType(8))],
+    return_type=ir.IntType(64),
+)
 def bpf_get_current_comm_emitter(
     call,
     map_ptr,
@@ -296,7 +330,11 @@ def bpf_get_current_comm_emitter(
     return result, None
 
 
-@HelperHandlerRegistry.register("pid")
+@HelperHandlerRegistry.register(
+    "pid",
+    param_types=[],
+    return_type=ir.IntType(64),
+)
 def bpf_get_current_pid_tgid_emitter(
     call,
     map_ptr,
@@ -318,12 +356,12 @@ def bpf_get_current_pid_tgid_emitter(
     result = builder.call(fn_ptr, [], tail=False)
 
     # Extract the lower 32 bits (PID) using bitwise AND with 0xFFFFFFFF
+    # TODO: return both PID and TGID if we end up needing TGID somewhere
     mask = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
     pid = builder.and_(result, mask)
     return pid, ir.IntType(64)
 
 
-@HelperHandlerRegistry.register("output")
 def bpf_perf_event_output_handler(
     call,
     map_ptr,
@@ -334,6 +372,10 @@ def bpf_perf_event_output_handler(
     struct_sym_tab=None,
     map_sym_tab=None,
 ):
+    """
+    Emit LLVM IR for bpf_perf_event_output helper function call.
+    """
+
     if len(call.args) != 1:
         raise ValueError(
             f"Perf event output expects exactly one argument, got {len(call.args)}"
@@ -371,6 +413,98 @@ def bpf_perf_event_output_handler(
     return result, None
 
 
+def bpf_ringbuf_output_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_ringbuf_output helper function call.
+    """
+
+    if len(call.args) != 1:
+        raise ValueError(
+            f"Ringbuf output expects exactly one argument, got {len(call.args)}"
+        )
+    data_arg = call.args[0]
+    data_ptr, size_val = get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab)
+    flags_val = ir.Constant(ir.IntType(64), 0)
+
+    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
+    data_void_ptr = builder.bitcast(data_ptr, ir.PointerType())
+    fn_type = ir.FunctionType(
+        ir.IntType(64),
+        [
+            ir.PointerType(),
+            ir.PointerType(),
+            ir.IntType(64),
+            ir.IntType(64),
+        ],
+        var_arg=False,
+    )
+    fn_ptr_type = ir.PointerType(fn_type)
+
+    # helper id
+    fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_RINGBUF_OUTPUT.value)
+    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+
+    result = builder.call(
+        fn_ptr, [map_void_ptr, data_void_ptr, size_val, flags_val], tail=False
+    )
+    return result, None
+
+
+@HelperHandlerRegistry.register(
+    "output",
+    param_types=[ir.PointerType(ir.IntType(8))],
+    return_type=ir.IntType(64),
+)
+def handle_output_helper(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Route output helper to the appropriate emitter based on map type.
+    """
+    match map_sym_tab[map_ptr.name].type:
+        case BPFMapType.PERF_EVENT_ARRAY:
+            return bpf_perf_event_output_handler(
+                call,
+                map_ptr,
+                module,
+                builder,
+                func,
+                local_sym_tab,
+                struct_sym_tab,
+                map_sym_tab,
+            )
+        case BPFMapType.RINGBUF:
+            return bpf_ringbuf_output_emitter(
+                call,
+                map_ptr,
+                module,
+                builder,
+                func,
+                local_sym_tab,
+                struct_sym_tab,
+                map_sym_tab,
+            )
+        case _:
+            logger.error("Unsupported map type for output helper.")
+    raise NotImplementedError("Output helper for this map type is not implemented.")
+
+
 def emit_probe_read_kernel_str_call(builder, dst_ptr, dst_size, src_ptr):
     """Emit LLVM IR call to bpf_probe_read_kernel_str"""
 
@@ -398,7 +532,14 @@ def emit_probe_read_kernel_str_call(builder, dst_ptr, dst_size, src_ptr):
     return result
 
 
-@HelperHandlerRegistry.register("probe_read_str")
+@HelperHandlerRegistry.register(
+    "probe_read_str",
+    param_types=[
+        ir.PointerType(ir.IntType(8)),
+        ir.PointerType(ir.IntType(8)),
+    ],
+    return_type=ir.IntType(64),
+)
 def bpf_probe_read_kernel_str_emitter(
     call,
     map_ptr,
@@ -417,8 +558,8 @@ def bpf_probe_read_kernel_str_emitter(
         )
 
     # Get destination buffer (char array -> i8*)
-    dst_ptr, dst_size = get_char_array_ptr_and_size(
-        call.args[0], builder, local_sym_tab, struct_sym_tab
+    dst_ptr, dst_size = get_or_create_ptr_from_arg(
+        func, module, call.args[0], builder, local_sym_tab, map_sym_tab, struct_sym_tab
     )
 
     # Get source pointer (evaluate expression)
@@ -433,6 +574,430 @@ def bpf_probe_read_kernel_str_emitter(
     return result, ir.IntType(64)
 
 
+@HelperHandlerRegistry.register(
+    "random",
+    param_types=[],
+    return_type=ir.IntType(32),
+)
+def bpf_get_prandom_u32_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_get_prandom_u32 helper function call.
+    """
+    helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_PRANDOM_U32.value)
+    fn_type = ir.FunctionType(ir.IntType(32), [], var_arg=False)
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
+    result = builder.call(fn_ptr, [], tail=False)
+    return result, ir.IntType(32)
+
+
+@HelperHandlerRegistry.register(
+    "probe_read",
+    param_types=[
+        ir.PointerType(ir.IntType(8)),
+        ir.IntType(32),
+        ir.PointerType(ir.IntType(8)),
+    ],
+    return_type=ir.IntType(64),
+)
+def bpf_probe_read_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_probe_read helper function
+    """
+
+    if len(call.args) != 3:
+        logger.warn("Expected 3 args for probe_read helper")
+        return
+    dst_ptr = get_or_create_ptr_from_arg(
+        func,
+        module,
+        call.args[0],
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+        ir.IntType(8),
+    )
+    size_val = get_int_value_from_arg(
+        call.args[1],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+    )
+    src_ptr = get_or_create_ptr_from_arg(
+        func,
+        module,
+        call.args[2],
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+        ir.IntType(8),
+    )
+    fn_type = ir.FunctionType(
+        ir.IntType(64),
+        [ir.PointerType(), ir.IntType(32), ir.PointerType()],
+        var_arg=False,
+    )
+    fn_ptr = builder.inttoptr(
+        ir.Constant(ir.IntType(64), BPFHelperID.BPF_PROBE_READ.value),
+        ir.PointerType(fn_type),
+    )
+    result = builder.call(
+        fn_ptr,
+        [
+            builder.bitcast(dst_ptr, ir.PointerType()),
+            builder.trunc(size_val, ir.IntType(32)),
+            builder.bitcast(src_ptr, ir.PointerType()),
+        ],
+        tail=False,
+    )
+    logger.info(f"Emitted bpf_probe_read (size={size_val})")
+    return result, ir.IntType(64)
+
+
+@HelperHandlerRegistry.register(
+    "smp_processor_id",
+    param_types=[],
+    return_type=ir.IntType(32),
+)
+def bpf_get_smp_processor_id_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_get_smp_processor_id helper function call.
+    """
+    helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_SMP_PROCESSOR_ID.value)
+    fn_type = ir.FunctionType(ir.IntType(32), [], var_arg=False)
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
+    result = builder.call(fn_ptr, [], tail=False)
+    logger.info("Emitted bpf_get_smp_processor_id call")
+    return result, ir.IntType(32)
+
+
+@HelperHandlerRegistry.register(
+    "uid",
+    param_types=[],
+    return_type=ir.IntType(64),
+)
+def bpf_get_current_uid_gid_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_get_current_uid_gid helper function call.
+    """
+    helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_UID_GID.value)
+    fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
+    result = builder.call(fn_ptr, [], tail=False)
+
+    # Extract the lower 32 bits (UID) using bitwise AND with 0xFFFFFFFF
+    # TODO: return both UID and GID if we end up needing GID somewhere
+    mask = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
+    pid = builder.and_(result, mask)
+    return pid, ir.IntType(64)
+
+
+@HelperHandlerRegistry.register(
+    "skb_store_bytes",
+    param_types=[
+        ir.IntType(32),
+        ir.PointerType(ir.IntType(8)),
+        ir.IntType(32),
+        ir.IntType(64),
+    ],
+    return_type=ir.IntType(64),
+)
+def bpf_skb_store_bytes_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_skb_store_bytes helper function call.
+    Expected call signature: skb_store_bytes(skb, offset, from, len, flags)
+    """
+
+    args_signature = [
+        ir.PointerType(),  # skb pointer
+        ir.IntType(32),  # offset
+        ir.PointerType(),  # from
+        ir.IntType(32),  # len
+        ir.IntType(64),  # flags
+    ]
+
+    if len(call.args) not in (3, 4):
+        raise ValueError(
+            f"skb_store_bytes expects 3 or 4 args (offset, from, len, flags), got {len(call.args)}"
+        )
+
+    skb_ptr = func.args[0]  # First argument to the function is skb
+    offset_val = get_int_value_from_arg(
+        call.args[0],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+    )
+    from_ptr = get_or_create_ptr_from_arg(
+        func,
+        module,
+        call.args[1],
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+        args_signature[2],
+    )
+    len_val = get_int_value_from_arg(
+        call.args[2],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+    )
+    if len(call.args) == 4:
+        flags_val = get_flags_val(call.args[3], builder, local_sym_tab)
+    else:
+        flags_val = 0
+    if isinstance(flags_val, int):
+        flags = ir.Constant(ir.IntType(64), flags_val)
+    else:
+        flags = flags_val
+    fn_type = ir.FunctionType(
+        ir.IntType(64),
+        args_signature,
+        var_arg=False,
+    )
+    fn_ptr = builder.inttoptr(
+        ir.Constant(ir.IntType(64), BPFHelperID.BPF_SKB_STORE_BYTES.value),
+        ir.PointerType(fn_type),
+    )
+    result = builder.call(
+        fn_ptr,
+        [
+            builder.bitcast(skb_ptr, ir.PointerType()),
+            builder.trunc(offset_val, ir.IntType(32)),
+            builder.bitcast(from_ptr, ir.PointerType()),
+            builder.trunc(len_val, ir.IntType(32)),
+            flags,
+        ],
+        tail=False,
+    )
+    logger.info("Emitted bpf_skb_store_bytes call")
+    return result, ir.IntType(64)
+
+
+@HelperHandlerRegistry.register(
+    "reserve",
+    param_types=[ir.IntType(64)],
+    return_type=ir.PointerType(ir.IntType(8)),
+)
+def bpf_ringbuf_reserve_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_ringbuf_reserve helper function call.
+    Expected call signature: ringbuf.reserve(size)
+    """
+
+    if len(call.args) != 1:
+        raise ValueError(
+            f"ringbuf.reserve expects exactly one argument (size), got {len(call.args)}"
+        )
+
+    size_val = get_int_value_from_arg(
+        call.args[0],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+    )
+
+    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
+    fn_type = ir.FunctionType(
+        ir.PointerType(ir.IntType(8)),
+        [ir.PointerType(), ir.IntType(64)],
+        var_arg=False,
+    )
+    fn_ptr_type = ir.PointerType(fn_type)
+
+    fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_RINGBUF_RESERVE.value)
+    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+
+    result = builder.call(fn_ptr, [map_void_ptr, size_val], tail=False)
+
+    return result, ir.PointerType(ir.IntType(8))
+
+
+@HelperHandlerRegistry.register(
+    "submit",
+    param_types=[ir.PointerType(ir.IntType(8)), ir.IntType(64)],
+    return_type=ir.VoidType(),
+)
+def bpf_ringbuf_submit_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_ringbuf_submit helper function call.
+    Expected call signature: ringbuf.submit(data, flags=0)
+    """
+
+    if len(call.args) not in (1, 2):
+        raise ValueError(
+            f"ringbuf.submit expects 1 or 2 args (data, flags), got {len(call.args)}"
+        )
+
+    data_arg = call.args[0]
+    flags_arg = call.args[1] if len(call.args) == 2 else None
+
+    data_ptr = get_or_create_ptr_from_arg(
+        func,
+        module,
+        data_arg,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+        ir.PointerType(ir.IntType(8)),
+    )
+
+    flags_const = get_flags_val(flags_arg, builder, local_sym_tab)
+    if isinstance(flags_const, int):
+        flags_const = ir.Constant(ir.IntType(64), flags_const)
+
+    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
+    fn_type = ir.FunctionType(
+        ir.VoidType(),
+        [ir.PointerType(), ir.PointerType(), ir.IntType(64)],
+        var_arg=False,
+    )
+    fn_ptr_type = ir.PointerType(fn_type)
+
+    fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_RINGBUF_SUBMIT.value)
+    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+
+    result = builder.call(fn_ptr, [map_void_ptr, data_ptr, flags_const], tail=False)
+
+    return result, None
+
+
+@HelperHandlerRegistry.register(
+    "get_stack",
+    param_types=[ir.PointerType(ir.IntType(8)), ir.IntType(64)],
+    return_type=ir.IntType(64),
+)
+def bpf_get_stack_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_get_stack helper function call.
+    """
+    if len(call.args) not in (1, 2):
+        raise ValueError(
+            f"get_stack expects atmost two arguments (buf, flags), got {len(call.args)}"
+        )
+    ctx_ptr = func.args[0]  # First argument to the function is ctx
+    buf_arg = call.args[0]
+    flags_arg = call.args[1] if len(call.args) == 2 else None
+    buf_ptr, buf_size = get_buffer_ptr_and_size(
+        buf_arg, builder, local_sym_tab, struct_sym_tab
+    )
+    flags_val = get_flags_val(flags_arg, builder, local_sym_tab)
+    if isinstance(flags_val, int):
+        flags_val = ir.Constant(ir.IntType(64), flags_val)
+
+    buf_void_ptr = builder.bitcast(buf_ptr, ir.PointerType())
+    fn_type = ir.FunctionType(
+        ir.IntType(64),
+        [
+            ir.PointerType(ir.IntType(8)),
+            ir.PointerType(),
+            ir.IntType(64),
+            ir.IntType(64),
+        ],
+        var_arg=False,
+    )
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_STACK.value)
+    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+    result = builder.call(
+        fn_ptr,
+        [ctx_ptr, buf_void_ptr, ir.Constant(ir.IntType(64), buf_size), flags_val],
+        tail=False,
+    )
+    return result, ir.IntType(64)
+
+
 def handle_helper_call(
     call,
     module,
@@ -487,6 +1052,6 @@ def handle_helper_call(
         if not map_sym_tab or map_name not in map_sym_tab:
             raise ValueError(f"Map '{map_name}' not found in symbol table")
 
-        return invoke_helper(method_name, map_sym_tab[map_name])
+        return invoke_helper(method_name, map_sym_tab[map_name].sym)
 
     return None

pythonbpf/helper/helper_registry.py

@@ -1,17 +1,31 @@
+from dataclasses import dataclass
+from llvmlite import ir
 from typing import Callable
 
 
+@dataclass
+class HelperSignature:
+    """Signature of a BPF helper function"""
+
+    arg_types: list[ir.Type]
+    return_type: ir.Type
+    func: Callable
+
+
 class HelperHandlerRegistry:
     """Registry for BPF helpers"""
 
-    _handlers: dict[str, Callable] = {}
+    _handlers: dict[str, HelperSignature] = {}
 
     @classmethod
-    def register(cls, helper_name):
+    def register(cls, helper_name, param_types=None, return_type=None):
         """Decorator to register a handler function for a helper"""
 
         def decorator(func):
-            cls._handlers[helper_name] = func
+            helper_sig = HelperSignature(
+                arg_types=param_types, return_type=return_type, func=func
+            )
+            cls._handlers[helper_name] = helper_sig
             return func
 
         return decorator
@@ -19,9 +33,29 @@ class HelperHandlerRegistry:
     @classmethod
     def get_handler(cls, helper_name):
         """Get the handler function for a helper"""
-        return cls._handlers.get(helper_name)
+        handler = cls._handlers.get(helper_name)
+        return handler.func if handler else None
 
     @classmethod
     def has_handler(cls, helper_name):
         """Check if a handler function is registered for a helper"""
         return helper_name in cls._handlers
+
+    @classmethod
+    def get_signature(cls, helper_name):
+        """Get the signature of a helper function"""
+        return cls._handlers.get(helper_name)
+
+    @classmethod
+    def get_param_type(cls, helper_name, index):
+        """Get the type of a parameter of a helper function by the index"""
+        signature = cls.get_signature(helper_name)
+        if signature and signature.arg_types and 0 <= index < len(signature.arg_types):
+            return signature.arg_types[index]
+        return None
+
+    @classmethod
+    def get_return_type(cls, helper_name):
+        """Get the return type of a helper function"""
+        signature = cls.get_signature(helper_name)
+        return signature.return_type if signature else None

pythonbpf/helper/helper_utils.py

@@ -14,26 +14,43 @@ class ScratchPoolManager:
     """Manage the temporary helper variables in local_sym_tab"""
 
     def __init__(self):
-        self._counter = 0
+        self._counters = {}
 
     @property
     def counter(self):
-        return self._counter
+        return sum(self._counters.values())
 
     def reset(self):
-        self._counter = 0
+        self._counters.clear()
         logger.debug("Scratch pool counter reset to 0")
 
-    def get_next_temp(self, local_sym_tab):
-        temp_name = f"__helper_temp_{self._counter}"
-        self._counter += 1
+    def _get_type_name(self, ir_type):
+        if isinstance(ir_type, ir.PointerType):
+            return "ptr"
+        elif isinstance(ir_type, ir.IntType):
+            return f"i{ir_type.width}"
+        elif isinstance(ir_type, ir.ArrayType):
+            return f"[{ir_type.count}x{self._get_type_name(ir_type.element)}]"
+        else:
+            return str(ir_type).replace(" ", "")
+
+    def get_next_temp(self, local_sym_tab, expected_type=None):
+        # Default to i64 if no expected type provided
+        type_name = self._get_type_name(expected_type) if expected_type else "i64"
+        if type_name not in self._counters:
+            self._counters[type_name] = 0
+
+        counter = self._counters[type_name]
+        temp_name = f"__helper_temp_{type_name}_{counter}"
+        self._counters[type_name] += 1
 
         if temp_name not in local_sym_tab:
             raise ValueError(
                 f"Scratch pool exhausted or inadequate: {temp_name}. "
-                f"Current counter: {self._counter}"
+                f"Type: {type_name} Counter: {counter}"
             )
 
+        logger.debug(f"Using {temp_name} for type {type_name}")
         return local_sym_tab[temp_name].var, temp_name
 
 
@@ -60,24 +77,73 @@ def get_var_ptr_from_name(var_name, local_sym_tab):
 def create_int_constant_ptr(value, builder, local_sym_tab, int_width=64):
     """Create a pointer to an integer constant."""
 
-    # Default to 64-bit integer
-    ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab)
+    int_type = ir.IntType(int_width)
+    ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab, int_type)
     logger.info(f"Using temp variable '{temp_name}' for int constant {value}")
-    const_val = ir.Constant(ir.IntType(int_width), value)
+    const_val = ir.Constant(int_type, value)
     builder.store(const_val, ptr)
     return ptr
 
 
 def get_or_create_ptr_from_arg(
-    func, module, arg, builder, local_sym_tab, map_sym_tab, struct_sym_tab=None
+    func,
+    module,
+    arg,
+    builder,
+    local_sym_tab,
+    map_sym_tab,
+    struct_sym_tab=None,
+    expected_type=None,
 ):
     """Extract or create pointer from the call arguments."""
 
+    logger.info(f"Getting pointer from arg: {ast.dump(arg)}")
+    sz = None
     if isinstance(arg, ast.Name):
+        # Stack space is already allocated
         ptr = get_var_ptr_from_name(arg.id, local_sym_tab)
     elif isinstance(arg, ast.Constant) and isinstance(arg.value, int):
-        ptr = create_int_constant_ptr(arg.value, builder, local_sym_tab)
+        int_width = 64  # Default to i64
+        if expected_type and isinstance(expected_type, ir.IntType):
+            int_width = expected_type.width
+        ptr = create_int_constant_ptr(arg.value, builder, local_sym_tab, int_width)
+    elif isinstance(arg, ast.Attribute):
+        # A struct field
+        struct_name = arg.value.id
+        field_name = arg.attr
+
+        if not local_sym_tab or struct_name not in local_sym_tab:
+            raise ValueError(f"Struct '{struct_name}' not found")
+
+        struct_type = local_sym_tab[struct_name].metadata
+        if not struct_sym_tab or struct_type not in struct_sym_tab:
+            raise ValueError(f"Struct type '{struct_type}' not found")
+
+        struct_info = struct_sym_tab[struct_type]
+        if field_name not in struct_info.fields:
+            raise ValueError(
+                f"Field '{field_name}' not found in struct '{struct_name}'"
+            )
+
+        field_type = struct_info.field_type(field_name)
+        struct_ptr = local_sym_tab[struct_name].var
+
+        # Special handling for char arrays
+        if (
+            isinstance(field_type, ir.ArrayType)
+            and isinstance(field_type.element, ir.IntType)
+            and field_type.element.width == 8
+        ):
+            ptr, sz = get_char_array_ptr_and_size(
+                arg, builder, local_sym_tab, struct_sym_tab
+            )
+            if not ptr:
+                raise ValueError("Failed to get char array pointer from struct field")
+        else:
+            ptr = struct_info.gep(builder, struct_ptr, field_name)
+
     else:
+        # NOTE: For any integer expression reaching this branch, it is probably a struct field or a binop
         # Evaluate the expression and store the result in a temp variable
         val = get_operand_value(
             func, module, arg, builder, local_sym_tab, map_sym_tab, struct_sym_tab
@@ -85,13 +151,20 @@ def get_or_create_ptr_from_arg(
         if val is None:
             raise ValueError("Failed to evaluate expression for helper arg.")
 
-        # NOTE: We assume the result is an int64 for now
-        # if isinstance(arg, ast.Attribute):
-        # return val
-        ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab)
+        ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab, expected_type)
         logger.info(f"Using temp variable '{temp_name}' for expression result")
+        if (
+            isinstance(val.type, ir.IntType)
+            and expected_type
+            and val.type.width > expected_type.width
+        ):
+            val = builder.trunc(val, expected_type)
         builder.store(val, ptr)
 
+    # NOTE: For char arrays, also return size
+    if sz:
+        return ptr, sz
+
     return ptr
 
 
@@ -214,7 +287,10 @@ def get_char_array_ptr_and_size(buf_arg, builder, local_sym_tab, struct_sym_tab)
 
         field_type = struct_info.field_type(field_name)
         if not _is_char_array(field_type):
-            raise ValueError("Expected char array field")
+            logger.info(
+                "Field is not a char array, falling back to int or ptr detection"
+            )
+            return None, 0
 
         struct_ptr = local_sym_tab[var_name].var
         field_ptr = struct_info.gep(builder, struct_ptr, field_name)
@@ -274,3 +350,23 @@ def get_ptr_from_arg(
         raise ValueError(f"Expected pointer type, got {val_type}")
 
     return val, val_type
+
+
+def get_int_value_from_arg(
+    arg, func, module, builder, local_sym_tab, map_sym_tab, struct_sym_tab
+):
+    """Evaluate argument and return integer value"""
+
+    result = eval_expr(
+        func, module, builder, arg, local_sym_tab, map_sym_tab, struct_sym_tab
+    )
+
+    if not result:
+        raise ValueError("Failed to evaluate argument")
+
+    val, val_type = result
+
+    if not isinstance(val_type, ir.IntType):
+        raise ValueError(f"Expected integer type, got {val_type}")
+
+    return val

pythonbpf/helper/helpers.py

@@ -27,6 +27,36 @@ def probe_read_str(dst, src):
     return ctypes.c_int64(0)
 
 
+def random():
+    """get a pseudorandom u32 number"""
+    return ctypes.c_int32(0)
+
+
+def probe_read(dst, size, src):
+    """Safely read data from kernel memory"""
+    return ctypes.c_int64(0)
+
+
+def smp_processor_id():
+    """get the current CPU id"""
+    return ctypes.c_int32(0)
+
+
+def uid():
+    """get current user id"""
+    return ctypes.c_int32(0)
+
+
+def skb_store_bytes(offset, from_buf, size, flags=0):
+    """store bytes into a socket buffer"""
+    return ctypes.c_int64(0)
+
+
+def get_stack(buf, flags=0):
+    """get the current stack trace"""
+    return ctypes.c_int64(0)
+
+
 XDP_ABORTED = ctypes.c_int64(0)
 XDP_DROP = ctypes.c_int64(1)
 XDP_PASS = ctypes.c_int64(2)

pythonbpf/helper/printk_formatter.py

@@ -4,6 +4,7 @@ import logging
 from llvmlite import ir
 from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
 from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
+from pythonbpf.helper.helper_utils import get_char_array_ptr_and_size
 
 logger = logging.getLogger(__name__)
 
@@ -219,11 +220,12 @@ def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_ta
     """Evaluate and prepare an expression to use as an arg for bpf_printk."""
 
     # Special case: struct field char array needs pointer to first element
-    char_array_ptr = _get_struct_char_array_ptr(
-        expr, builder, local_sym_tab, struct_sym_tab
-    )
-    if char_array_ptr:
-        return char_array_ptr
+    if isinstance(expr, ast.Attribute):
+        char_array_ptr, _ = get_char_array_ptr_and_size(
+            expr, builder, local_sym_tab, struct_sym_tab
+        )
+        if char_array_ptr:
+            return char_array_ptr
 
     # Regular expression evaluation
     val, _ = eval_expr(func, module, builder, expr, local_sym_tab, None, struct_sym_tab)
@@ -242,52 +244,6 @@ def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_ta
         return ir.Constant(ir.IntType(64), 0)
 
 
-def _get_struct_char_array_ptr(expr, builder, local_sym_tab, struct_sym_tab):
-    """Get pointer to first element of char array in struct field, or None."""
-    if not (isinstance(expr, ast.Attribute) and isinstance(expr.value, ast.Name)):
-        return None
-
-    var_name = expr.value.id
-    field_name = expr.attr
-
-    # Check if it's a valid struct field
-    if not (
-        local_sym_tab
-        and var_name in local_sym_tab
-        and struct_sym_tab
-        and local_sym_tab[var_name].metadata in struct_sym_tab
-    ):
-        return None
-
-    struct_type = local_sym_tab[var_name].metadata
-    struct_info = struct_sym_tab[struct_type]
-
-    if field_name not in struct_info.fields:
-        return None
-
-    field_type = struct_info.field_type(field_name)
-
-    # Check if it's a char array
-    is_char_array = (
-        isinstance(field_type, ir.ArrayType)
-        and isinstance(field_type.element, ir.IntType)
-        and field_type.element.width == 8
-    )
-
-    if not is_char_array:
-        return None
-
-    # Get field pointer and GEP to first element: [N x i8]* -> i8*
-    struct_ptr = local_sym_tab[var_name].var
-    field_ptr = struct_info.gep(builder, struct_ptr, field_name)
-
-    return builder.gep(
-        field_ptr,
-        [ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), 0)],
-        inbounds=True,
-    )
-
-
 def _handle_pointer_arg(val, func, builder):
     """Convert pointer type for bpf_printk."""
     target, depth = get_base_type_and_depth(val.type)

pythonbpf/maps/__init__.py

@@ -1,4 +1,5 @@
-from .maps import HashMap, PerfEventArray, RingBuf
+from .maps import HashMap, PerfEventArray, RingBuffer
 from .maps_pass import maps_proc
+from .map_types import BPFMapType
 
-__all__ = ["HashMap", "PerfEventArray", "maps_proc", "RingBuf"]
+__all__ = ["HashMap", "PerfEventArray", "maps_proc", "RingBuffer", "BPFMapType"]

pythonbpf/maps/map_debug_info.py

@@ -1,22 +1,31 @@
+import logging
+from llvmlite import ir
 from pythonbpf.debuginfo import DebugInfoGenerator
 from .map_types import BPFMapType
 
+logger: logging.Logger = logging.getLogger(__name__)
 
-def create_map_debug_info(module, map_global, map_name, map_params):
+
+def create_map_debug_info(module, map_global, map_name, map_params, structs_sym_tab):
     """Generate debug info metadata for BPF maps HASH and PERF_EVENT_ARRAY"""
     generator = DebugInfoGenerator(module)
-
+    logger.info(f"Creating debug info for map {map_name} with params {map_params}")
     uint_type = generator.get_uint32_type()
-    ulong_type = generator.get_uint64_type()
     array_type = generator.create_array_type(
         uint_type, map_params.get("type", BPFMapType.UNSPEC).value
     )
     type_ptr = generator.create_pointer_type(array_type, 64)
     key_ptr = generator.create_pointer_type(
-        array_type if "key_size" in map_params else ulong_type, 64
+        array_type
+        if "key_size" in map_params
+        else _get_key_val_dbg_type(map_params.get("key"), generator, structs_sym_tab),
+        64,
     )
     value_ptr = generator.create_pointer_type(
-        array_type if "value_size" in map_params else ulong_type, 64
+        array_type
+        if "value_size" in map_params
+        else _get_key_val_dbg_type(map_params.get("value"), generator, structs_sym_tab),
+        64,
     )
 
     elements_arr = []
@@ -64,7 +73,13 @@ def create_map_debug_info(module, map_global, map_name, map_params):
     return global_var
 
 
-def create_ringbuf_debug_info(module, map_global, map_name, map_params):
+# TODO: This should not be exposed outside of the module.
+# Ideally we should expose a single create_map_debug_info function that handles all map types.
+# We can probably use a registry pattern to register different map types and their debug info generators.
+# map_params["type"] will be used to determine which generator to use.
+def create_ringbuf_debug_info(
+    module, map_global, map_name, map_params, structs_sym_tab
+):
     """Generate debug information metadata for BPF RINGBUF map"""
     generator = DebugInfoGenerator(module)
 
@@ -91,3 +106,65 @@ def create_ringbuf_debug_info(module, map_global, map_name, map_params):
     )
     map_global.set_metadata("dbg", global_var)
     return global_var
+
+
+def _get_key_val_dbg_type(name, generator, structs_sym_tab):
+    """Get the debug type for key/value based on type object"""
+
+    if not name:
+        logger.warn("No name provided for key/value type, defaulting to uint64")
+        return generator.get_uint64_type()
+
+    type_obj = structs_sym_tab.get(name)
+    if type_obj:
+        return _get_struct_debug_type(type_obj, generator, structs_sym_tab)
+
+    # Fallback to basic types
+    logger.info(f"No struct named {name}, falling back to basic type")
+
+    # NOTE: Only handling int and long for now
+    if name in ["c_int32", "c_uint32"]:
+        return generator.get_uint32_type()
+
+    # Default fallback for now
+    return generator.get_uint64_type()
+
+
+def _get_struct_debug_type(struct_obj, generator, structs_sym_tab):
+    """Recursively create debug type for struct"""
+    elements_arr = []
+    for fld in struct_obj.fields.keys():
+        fld_type = struct_obj.field_type(fld)
+        if isinstance(fld_type, ir.IntType):
+            if fld_type.width == 32:
+                fld_dbg_type = generator.get_uint32_type()
+            else:
+                # NOTE: Assuming 64-bit for all other int types
+                fld_dbg_type = generator.get_uint64_type()
+        elif isinstance(fld_type, ir.ArrayType):
+            # NOTE: Array types have u8 elements only for now
+            # Debug info generation should fail for other types
+            elem_type = fld_type.element
+            if isinstance(elem_type, ir.IntType) and elem_type.width == 8:
+                char_type = generator.get_uint8_type()
+                fld_dbg_type = generator.create_array_type(char_type, fld_type.count)
+            else:
+                logger.warning(
+                    f"Array element type {str(elem_type)} not supported for debug info, skipping"
+                )
+                continue
+        else:
+            # NOTE: Only handling int and char arrays for now
+            logger.warning(
+                f"Field type {str(fld_type)} not supported for debug info, skipping"
+            )
+            continue
+
+        member = generator.create_struct_member(
+            fld, fld_dbg_type, struct_obj.field_size(fld)
+        )
+        elements_arr.append(member)
+    struct_type = generator.create_struct_type(
+        elements_arr, struct_obj.size, is_distinct=True
+    )
+    return struct_type

pythonbpf/maps/maps.py

@@ -36,11 +36,14 @@ class PerfEventArray:
         pass  # Placeholder for output method
 
 
-class RingBuf:
+class RingBuffer:
     def __init__(self, max_entries):
         self.max_entries = max_entries
 
-    def reserve(self, size: int, flags=0):
+    def output(self, data, flags=0):
+        pass
+
+    def reserve(self, size: int):
         if size > self.max_entries:
             raise ValueError("size cannot be greater than set maximum entries")
         return 0
@@ -48,4 +51,7 @@ class RingBuf:
     def submit(self, data, flags=0):
         pass
 
+    def discard(self, data, flags=0):
+        pass
+
     # add discard, output and also give names to flags and stuff

pythonbpf/maps/maps_pass.py

@@ -3,7 +3,7 @@ import logging
 from logging import Logger
 from llvmlite import ir
 
-from .maps_utils import MapProcessorRegistry
+from .maps_utils import MapProcessorRegistry, MapSymbol
 from .map_types import BPFMapType
 from .map_debug_info import create_map_debug_info, create_ringbuf_debug_info
 from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
@@ -12,13 +12,15 @@ from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
 logger: Logger = logging.getLogger(__name__)
 
 
-def maps_proc(tree, module, chunks):
+def maps_proc(tree, module, chunks, structs_sym_tab):
     """Process all functions decorated with @map to find BPF maps"""
     map_sym_tab = {}
     for func_node in chunks:
         if is_map(func_node):
             logger.info(f"Found BPF map: {func_node.name}")
-            map_sym_tab[func_node.name] = process_bpf_map(func_node, module)
+            map_sym_tab[func_node.name] = process_bpf_map(
+                func_node, module, structs_sym_tab
+            )
     return map_sym_tab
 
 
@@ -46,7 +48,7 @@ def create_bpf_map(module, map_name, map_params):
     map_global.align = 8
 
     logger.info(f"Created BPF map: {map_name} with params {map_params}")
-    return map_global
+    return MapSymbol(type=map_params["type"], sym=map_global, params=map_params)
 
 
 def _parse_map_params(rval, expected_args=None):
@@ -60,7 +62,8 @@ def _parse_map_params(rval, expected_args=None):
             if i < len(rval.args):
                 arg = rval.args[i]
                 if isinstance(arg, ast.Name):
-                    params[arg_name] = arg.id
+                    result = _get_vmlinux_enum(handler, arg.id)
+                    params[arg_name] = result if result is not None else arg.id
                 elif isinstance(arg, ast.Constant):
                     params[arg_name] = arg.value
 
@@ -68,33 +71,48 @@ def _parse_map_params(rval, expected_args=None):
     for keyword in rval.keywords:
         if isinstance(keyword.value, ast.Name):
             name = keyword.value.id
-            if handler and handler.is_vmlinux_enum(name):
-                result = handler.get_vmlinux_enum_value(name)
-                params[keyword.arg] = result if result is not None else name
-            else:
-                params[keyword.arg] = name
+            result = _get_vmlinux_enum(handler, name)
+            params[keyword.arg] = result if result is not None else name
         elif isinstance(keyword.value, ast.Constant):
             params[keyword.arg] = keyword.value.value
 
     return params
 
 
-@MapProcessorRegistry.register("RingBuf")
-def process_ringbuf_map(map_name, rval, module):
+def _get_vmlinux_enum(handler, name):
+    if handler and handler.is_vmlinux_enum(name):
+        return handler.get_vmlinux_enum_value(name)
+
+
+@MapProcessorRegistry.register("RingBuffer")
+def process_ringbuf_map(map_name, rval, module, structs_sym_tab):
     """Process a BPF_RINGBUF map declaration"""
     logger.info(f"Processing Ringbuf: {map_name}")
     map_params = _parse_map_params(rval, expected_args=["max_entries"])
     map_params["type"] = BPFMapType.RINGBUF
 
+    # NOTE: constraints borrowed from https://docs.ebpf.io/linux/map-type/BPF_MAP_TYPE_RINGBUF/
+    max_entries = map_params.get("max_entries")
+    if (
+        not isinstance(max_entries, int)
+        or max_entries < 4096
+        or (max_entries & (max_entries - 1)) != 0
+    ):
+        raise ValueError(
+            "Ringbuf max_entries must be a power of two greater than or equal to the page size (4096)"
+        )
+
     logger.info(f"Ringbuf map parameters: {map_params}")
 
     map_global = create_bpf_map(module, map_name, map_params)
-    create_ringbuf_debug_info(module, map_global, map_name, map_params)
+    create_ringbuf_debug_info(
+        module, map_global.sym, map_name, map_params, structs_sym_tab
+    )
     return map_global
 
 
 @MapProcessorRegistry.register("HashMap")
-def process_hash_map(map_name, rval, module):
+def process_hash_map(map_name, rval, module, structs_sym_tab):
     """Process a BPF_HASH map declaration"""
     logger.info(f"Processing HashMap: {map_name}")
     map_params = _parse_map_params(rval, expected_args=["key", "value", "max_entries"])
@@ -103,12 +121,12 @@ def process_hash_map(map_name, rval, module):
     logger.info(f"Map parameters: {map_params}")
     map_global = create_bpf_map(module, map_name, map_params)
     # Generate debug info for BTF
-    create_map_debug_info(module, map_global, map_name, map_params)
+    create_map_debug_info(module, map_global.sym, map_name, map_params, structs_sym_tab)
     return map_global
 
 
 @MapProcessorRegistry.register("PerfEventArray")
-def process_perf_event_map(map_name, rval, module):
+def process_perf_event_map(map_name, rval, module, structs_sym_tab):
     """Process a BPF_PERF_EVENT_ARRAY map declaration"""
     logger.info(f"Processing PerfEventArray: {map_name}")
     map_params = _parse_map_params(rval, expected_args=["key_size", "value_size"])
@@ -117,11 +135,11 @@ def process_perf_event_map(map_name, rval, module):
     logger.info(f"Map parameters: {map_params}")
     map_global = create_bpf_map(module, map_name, map_params)
     # Generate debug info for BTF
-    create_map_debug_info(module, map_global, map_name, map_params)
+    create_map_debug_info(module, map_global.sym, map_name, map_params)
     return map_global
 
 
-def process_bpf_map(func_node, module):
+def process_bpf_map(func_node, module, structs_sym_tab):
     """Process a BPF map (a function decorated with @map)"""
     map_name = func_node.name
     logger.info(f"Processing BPF map: {map_name}")
@@ -140,7 +158,7 @@ def process_bpf_map(func_node, module):
     if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name):
         handler = MapProcessorRegistry.get_processor(rval.func.id)
         if handler:
-            return handler(map_name, rval, module)
+            return handler(map_name, rval, module, structs_sym_tab)
         else:
             logger.warning(f"Unknown map type {rval.func.id}, defaulting to HashMap")
             return process_hash_map(map_name, rval, module)

pythonbpf/maps/maps_utils.py

@@ -1,5 +1,17 @@
 from collections.abc import Callable
+from dataclasses import dataclass
+from llvmlite import ir
 from typing import Any
+from .map_types import BPFMapType
+
+
+@dataclass
+class MapSymbol:
+    """Class representing a symbol on the map"""
+
+    type: BPFMapType
+    sym: ir.GlobalVariable
+    params: dict[str, Any] | None = None
 
 
 class MapProcessorRegistry:

tests/passing_tests/hash_map_struct.py

@@ -0,0 +1,42 @@
+from pythonbpf import bpf, section, struct, bpfglobal, compile, map
+from pythonbpf.maps import HashMap
+from pythonbpf.helper import pid
+from ctypes import c_void_p, c_int64
+
+
+@bpf
+@struct
+class val_type:
+    counter: c_int64
+    shizzle: c_int64
+
+
+@bpf
+@map
+def last() -> HashMap:
+    return HashMap(key=val_type, value=c_int64, max_entries=16)
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_clone")
+def hello_world(ctx: c_void_p) -> c_int64:
+    obj = val_type()
+    obj.counter, obj.shizzle = 42, 96
+    t = last.lookup(obj)
+    if t:
+        print(f"Found existing entry: counter={obj.counter}, pid={t}")
+        last.delete(obj)
+        return 0  # type: ignore [return-value]
+    val = pid()
+    last.update(obj, val)
+    print(f"Map updated!, {obj.counter}, {obj.shizzle}, {val}")
+    return 0  # type: ignore [return-value]
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

tests/passing_tests/helpers/bpf_probe_read.py

@@ -0,0 +1,29 @@
+from pythonbpf import bpf, section, bpfglobal, compile, struct
+from ctypes import c_void_p, c_int64, c_uint64, c_uint32
+from pythonbpf.helper import probe_read
+
+
+@bpf
+@struct
+class data_t:
+    pid: c_uint32
+    value: c_uint64
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def test_probe_read(ctx: c_void_p) -> c_int64:
+    """Test bpf_probe_read helper function"""
+    data = data_t()
+    probe_read(data.value, 8, ctx)
+    probe_read(data.pid, 4, ctx)
+    return 0
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

tests/passing_tests/helpers/prandom.py

@@ -0,0 +1,25 @@
+from pythonbpf import bpf, bpfglobal, section, BPF, trace_pipe
+from ctypes import c_void_p, c_int64
+from pythonbpf.helper import random
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_clone")
+def hello_world(ctx: c_void_p) -> c_int64:
+    r = random()
+    print(f"Hello, World!, {r}")
+    return 0  # type: ignore [return-value]
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+# Compile and load
+b = BPF()
+b.load()
+b.attach_all()
+
+trace_pipe()

tests/passing_tests/helpers/smp_processor_id.py

@@ -0,0 +1,40 @@
+from pythonbpf import bpf, section, bpfglobal, compile, struct
+from ctypes import c_void_p, c_int64, c_uint32, c_uint64
+from pythonbpf.helper import smp_processor_id, ktime
+
+
+@bpf
+@struct
+class cpu_event_t:
+    cpu_id: c_uint32
+    timestamp: c_uint64
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def trace_with_cpu(ctx: c_void_p) -> c_int64:
+    """Test bpf_get_smp_processor_id helper function"""
+
+    # Get the current CPU ID
+    cpu = smp_processor_id()
+
+    # Print it
+    print(f"Running on CPU {cpu}")
+
+    # Use it in a struct
+    event = cpu_event_t()
+    event.cpu_id = smp_processor_id()
+    event.timestamp = ktime()
+
+    print(f"Event on CPU {event.cpu_id} at time {event.timestamp}")
+
+    return 0
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

tests/passing_tests/helpers/uid_gid.py

@@ -0,0 +1,31 @@
+from pythonbpf import bpf, section, bpfglobal, compile
+from ctypes import c_void_p, c_int64
+from pythonbpf.helper import uid, pid
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def filter_by_user(ctx: c_void_p) -> c_int64:
+    """Filter events by specific user ID"""
+
+    current_uid = uid()
+
+    # Only trace root user (UID 0)
+    if current_uid == 0:
+        process_id = pid()
+        print(f"Root process {process_id} executed")
+
+    # Or trace specific user (e.g., UID 1000)
+    if current_uid == 1002:
+        print("User 1002 executed something")
+
+    return 0
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()