From 573bbb350eb2641830c383c30f64bd80cb0ee50c Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyanshchaturvedi18@gmail.com>
Date: Wed, 19 Nov 2025 00:08:15 +0530
Subject: [PATCH] Allow structs to be key/val type for hashmaps

---
 pythonbpf/maps/map_debug_info.py | 57 +++++++++++++++++++++++++++++---
 pythonbpf/maps/maps_pass.py      |  6 ++--
 2 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/pythonbpf/maps/map_debug_info.py b/pythonbpf/maps/map_debug_info.py
index c6d46cc..77a5888 100644
--- a/pythonbpf/maps/map_debug_info.py
+++ b/pythonbpf/maps/map_debug_info.py
@@ -1,4 +1,5 @@
 import logging
+from llvmlite import ir
 from pythonbpf.debuginfo import DebugInfoGenerator
 from .map_types import BPFMapType
 
@@ -8,18 +9,23 @@ 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 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 = []
@@ -104,6 +110,11 @@ def create_ringbuf_debug_info(
 
 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)
@@ -120,4 +131,40 @@ def _get_key_val_dbg_type(name, generator, structs_sym_tab):
 
 
 def _get_struct_debug_type(struct_obj, generator, structs_sym_tab):
-    pass
+    """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
diff --git a/pythonbpf/maps/maps_pass.py b/pythonbpf/maps/maps_pass.py
index 907be57..ed60958 100644
--- a/pythonbpf/maps/maps_pass.py
+++ b/pythonbpf/maps/maps_pass.py
@@ -105,7 +105,9 @@ 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(module, map_global.sym, map_name, map_params)
+    create_ringbuf_debug_info(
+        module, map_global.sym, map_name, map_params, structs_sym_tab
+    )
     return map_global
 
 
@@ -119,7 +121,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)
+    create_map_debug_info(module, map_global.sym, map_name, map_params, structs_sym_tab)
     return map_global