From 61f6743f0a81479afb82f988e545d6994b965261 Mon Sep 17 00:00:00 2001 From: Pragyansh Chaturvedi Date: Wed, 1 Oct 2025 03:53:11 +0530 Subject: [PATCH] Use HelperHandleRegitry --- pythonbpf/expr_pass.py | 8 ++--- pythonbpf/functions_pass.py | 8 ++--- pythonbpf/helper/__init__.py | 2 ++ pythonbpf/helper/bpf_helper_handler.py | 48 ++++++++++++++------------ pythonbpf/{helper => }/helpers.py | 0 5 files changed, 36 insertions(+), 30 deletions(-) create mode 100644 pythonbpf/helper/__init__.py rename pythonbpf/{helper => }/helpers.py (100%) diff --git a/pythonbpf/expr_pass.py b/pythonbpf/expr_pass.py index 1c497be..abbba23 100644 --- a/pythonbpf/expr_pass.py +++ b/pythonbpf/expr_pass.py @@ -23,7 +23,7 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_s return None elif isinstance(expr, ast.Call): # delayed import to avoid circular dependency - from .helper.bpf_helper_handler import helper_func_list, handle_helper_call + from pythonbpf.helper import HelperHandlerRegistry, handle_helper_call if isinstance(expr.func, ast.Name): # check deref @@ -50,21 +50,21 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_s return val, local_sym_tab[expr.args[0].id][1] # check for helpers - if expr.func.id in helper_func_list: + if expr.func.id in HelperHandlerRegistry._handlers: return handle_helper_call( expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata) elif isinstance(expr.func, ast.Attribute): print(f"Handling method call: {ast.dump(expr.func)}") if isinstance(expr.func.value, ast.Call) and isinstance(expr.func.value.func, ast.Name): method_name = expr.func.attr - if method_name in helper_func_list: + if method_name in HelperHandlerRegistry._handlers: return handle_helper_call( expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata) elif isinstance(expr.func.value, ast.Name): obj_name = expr.func.value.id method_name = expr.func.attr if obj_name in map_sym_tab: - if method_name in helper_func_list: + if method_name in HelperHandlerRegistry._handlers: return handle_helper_call( expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata) elif isinstance(expr, ast.Attribute): diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index 04e9c52..38f6349 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -2,7 +2,7 @@ from llvmlite import ir import ast -from .helper.bpf_helper_handler import helper_func_list, handle_helper_call +from .helper import HelperHandlerRegistry, handle_helper_call from .type_deducer import ctypes_to_ir from .binary_ops import handle_binary_op from .expr_pass import eval_expr, handle_expr @@ -113,7 +113,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc print(f"Assigned {call_type} constant " f"{rval.args[0].value} to {var_name}") # local_sym_tab[var_name] = var - elif call_type in helper_func_list: + elif call_type in HelperHandlerRegistry._handlers: # var = builder.alloca(ir.IntType(64), name=var_name) # var.align = 8 val = handle_helper_call( @@ -154,7 +154,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc method_name = rval.func.attr if map_name in map_sym_tab: map_ptr = map_sym_tab[map_name] - if method_name in helper_func_list: + if method_name in HelperHandlerRegistry._handlers: val = handle_helper_call( rval, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata) # var = builder.alloca(ir.IntType(64), name=var_name) @@ -344,7 +344,7 @@ def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_t var.align = ir_type.width // 8 print( f"Pre-allocated variable {var_name} of type {call_type}") - elif call_type in helper_func_list: + elif call_type in HelperHandlerRegistry._handlers: # Assume return type is int64 for now ir_type = ir.IntType(64) var = builder.alloca(ir_type, name=var_name) diff --git a/pythonbpf/helper/__init__.py b/pythonbpf/helper/__init__.py new file mode 100644 index 0000000..5e538d3 --- /dev/null +++ b/pythonbpf/helper/__init__.py @@ -0,0 +1,2 @@ +from .helper_utils import HelperHandlerRegistry +from .bpf_helper_handler import handle_helper_call diff --git a/pythonbpf/helper/bpf_helper_handler.py b/pythonbpf/helper/bpf_helper_handler.py index fe8fbe0..1a0f065 100644 --- a/pythonbpf/helper/bpf_helper_handler.py +++ b/pythonbpf/helper/bpf_helper_handler.py @@ -2,6 +2,7 @@ import ast from llvmlite import ir from pythonbpf.expr_pass import eval_expr from enum import Enum +from .helper_utils import HelperHandlerRegistry class BPFHelperID(Enum): @@ -14,6 +15,7 @@ class BPFHelperID(Enum): BPF_PERF_EVENT_OUTPUT = 25 +@HelperHandlerRegistry.register("ktime") def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): """ Emit LLVM IR for bpf_ktime_get_ns helper function call. @@ -27,6 +29,7 @@ def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab return result, ir.IntType(64) +@HelperHandlerRegistry.register("lookup") def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): """ Emit LLVM IR for bpf_map_lookup_elem helper function call. @@ -66,7 +69,8 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_ 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_addr = ir.Constant(ir.IntType( + 64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False) @@ -74,6 +78,7 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_ return result, ir.PointerType() +@HelperHandlerRegistry.register("print") def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): if not hasattr(func, "_fmt_counter"): func._fmt_counter = 0 @@ -233,6 +238,7 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, return None +@HelperHandlerRegistry.register("update") def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): """ Emit LLVM IR for bpf_map_update_elem helper function call. @@ -315,7 +321,8 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_ fn_ptr_type = ir.PointerType(fn_type) # helper id - fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value) + fn_addr = ir.Constant(ir.IntType( + 64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) if isinstance(flags_val, int): @@ -329,6 +336,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_ return result, None +@HelperHandlerRegistry.register("delete") def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): """ Emit LLVM IR for bpf_map_delete_elem helper function call. @@ -375,7 +383,8 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_ 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_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 @@ -384,12 +393,14 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_ return result, None +@HelperHandlerRegistry.register("pid") def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): """ Emit LLVM IR for bpf_get_current_pid_tgid helper function call. """ # func is an arg to just have a uniform signature with other emitters - helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value) + helper_id = ir.Constant(ir.IntType( + 64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.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) @@ -442,7 +453,8 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sy fn_ptr_type = ir.PointerType(fn_type) # helper id - fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value) + fn_addr = ir.Constant(ir.IntType( + 64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) result = builder.call( @@ -453,24 +465,14 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sy "Only simple object names are supported as data in perf event output.") -helper_func_list = { - "lookup": bpf_map_lookup_elem_emitter, - "print": bpf_printk_emitter, - "ktime": bpf_ktime_get_ns_emitter, - "update": bpf_map_update_elem_emitter, - "delete": bpf_map_delete_elem_emitter, - "pid": bpf_get_current_pid_tgid_emitter, - "output": bpf_perf_event_output_handler, -} - - def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): print(local_var_metadata) if isinstance(call.func, ast.Name): func_name = call.func.id - if func_name in helper_func_list: + hdl_func = HelperHandlerRegistry.get_handler(func_name) + if hdl_func: # it is not a map method call - return helper_func_list[func_name](call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata) + return hdl_func(call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata) else: raise NotImplementedError( f"Function {func_name} is not implemented as a helper function.") @@ -481,9 +483,10 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_ method_name = call.func.attr if map_sym_tab and map_name in map_sym_tab: map_ptr = map_sym_tab[map_name] - if method_name in helper_func_list: + hdl_func = HelperHandlerRegistry.get_handler(method_name) + if hdl_func: print(local_var_metadata) - return helper_func_list[method_name]( + return hdl_func( call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata) else: raise NotImplementedError( @@ -496,8 +499,9 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_ method_name = call.func.attr if map_sym_tab and obj_name in map_sym_tab: map_ptr = map_sym_tab[obj_name] - if method_name in helper_func_list: - return helper_func_list[method_name]( + hdl_func = HelperHandlerRegistry.get_handler(method_name) + if hdl_func: + return hdl_func( call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata) else: raise NotImplementedError( diff --git a/pythonbpf/helper/helpers.py b/pythonbpf/helpers.py similarity index 100% rename from pythonbpf/helper/helpers.py rename to pythonbpf/helpers.py