Use get_flags_val in bpf_map_update_elem_emitter

examples/sys_sync.py

@@ -1,5 +1,5 @@
 from pythonbpf import bpf, map, section, bpfglobal, compile
-from pythonbpf.helper.helpers import ktime
+from pythonbpf.helpers import ktime
 from pythonbpf.maps import HashMap
 
 from ctypes import c_void_p, c_int64, c_uint64

pythonbpf/helper/bpf_helper_handler.py

@@ -2,7 +2,7 @@ import ast
 from llvmlite import ir
 from pythonbpf.expr_pass import eval_expr
 from enum import Enum
-from .helper_utils import HelperHandlerRegistry
+from .helper_utils import HelperHandlerRegistry, get_or_create_ptr_from_arg, get_flags_val
 
 
 class BPFHelperID(Enum):
@@ -38,31 +38,10 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func,
     """
     Emit LLVM IR for bpf_map_lookup_elem helper function call.
     """
-    if call.args and len(call.args) != 1:
-        raise ValueError("Map lookup expects exactly one argument, got "
+    if not call.args or len(call.args) != 1:
+        raise ValueError("Map lookup expects exactly one argument (key), got "
                          f"{len(call.args)}")
-    key_arg = call.args[0]
-    if isinstance(key_arg, ast.Name):
-        key_name = key_arg.id
-        if local_sym_tab and key_name in local_sym_tab:
-            key_ptr = local_sym_tab[key_name][0]
-        else:
-            raise ValueError(
-                f"Key variable {key_name} not found in local symbol table.")
-    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
-        # handle constant integer keys
-        key_val = key_arg.value
-        key_type = ir.IntType(64)
-        key_ptr = builder.alloca(key_type)
-        key_ptr.align = key_type // 8
-        builder.store(ir.Constant(key_type, key_val), key_ptr)
-    else:
-        raise NotImplementedError(
-            "Only simple variable names are supported as keys in map lookup.")
-
-    if key_ptr is None:
-        raise ValueError("Key pointer is None.")
-
+    key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
     map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
 
     fn_type = ir.FunctionType(
@@ -72,7 +51,6 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func,
     )
     fn_ptr_type = ir.PointerType(fn_type)
 
-    # Helper ID 1 is bpf_map_lookup_elem
     fn_addr = ir.Constant(ir.IntType(
         64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
     fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
@@ -264,65 +242,9 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func,
     value_arg = call.args[1]
     flags_arg = call.args[2] if len(call.args) > 2 else None
 
-    # Handle key
-    if isinstance(key_arg, ast.Name):
-        key_name = key_arg.id
-        if local_sym_tab and key_name in local_sym_tab:
-            key_ptr = local_sym_tab[key_name][0]
-        else:
-            raise ValueError(
-                f"Key variable {key_name} not found in local symbol table.")
-    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
-        # Handle constant integer keys
-        key_val = key_arg.value
-        key_type = ir.IntType(64)
-        key_ptr = builder.alloca(key_type)
-        key_ptr.align = key_type.width // 8
-        builder.store(ir.Constant(key_type, key_val), key_ptr)
-    else:
-        raise NotImplementedError(
-            "Only simple variable names and integer constants are supported as keys in map update.")
-
-    # Handle value
-    if isinstance(value_arg, ast.Name):
-        value_name = value_arg.id
-        if local_sym_tab and value_name in local_sym_tab:
-            value_ptr = local_sym_tab[value_name][0]
-        else:
-            raise ValueError(
-                f"Value variable {value_name} not found in local symbol table.")
-    elif isinstance(value_arg, ast.Constant) and isinstance(value_arg.value, int):
-        # Handle constant integers
-        value_val = value_arg.value
-        value_type = ir.IntType(64)
-        value_ptr = builder.alloca(value_type)
-        value_ptr.align = value_type.width // 8
-        builder.store(ir.Constant(value_type, value_val), value_ptr)
-    else:
-        raise NotImplementedError(
-            "Only simple variable names and integer constants are supported as values in map update.")
-
-    # Handle flags argument (defaults to 0)
-    if flags_arg is not None:
-        if isinstance(flags_arg, ast.Constant) and isinstance(flags_arg.value, int):
-            flags_val = flags_arg.value
-        elif isinstance(flags_arg, ast.Name):
-            flags_name = flags_arg.id
-            if local_sym_tab and flags_name in local_sym_tab:
-                # Assume it's a stored integer value, load it
-                flags_ptr = local_sym_tab[flags_name][0]
-                flags_val = builder.load(flags_ptr)
-            else:
-                raise ValueError(
-                    f"Flags variable {flags_name} not found in local symbol table.")
-        else:
-            raise NotImplementedError(
-                "Only integer constants and simple variable names are supported as flags in map update.")
-    else:
-        flags_val = 0
-
-    if key_ptr is None or value_ptr is None:
-        raise ValueError("Key pointer or value pointer is None.")
+    key_ptr = get_or_create_ptr_from_arg(key_arg, builder, local_sym_tab)
+    value_ptr = get_or_create_ptr_from_arg(value_arg, builder, local_sym_tab)
+    flags_val = get_flags_val(flags_arg, builder, local_sym_tab)
 
     map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
     fn_type = ir.FunctionType(
@@ -332,7 +254,6 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func,
     )
     fn_ptr_type = ir.PointerType(fn_type)
 
-    # helper id
     fn_addr = ir.Constant(ir.IntType(
         64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
     fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
@@ -356,36 +277,10 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func,
     Emit LLVM IR for bpf_map_delete_elem helper function call.
     Expected call signature: map.delete(key)
     """
-    # Check for correct number of arguments
     if not call.args or len(call.args) != 1:
-        raise ValueError("Map delete expects exactly 1 argument (key), got "
+        raise ValueError("Map delete expects exactly one argument (key), got "
                          f"{len(call.args)}")
-
-    key_arg = call.args[0]
-
-    # Handle key argument
-    if isinstance(key_arg, ast.Name):
-        key_name = key_arg.id
-        if local_sym_tab and key_name in local_sym_tab:
-            key_ptr = local_sym_tab[key_name][0]
-        else:
-            raise ValueError(
-                f"Key variable {key_name} not found in local symbol table.")
-    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
-        # Handle constant integer keys
-        key_val = key_arg.value
-        key_type = ir.IntType(64)
-        key_ptr = builder.alloca(key_type)
-        key_ptr.align = key_type.width // 8
-        builder.store(ir.Constant(key_type, key_val), key_ptr)
-    else:
-        raise NotImplementedError(
-            "Only simple variable names and integer constants are supported as keys in map delete.")
-
-    if key_ptr is None:
-        raise ValueError("Key pointer is None.")
-
-    # Cast map pointer to void*
+    key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
     map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
 
     # Define function type for bpf_map_delete_elem
@@ -396,12 +291,10 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func,
     )
     fn_ptr_type = ir.PointerType(fn_type)
 
-    # Helper ID 3 is bpf_map_delete_elem
     fn_addr = ir.Constant(ir.IntType(
         64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
     fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
 
-    # Call the helper function
     result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
 
     return result, None

pythonbpf/helper/helper_utils.py

@@ -1,3 +1,7 @@
+import ast
+from llvmlite import ir
+
+
 class HelperHandlerRegistry:
     """Registry for BPF helpers"""
     _handlers = {}
@@ -14,3 +18,52 @@ class HelperHandlerRegistry:
     def get_handler(cls, helper_name):
         """Get the handler function for a helper"""
         return cls._handlers.get(helper_name)
+
+
+def get_var_ptr_from_name(var_name, local_sym_tab):
+    """Get a pointer to a variable from the symbol table."""
+    if local_sym_tab and var_name in local_sym_tab:
+        return local_sym_tab[var_name][0]
+    raise ValueError(f"Variable '{var_name}' not found in local symbol table")
+
+
+def create_int_constant_ptr(value, builder, int_width=64):
+    """Create a pointer to an integer constant."""
+    # Default to 64-bit integer
+    int_type = ir.IntType(int_width)
+    ptr = builder.alloca(int_type)
+    ptr.align = int_type.width // 8
+    builder.store(ir.Constant(int_type, value), ptr)
+    return ptr
+
+
+def get_or_create_ptr_from_arg(arg, builder, local_sym_tab):
+    """Extract or create pointer from the call arguments."""
+
+    if isinstance(arg, ast.Name):
+        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)
+    else:
+        raise NotImplementedError(
+            "Only simple variable names are supported as args in map helpers.")
+    return ptr
+
+
+def get_flags_val(arg, builder, local_sym_tab):
+    """Extract or create flags value from the call arguments."""
+    if not arg:
+        return 0
+
+    if isinstance(arg, ast.Name):
+        if local_sym_tab and arg.id in local_sym_tab:
+            flags_ptr = local_sym_tab[arg.id][0]
+            return builder.load(flags_ptr)
+        else:
+            raise ValueError(
+                f"Variable '{arg.id}' not found in local symbol table")
+    elif isinstance(arg, ast.Constant) and isinstance(arg.value, int):
+        return arg.value
+
+    raise NotImplementedError(
+        "Only simple variable names or integer constants are supported as flags in map helpers.")