Skip to content

Commit 2a3f367

Browse files
varun-r-mallyavarun-r-mallya
committed
separate map creation logic
1 parent 08ff076 commit 2a3f367

3 files changed

Lines changed: 97 additions & 83 deletions

File tree

‎pythonbpf/codegen.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from llvmlite import ir
33
from .license_pass import license_processing
44
from .functions_pass import func_proc
5+
from .maps_pass import maps_proc
56
# from .constants_pass import constants_processing
67
from .globals_pass import globals_processing
78

@@ -26,6 +27,7 @@ def processor(source_code, filename, module):
2627
for func_node in bpf_chunks:
2728
print(f"Found BPF function: {func_node.name}")
2829

30+
maps_proc(tree, module, bpf_chunks)
2931
func_proc(tree, module, bpf_chunks)
3032
# For now, we will parse the BPF specific parts of AST
3133

‎pythonbpf/functions_pass.py‎

Lines changed: 1 addition & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -86,96 +86,14 @@ def process_bpf_chunk(func_node, module, return_type):
8686

8787
return func
8888

89-
90-
def create_bpf_map(module, map_name, map_params):
91-
"""Create a BPF map in the module with the given parameters"""
92-
93-
key_type_str = map_params.get('key_type', 'c_uint32')
94-
value_type_str = map_params.get('value_type', 'c_uint32')
95-
96-
key_type = ctypes_to_ir(key_type_str)
97-
value_type = ctypes_to_ir(value_type_str)
98-
99-
map_struct_type = ir.LiteralStructType([
100-
ir.PointerType(), # type
101-
ir.PointerType(), # max_entries
102-
ir.PointerType(), # key_type
103-
ir.PointerType() # value_type
104-
])
105-
106-
map_global = ir.GlobalVariable(module, map_struct_type, name=map_name)
107-
map_global.linkage = 'external'
108-
map_global.initializer = ir.Constant(
109-
map_struct_type, [None, None, None, None])
110-
map_global.section = ".maps"
111-
map_global.align = 8
112-
113-
# TODO: Store map parameters in metadata or a suitable structure
114-
# maps[map_name] = {
115-
# 'global': map_global,
116-
# 'key_type': key_type,
117-
# 'value_type': value_type,
118-
# 'max_entries': map_params.get('max_entries', 1),
119-
# 'map_type': map_params.get('map_type', 'BPF_MAP_TYPE_HASH')
120-
# }
121-
122-
print(f"Created BPF map: {map_name}")
123-
return map_global
124-
125-
126-
def process_bpf_global(func_node, module):
127-
"""Process a BPF global (a function decorated with @bpfglobal)"""
128-
global_name = func_node.name
129-
print(f"Processing BPF global: {global_name}")
130-
131-
# For now, assume single return statement
132-
return_stmt = None
133-
for stmt in func_node.body:
134-
if isinstance(stmt, ast.Return):
135-
return_stmt = stmt
136-
break
137-
if return_stmt is None:
138-
raise ValueError("BPF global must have a return statement")
139-
140-
rval = return_stmt.value
141-
142-
# For now, just handle maps
143-
if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name) and rval.func.id == "HashMap":
144-
print(f"Creating HashMap global: {global_name}")
145-
map_params = {'map_type': 'HASH'}
146-
# Handle positional arguments
147-
if rval.args:
148-
# Assuming order is: key_type, value_type, max_entries
149-
if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name):
150-
map_params['key_type'] = rval.args[0].id
151-
if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name):
152-
map_params['value_type'] = rval.args[1].id
153-
if len(rval.args) >= 3 and isinstance(rval.args[2], ast.Constant):
154-
map_params['max_entries'] = rval.args[2].value
155-
156-
# Handle keyword arguments (these will override any positional args)
157-
for keyword in rval.keywords:
158-
if keyword.arg == "key_type" and isinstance(keyword.value, ast.Name):
159-
map_params['key_type'] = keyword.value.id
160-
elif keyword.arg == "value_type" and isinstance(keyword.value, ast.Name):
161-
map_params['value_type'] = keyword.value.id
162-
elif keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant):
163-
map_params['max_entries'] = keyword.value.value
164-
print(f"Map parameters: {map_params}")
165-
print(create_bpf_map(module, global_name, map_params))
166-
167-
16889
def func_proc(tree, module, chunks):
16990
for func_node in chunks:
170-
# Check if this function is a global
17191
is_global = False
17292
for decorator in func_node.decorator_list:
173-
if isinstance(decorator, ast.Name) and decorator.id == "bpfglobal":
93+
if isinstance(decorator, ast.Name) and decorator.id == "map":
17494
is_global = True
17595
break
17696
if is_global:
177-
print(f"Found BPF global: {func_node.name}")
178-
process_bpf_global(func_node, module)
17997
continue
18098
func_type = get_probe_string(func_node)
18199
print(f"Found probe_string of {func_node.name}: {func_type}")

‎pythonbpf/maps_pass.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
import ast
2+
from llvmlite import ir
3+
from .type_deducer import ctypes_to_ir
4+
5+
def maps_proc(tree, module, chunks):
6+
for func_node in chunks:
7+
# Check if this function is a global
8+
is_global = False
9+
for decorator in func_node.decorator_list:
10+
if isinstance(decorator, ast.Name) and decorator.id == "map":
11+
is_global = True
12+
break
13+
if is_global:
14+
print(f"Found BPF map: {func_node.name}")
15+
process_bpf_global(func_node, module)
16+
continue
17+
18+
19+
def create_bpf_map(module, map_name, map_params):
20+
"""Create a BPF map in the module with the given parameters"""
21+
22+
key_type_str = map_params.get('key_type', 'c_uint32')
23+
value_type_str = map_params.get('value_type', 'c_uint32')
24+
25+
key_type = ctypes_to_ir(key_type_str)
26+
value_type = ctypes_to_ir(value_type_str)
27+
28+
map_struct_type = ir.LiteralStructType([
29+
ir.PointerType(), # type
30+
ir.PointerType(), # max_entries
31+
ir.PointerType(), # key_type
32+
ir.PointerType() # value_type
33+
])
34+
35+
map_global = ir.GlobalVariable(module, map_struct_type, name=map_name)
36+
map_global.linkage = 'external'
37+
map_global.initializer = ir.Constant(
38+
map_struct_type, [None, None, None, None])
39+
map_global.section = ".maps"
40+
map_global.align = 8
41+
42+
# TODO: Store map parameters in metadata or a suitable structure
43+
# maps[map_name] = {
44+
# 'global': map_global,
45+
# 'key_type': key_type,
46+
# 'value_type': value_type,
47+
# 'max_entries': map_params.get('max_entries', 1),
48+
# 'map_type': map_params.get('map_type', 'BPF_MAP_TYPE_HASH')
49+
# }
50+
51+
print(f"Created BPF map: {map_name}")
52+
return map_global
53+
54+
55+
def process_bpf_global(func_node, module):
56+
"""Process a BPF global (a function decorated with @bpfglobal)"""
57+
global_name = func_node.name
58+
print(f"Processing BPF global: {global_name}")
59+
60+
# For now, assume single return statement
61+
return_stmt = None
62+
for stmt in func_node.body:
63+
if isinstance(stmt, ast.Return):
64+
return_stmt = stmt
65+
break
66+
if return_stmt is None:
67+
raise ValueError("BPF global must have a return statement")
68+
69+
rval = return_stmt.value
70+
71+
# For now, just handle maps
72+
if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name) and rval.func.id == "HashMap":
73+
print(f"Creating HashMap global: {global_name}")
74+
map_params = {'map_type': 'HASH'}
75+
# Handle positional arguments
76+
if rval.args:
77+
# Assuming order is: key_type, value_type, max_entries
78+
if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name):
79+
map_params['key_type'] = rval.args[0].id
80+
if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name):
81+
map_params['value_type'] = rval.args[1].id
82+
if len(rval.args) >= 3 and isinstance(rval.args[2], ast.Constant):
83+
map_params['max_entries'] = rval.args[2].value
84+
85+
# Handle keyword arguments (these will override any positional args)
86+
for keyword in rval.keywords:
87+
if keyword.arg == "key_type" and isinstance(keyword.value, ast.Name):
88+
map_params['key_type'] = keyword.value.id
89+
elif keyword.arg == "value_type" and isinstance(keyword.value, ast.Name):
90+
map_params['value_type'] = keyword.value.id
91+
elif keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant):
92+
map_params['max_entries'] = keyword.value.value
93+
print(f"Map parameters: {map_params}")
94+
print(create_bpf_map(module, global_name, map_params))

0 commit comments

Comments
 (0)