diff --git a/pythonbpf/expr_pass.py b/pythonbpf/expr_pass.py index abbba23..685484d 100644 --- a/pythonbpf/expr_pass.py +++ b/pythonbpf/expr_pass.py @@ -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 HelperHandlerRegistry._handlers: + if HelperHandlerRegistry.has_handler(expr.func.id): 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 HelperHandlerRegistry._handlers: + if HelperHandlerRegistry.has_handler(method_name): 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 HelperHandlerRegistry._handlers: + if HelperHandlerRegistry.has_handler(method_name): 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 38f6349..e155666 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -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 HelperHandlerRegistry._handlers: + elif HelperHandlerRegistry.has_handler(call_type): # 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 HelperHandlerRegistry._handlers: + if HelperHandlerRegistry.has_handler(method_name): 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 HelperHandlerRegistry._handlers: + elif HelperHandlerRegistry.has_handler(call_type): # 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/helper_utils.py b/pythonbpf/helper/helper_utils.py index 5c9248e..89aa5ca 100644 --- a/pythonbpf/helper/helper_utils.py +++ b/pythonbpf/helper/helper_utils.py @@ -23,6 +23,11 @@ class HelperHandlerRegistry: """Get the handler function for a helper""" return cls._handlers.get(helper_name) + @classmethod + def has_handler(cls, helper_name): + """Check if a handler function is registered for a helper""" + return helper_name in cls._handlers + def get_var_ptr_from_name(var_name, local_sym_tab): """Get a pointer to a variable from the symbol table."""