Pass down structs_sym_tab to maps_debug_info, allow vmlinux enums to be used in an indexed format for map declaration

pythonbpf/allocation_pass.py

@@ -7,7 +7,6 @@ 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__)
 
@@ -26,9 +25,7 @@ def create_targets_and_rvals(stmt):
     return stmt.targets, [stmt.value]
 
 
-def handle_assign_allocation(
+def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
-    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)}")
@@ -58,9 +55,7 @@ def handle_assign_allocation(
 
         # Determine type and allocate based on rval
         if isinstance(rval, ast.Call):
-            _allocate_for_call(
+            _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab)
-                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):
@@ -79,9 +74,7 @@ def handle_assign_allocation(
             )
 
 
-def _allocate_for_call(
+def _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab):
-    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):
@@ -123,74 +116,15 @@ def _allocate_for_call(
 
     elif isinstance(rval.func, ast.Attribute):
         # Map method calls - need double allocation for ptr handling
-        _allocate_for_map_method(
+        _allocate_for_map_method(builder, var_name, local_sym_tab)
-            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(
+def _allocate_for_map_method(builder, var_name, local_sym_tab):
-    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)
@@ -201,9 +135,7 @@ def _allocate_for_map_method_fallback(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(
+    logger.info(f"Pre-allocated {var_name} and {var_name}_tmp for map method")
-        f"Pre-allocated {var_name} and {var_name}_tmp for map method (fallback)"
-    )
 
 
 def _allocate_for_constant(builder, var_name, rval, local_sym_tab):

pythonbpf/debuginfo/debug_info_generator.py

@@ -49,10 +49,6 @@ 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

@@ -147,9 +147,7 @@ def allocate_mem(
                 structs_sym_tab,
             )
         elif isinstance(stmt, ast.Assign):
-            handle_assign_allocation(
+            handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab)
-                builder, stmt, local_sym_tab, map_sym_tab, structs_sym_tab
-            )
 
     allocate_temp_pool(builder, max_temps_needed, local_sym_tab)
 

pythonbpf/maps/map_debug_info.py

@@ -1,31 +1,22 @@
-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, 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
+        array_type if "key_size" in map_params else ulong_type, 64
-        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
+        array_type if "value_size" in map_params else ulong_type, 64
-        if "value_size" in map_params
-        else _get_key_val_dbg_type(map_params.get("value"), generator, structs_sym_tab),
-        64,
     )
 
     elements_arr = []
@@ -106,65 +97,3 @@ def create_ringbuf_debug_info(
     )
     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_pass.py

@@ -48,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 MapSymbol(type=map_params["type"], sym=map_global, params=map_params)
+    return MapSymbol(type=map_params["type"], sym=map_global)
 
 
 def _parse_map_params(rval, expected_args=None):
@@ -105,9 +105,7 @@ def process_ringbuf_map(map_name, rval, module, structs_sym_tab):
     logger.info(f"Ringbuf map parameters: {map_params}")
 
     map_global = create_bpf_map(module, map_name, map_params)
-    create_ringbuf_debug_info(
+    create_ringbuf_debug_info(module, map_global.sym, map_name, map_params)
-        module, map_global.sym, map_name, map_params, structs_sym_tab
-    )
     return map_global
 
 
@@ -121,7 +119,7 @@ def process_hash_map(map_name, rval, module, structs_sym_tab):
     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.sym, map_name, map_params, structs_sym_tab)
+    create_map_debug_info(module, map_global.sym, map_name, map_params)
     return map_global
 
 

pythonbpf/maps/maps_utils.py

@@ -11,7 +11,6 @@ class MapSymbol:
 
     type: BPFMapType
     sym: ir.GlobalVariable
-    params: dict[str, Any] | None = None
 
 
 class MapProcessorRegistry:

tests/passing_tests/hash_map_struct.py

@@ -1,42 +0,0 @@
-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()