Skip to content

Commit 61f6743

Browse files
r41k0ur41k0u
committed
Use HelperHandleRegitry
1 parent 6cd0749 commit 61f6743

5 files changed

Lines changed: 36 additions & 30 deletions

File tree

‎pythonbpf/expr_pass.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_s
2323
return None
2424
elif isinstance(expr, ast.Call):
2525
# delayed import to avoid circular dependency
26-
from .helper.bpf_helper_handler import helper_func_list, handle_helper_call
26+
from pythonbpf.helper import HelperHandlerRegistry, handle_helper_call
2727

2828
if isinstance(expr.func, ast.Name):
2929
# check deref
@@ -50,21 +50,21 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_s
5050
return val, local_sym_tab[expr.args[0].id][1]
5151

5252
# check for helpers
53-
if expr.func.id in helper_func_list:
53+
if expr.func.id in HelperHandlerRegistry._handlers:
5454
return handle_helper_call(
5555
expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
5656
elif isinstance(expr.func, ast.Attribute):
5757
print(f"Handling method call: {ast.dump(expr.func)}")
5858
if isinstance(expr.func.value, ast.Call) and isinstance(expr.func.value.func, ast.Name):
5959
method_name = expr.func.attr
60-
if method_name in helper_func_list:
60+
if method_name in HelperHandlerRegistry._handlers:
6161
return handle_helper_call(
6262
expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
6363
elif isinstance(expr.func.value, ast.Name):
6464
obj_name = expr.func.value.id
6565
method_name = expr.func.attr
6666
if obj_name in map_sym_tab:
67-
if method_name in helper_func_list:
67+
if method_name in HelperHandlerRegistry._handlers:
6868
return handle_helper_call(
6969
expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
7070
elif isinstance(expr, ast.Attribute):

‎pythonbpf/functions_pass.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import ast
33

44

5-
from .helper.bpf_helper_handler import helper_func_list, handle_helper_call
5+
from .helper import HelperHandlerRegistry, handle_helper_call
66
from .type_deducer import ctypes_to_ir
77
from .binary_ops import handle_binary_op
88
from .expr_pass import eval_expr, handle_expr
@@ -113,7 +113,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
113113
print(f"Assigned {call_type} constant "
114114
f"{rval.args[0].value} to {var_name}")
115115
# local_sym_tab[var_name] = var
116-
elif call_type in helper_func_list:
116+
elif call_type in HelperHandlerRegistry._handlers:
117117
# var = builder.alloca(ir.IntType(64), name=var_name)
118118
# var.align = 8
119119
val = handle_helper_call(
@@ -154,7 +154,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
154154
method_name = rval.func.attr
155155
if map_name in map_sym_tab:
156156
map_ptr = map_sym_tab[map_name]
157-
if method_name in helper_func_list:
157+
if method_name in HelperHandlerRegistry._handlers:
158158
val = handle_helper_call(
159159
rval, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
160160
# 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
344344
var.align = ir_type.width // 8
345345
print(
346346
f"Pre-allocated variable {var_name} of type {call_type}")
347-
elif call_type in helper_func_list:
347+
elif call_type in HelperHandlerRegistry._handlers:
348348
# Assume return type is int64 for now
349349
ir_type = ir.IntType(64)
350350
var = builder.alloca(ir_type, name=var_name)

‎pythonbpf/helper/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
from .helper_utils import HelperHandlerRegistry
2+
from .bpf_helper_handler import handle_helper_call

‎pythonbpf/helper/bpf_helper_handler.py‎

Lines changed: 26 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from llvmlite import ir
33
from pythonbpf.expr_pass import eval_expr
44
from enum import Enum
5+
from .helper_utils import HelperHandlerRegistry
56

67

78
class BPFHelperID(Enum):
@@ -14,6 +15,7 @@ class BPFHelperID(Enum):
1415
BPF_PERF_EVENT_OUTPUT = 25
1516

1617

18+
@HelperHandlerRegistry.register("ktime")
1719
def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
1820
"""
1921
Emit LLVM IR for bpf_ktime_get_ns helper function call.
@@ -27,6 +29,7 @@ def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab
2729
return result, ir.IntType(64)
2830

2931

32+
@HelperHandlerRegistry.register("lookup")
3033
def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
3134
"""
3235
Emit LLVM IR for bpf_map_lookup_elem helper function call.
@@ -66,14 +69,16 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_
6669
fn_ptr_type = ir.PointerType(fn_type)
6770

6871
# Helper ID 1 is bpf_map_lookup_elem
69-
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
72+
fn_addr = ir.Constant(ir.IntType(
73+
64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
7074
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
7175

7276
result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
7377

7478
return result, ir.PointerType()
7579

7680

81+
@HelperHandlerRegistry.register("print")
7782
def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
7883
if not hasattr(func, "_fmt_counter"):
7984
func._fmt_counter = 0
@@ -233,6 +238,7 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None,
233238
return None
234239

235240

241+
@HelperHandlerRegistry.register("update")
236242
def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
237243
"""
238244
Emit LLVM IR for bpf_map_update_elem helper function call.
@@ -315,7 +321,8 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_
315321
fn_ptr_type = ir.PointerType(fn_type)
316322

317323
# helper id
318-
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
324+
fn_addr = ir.Constant(ir.IntType(
325+
64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
319326
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
320327

321328
if isinstance(flags_val, int):
@@ -329,6 +336,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_
329336
return result, None
330337

331338

339+
@HelperHandlerRegistry.register("delete")
332340
def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
333341
"""
334342
Emit LLVM IR for bpf_map_delete_elem helper function call.
@@ -375,7 +383,8 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_
375383
fn_ptr_type = ir.PointerType(fn_type)
376384

377385
# Helper ID 3 is bpf_map_delete_elem
378-
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
386+
fn_addr = ir.Constant(ir.IntType(
387+
64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
379388
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
380389

381390
# Call the helper function
@@ -384,12 +393,14 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_
384393
return result, None
385394

386395

396+
@HelperHandlerRegistry.register("pid")
387397
def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
388398
"""
389399
Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
390400
"""
391401
# func is an arg to just have a uniform signature with other emitters
392-
helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value)
402+
helper_id = ir.Constant(ir.IntType(
403+
64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value)
393404
fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
394405
fn_ptr_type = ir.PointerType(fn_type)
395406
fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
@@ -442,7 +453,8 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sy
442453
fn_ptr_type = ir.PointerType(fn_type)
443454

444455
# helper id
445-
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value)
456+
fn_addr = ir.Constant(ir.IntType(
457+
64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value)
446458
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
447459

448460
result = builder.call(
@@ -453,24 +465,14 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sy
453465
"Only simple object names are supported as data in perf event output.")
454466

455467

456-
helper_func_list = {
457-
"lookup": bpf_map_lookup_elem_emitter,
458-
"print": bpf_printk_emitter,
459-
"ktime": bpf_ktime_get_ns_emitter,
460-
"update": bpf_map_update_elem_emitter,
461-
"delete": bpf_map_delete_elem_emitter,
462-
"pid": bpf_get_current_pid_tgid_emitter,
463-
"output": bpf_perf_event_output_handler,
464-
}
465-
466-
467468
def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
468469
print(local_var_metadata)
469470
if isinstance(call.func, ast.Name):
470471
func_name = call.func.id
471-
if func_name in helper_func_list:
472+
hdl_func = HelperHandlerRegistry.get_handler(func_name)
473+
if hdl_func:
472474
# it is not a map method call
473-
return helper_func_list[func_name](call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
475+
return hdl_func(call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
474476
else:
475477
raise NotImplementedError(
476478
f"Function {func_name} is not implemented as a helper function.")
@@ -481,9 +483,10 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
481483
method_name = call.func.attr
482484
if map_sym_tab and map_name in map_sym_tab:
483485
map_ptr = map_sym_tab[map_name]
484-
if method_name in helper_func_list:
486+
hdl_func = HelperHandlerRegistry.get_handler(method_name)
487+
if hdl_func:
485488
print(local_var_metadata)
486-
return helper_func_list[method_name](
489+
return hdl_func(
487490
call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
488491
else:
489492
raise NotImplementedError(
@@ -496,8 +499,9 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
496499
method_name = call.func.attr
497500
if map_sym_tab and obj_name in map_sym_tab:
498501
map_ptr = map_sym_tab[obj_name]
499-
if method_name in helper_func_list:
500-
return helper_func_list[method_name](
502+
hdl_func = HelperHandlerRegistry.get_handler(method_name)
503+
if hdl_func:
504+
return hdl_func(
501505
call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
502506
else:
503507
raise NotImplementedError(

pythonbpf/helper/helpers.py renamed to pythonbpf/helpers.py

File renamed without changes.

0 commit comments

Comments
 (0)