Add map decorator and simplify type conversion logic

Author: varun-r-mallya

examples/execve2.py

@@ -1,12 +1,12 @@
-from pythonbpf.decorators import bpf, bpfglobal, section
+from pythonbpf.decorators import bpf, map, section, bpfglobal
 from ctypes import c_void_p, c_int64, c_int32, c_uint64
 from pythonbpf.helpers import bpf_ktime_get_ns
 from pythonbpf.maps import HashMap
 
 
 @bpf
-@bpfglobal
-def last():
+@map
+def last() -> HashMap:
     return HashMap(key_type=c_uint64, value_type=c_uint64, max_entries=1)
 
 
@@ -25,5 +25,9 @@ def hello_again(ctx: c_void_p) -> c_int64:
     ts = bpf_ktime_get_ns()
     return c_int64(0)
 
+# @bpf
+# @bpfglobal
+# def LICENSE() -> str:
+#     return "GPL"
 
 LICENSE = "GPL"

pythonbpf/decorators.py

@@ -9,6 +9,10 @@ def bpfglobal(func):
     func._is_bpfglobal = True
     return func
 
+def map(func):
+    """Decorator to mark a function as a BPF map."""
+    func._is_map = True
+    return func
 
 def section(name: str):
     def wrapper(fn):

pythonbpf/functions_pass.py

@@ -90,19 +90,11 @@ def process_bpf_chunk(func_node, module, return_type):
 def create_bpf_map(module, map_name, map_params):
     """Create a BPF map in the module with the given parameters"""
 
-    type_mapping = {
-        'c_uint32': ir.IntType(32),
-        'c_uint64': ir.IntType(64),
-        'c_int32': ir.IntType(32),
-        'c_int64': ir.IntType(64),
-        # Add more mappings as needed
-    }
-
     key_type_str = map_params.get('key_type', 'c_uint32')
     value_type_str = map_params.get('value_type', 'c_uint32')
 
-    key_type = type_mapping.get(key_type_str, ir.IntType(32))
-    value_type = type_mapping.get(value_type_str, ir.IntType(32))
+    key_type = ctypes_to_ir(key_type_str)
+    value_type = ctypes_to_ir(value_type_str)
 
     map_struct_type = ir.LiteralStructType([
         ir.PointerType(),  # type

pythonbpf/globals_pass.py

@@ -42,5 +42,8 @@ def globals_processing(tree, module: ir.Module):
 
                 elif isinstance(dec, ast.Name) and dec.id == "bpfglobal":
                     collected.append(node.name)
+                
+                elif isinstance(dec, ast.Name) and dec.id == "map":
+                    collected.append(node.name)
 
     emit_globals(module, collected)