Add has_handler to HelperHandlerRegistry

Author: r41k0u

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):

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)

pythonbpf/helper/helper_utils.py

@@ -23,6 +23,11 @@ def get_handler(cls, helper_name):
         """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."""