From 2a3f367e230be20b9a2ff2a5d1a6ff560e65d65f Mon Sep 17 00:00:00 2001 From: varun-r-mallya Date: Mon, 8 Sep 2025 21:26:27 +0530 Subject: [PATCH] separate map creation logic --- pythonbpf/codegen.py | 2 + pythonbpf/functions_pass.py | 84 +-------------------------------- pythonbpf/maps_pass.py | 94 +++++++++++++++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 83 deletions(-) create mode 100644 pythonbpf/maps_pass.py diff --git a/pythonbpf/codegen.py b/pythonbpf/codegen.py index 921a553..7ce974b 100644 --- a/pythonbpf/codegen.py +++ b/pythonbpf/codegen.py @@ -2,6 +2,7 @@ import ast from llvmlite import ir from .license_pass import license_processing from .functions_pass import func_proc +from .maps_pass import maps_proc # from .constants_pass import constants_processing from .globals_pass import globals_processing @@ -26,6 +27,7 @@ def processor(source_code, filename, module): for func_node in bpf_chunks: print(f"Found BPF function: {func_node.name}") + maps_proc(tree, module, bpf_chunks) func_proc(tree, module, bpf_chunks) # For now, we will parse the BPF specific parts of AST diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index 08b068a..69b488d 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -86,96 +86,14 @@ def process_bpf_chunk(func_node, module, return_type): return func - -def create_bpf_map(module, map_name, map_params): - """Create a BPF map in the module with the given parameters""" - - key_type_str = map_params.get('key_type', 'c_uint32') - value_type_str = map_params.get('value_type', 'c_uint32') - - key_type = ctypes_to_ir(key_type_str) - value_type = ctypes_to_ir(value_type_str) - - map_struct_type = ir.LiteralStructType([ - ir.PointerType(), # type - ir.PointerType(), # max_entries - ir.PointerType(), # key_type - ir.PointerType() # value_type - ]) - - map_global = ir.GlobalVariable(module, map_struct_type, name=map_name) - map_global.linkage = 'external' - map_global.initializer = ir.Constant( - map_struct_type, [None, None, None, None]) - map_global.section = ".maps" - map_global.align = 8 - - # TODO: Store map parameters in metadata or a suitable structure - # maps[map_name] = { - # 'global': map_global, - # 'key_type': key_type, - # 'value_type': value_type, - # 'max_entries': map_params.get('max_entries', 1), - # 'map_type': map_params.get('map_type', 'BPF_MAP_TYPE_HASH') - # } - - print(f"Created BPF map: {map_name}") - return map_global - - -def process_bpf_global(func_node, module): - """Process a BPF global (a function decorated with @bpfglobal)""" - global_name = func_node.name - print(f"Processing BPF global: {global_name}") - - # For now, assume single return statement - return_stmt = None - for stmt in func_node.body: - if isinstance(stmt, ast.Return): - return_stmt = stmt - break - if return_stmt is None: - raise ValueError("BPF global must have a return statement") - - rval = return_stmt.value - - # For now, just handle maps - if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name) and rval.func.id == "HashMap": - print(f"Creating HashMap global: {global_name}") - map_params = {'map_type': 'HASH'} - # Handle positional arguments - if rval.args: - # Assuming order is: key_type, value_type, max_entries - if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name): - map_params['key_type'] = rval.args[0].id - if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name): - map_params['value_type'] = rval.args[1].id - if len(rval.args) >= 3 and isinstance(rval.args[2], ast.Constant): - map_params['max_entries'] = rval.args[2].value - - # Handle keyword arguments (these will override any positional args) - for keyword in rval.keywords: - if keyword.arg == "key_type" and isinstance(keyword.value, ast.Name): - map_params['key_type'] = keyword.value.id - elif keyword.arg == "value_type" and isinstance(keyword.value, ast.Name): - map_params['value_type'] = keyword.value.id - elif keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant): - map_params['max_entries'] = keyword.value.value - print(f"Map parameters: {map_params}") - print(create_bpf_map(module, global_name, map_params)) - - def func_proc(tree, module, chunks): for func_node in chunks: - # Check if this function is a global is_global = False for decorator in func_node.decorator_list: - if isinstance(decorator, ast.Name) and decorator.id == "bpfglobal": + if isinstance(decorator, ast.Name) and decorator.id == "map": is_global = True break if is_global: - print(f"Found BPF global: {func_node.name}") - process_bpf_global(func_node, module) continue func_type = get_probe_string(func_node) print(f"Found probe_string of {func_node.name}: {func_type}") diff --git a/pythonbpf/maps_pass.py b/pythonbpf/maps_pass.py new file mode 100644 index 0000000..dce5860 --- /dev/null +++ b/pythonbpf/maps_pass.py @@ -0,0 +1,94 @@ +import ast +from llvmlite import ir +from .type_deducer import ctypes_to_ir + +def maps_proc(tree, module, chunks): + for func_node in chunks: + # Check if this function is a global + is_global = False + for decorator in func_node.decorator_list: + if isinstance(decorator, ast.Name) and decorator.id == "map": + is_global = True + break + if is_global: + print(f"Found BPF map: {func_node.name}") + process_bpf_global(func_node, module) + continue + + +def create_bpf_map(module, map_name, map_params): + """Create a BPF map in the module with the given parameters""" + + key_type_str = map_params.get('key_type', 'c_uint32') + value_type_str = map_params.get('value_type', 'c_uint32') + + key_type = ctypes_to_ir(key_type_str) + value_type = ctypes_to_ir(value_type_str) + + map_struct_type = ir.LiteralStructType([ + ir.PointerType(), # type + ir.PointerType(), # max_entries + ir.PointerType(), # key_type + ir.PointerType() # value_type + ]) + + map_global = ir.GlobalVariable(module, map_struct_type, name=map_name) + map_global.linkage = 'external' + map_global.initializer = ir.Constant( + map_struct_type, [None, None, None, None]) + map_global.section = ".maps" + map_global.align = 8 + + # TODO: Store map parameters in metadata or a suitable structure + # maps[map_name] = { + # 'global': map_global, + # 'key_type': key_type, + # 'value_type': value_type, + # 'max_entries': map_params.get('max_entries', 1), + # 'map_type': map_params.get('map_type', 'BPF_MAP_TYPE_HASH') + # } + + print(f"Created BPF map: {map_name}") + return map_global + + +def process_bpf_global(func_node, module): + """Process a BPF global (a function decorated with @bpfglobal)""" + global_name = func_node.name + print(f"Processing BPF global: {global_name}") + + # For now, assume single return statement + return_stmt = None + for stmt in func_node.body: + if isinstance(stmt, ast.Return): + return_stmt = stmt + break + if return_stmt is None: + raise ValueError("BPF global must have a return statement") + + rval = return_stmt.value + + # For now, just handle maps + if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name) and rval.func.id == "HashMap": + print(f"Creating HashMap global: {global_name}") + map_params = {'map_type': 'HASH'} + # Handle positional arguments + if rval.args: + # Assuming order is: key_type, value_type, max_entries + if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name): + map_params['key_type'] = rval.args[0].id + if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name): + map_params['value_type'] = rval.args[1].id + if len(rval.args) >= 3 and isinstance(rval.args[2], ast.Constant): + map_params['max_entries'] = rval.args[2].value + + # Handle keyword arguments (these will override any positional args) + for keyword in rval.keywords: + if keyword.arg == "key_type" and isinstance(keyword.value, ast.Name): + map_params['key_type'] = keyword.value.id + elif keyword.arg == "value_type" and isinstance(keyword.value, ast.Name): + map_params['value_type'] = keyword.value.id + elif keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant): + map_params['max_entries'] = keyword.value.value + print(f"Map parameters: {map_params}") + print(create_bpf_map(module, global_name, map_params))