PythonBPF: Add Compilation Context to allow parallel compilation of multiple bpf programs

pythonbpf/helper/__init__.py

@@ -1,5 +1,5 @@
 from .helper_registry import HelperHandlerRegistry
-from .helper_utils import reset_scratch_pool
+
 from .bpf_helper_handler import (
     handle_helper_call,
     emit_probe_read_kernel_str_call,
@@ -28,9 +28,7 @@ def _register_helper_handler():
     """Register helper call handler with the expression evaluator"""
     from pythonbpf.expr.expr_pass import CallHandlerRegistry
 
-    def helper_call_handler(
-        call, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab
-    ):
+    def helper_call_handler(call, compilation_context, builder, func, local_sym_tab):
         """Check if call is a helper and handle it"""
         import ast
 
@@ -39,17 +37,16 @@ def _register_helper_handler():
             if HelperHandlerRegistry.has_handler(call.func.id):
                 return handle_helper_call(
                     call,
-                    module,
+                    compilation_context,
                     builder,
                     func,
                     local_sym_tab,
-                    map_sym_tab,
-                    structs_sym_tab,
                 )
 
         # Check for method calls (e.g., map.lookup())
         elif isinstance(call.func, ast.Attribute):
             method_name = call.func.attr
+            map_sym_tab = compilation_context.map_sym_tab
 
             # Handle: my_map.lookup(key)
             if isinstance(call.func.value, ast.Name):
@@ -58,12 +55,10 @@ def _register_helper_handler():
                     if HelperHandlerRegistry.has_handler(method_name):
                         return handle_helper_call(
                             call,
-                            module,
+                            compilation_context,
                             builder,
                             func,
                             local_sym_tab,
-                            map_sym_tab,
-                            structs_sym_tab,
                         )
 
         return None
@@ -76,7 +71,6 @@ _register_helper_handler()
 
 __all__ = [
     "HelperHandlerRegistry",
-    "reset_scratch_pool",
     "handle_helper_call",
     "emit_probe_read_kernel_str_call",
     "emit_probe_read_kernel_call",