varun-r-mallya/python-bpf
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/Python-BPF.git synced 2025-12-31 21:06:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Cleanup handle_assign

Browse Source
This commit is contained in:
Pragyansh Chaturvedi
2025-10-12 10:12:45 +05:30
parent 933d2a5c77
commit 105c5a7bd0
1 changed files with 34 additions and 180 deletions
Download Patch File Download Diff File Expand all files Collapse all files

214
pythonbpf/functions/functions_pass.py
View File

@ -6,11 +6,9 @@ from dataclasses import dataclass
from pythonbpf.helper import ( from pythonbpf.helper import (
HelperHandlerRegistry, HelperHandlerRegistry,
handle_helper_call,
reset_scratch_pool, reset_scratch_pool,
) )
from pythonbpf.type_deducer import ctypes_to_ir from pythonbpf.type_deducer import ctypes_to_ir
from pythonbpf.binary_ops import handle_binary_op
from pythonbpf.expr import eval_expr, handle_expr, convert_to_bool from pythonbpf.expr import eval_expr, handle_expr, convert_to_bool
from pythonbpf.assign_pass import handle_variable_assignment from pythonbpf.assign_pass import handle_variable_assignment
@ -80,191 +78,47 @@ def handle_assign(
logger.error(f"Failed to handle assignment to {var_name}") logger.error(f"Failed to handle assignment to {var_name}")
return return
num_types = ("c_int32", "c_int64", "c_uint32", "c_uint64")
logger.info(f"Handling assignment to {ast.dump(target)}") logger.info(f"Handling assignment to {ast.dump(target)}")
if not isinstance(target, ast.Name) and not isinstance(target, ast.Attribute): if not isinstance(target, ast.Name) and not isinstance(target, ast.Attribute):
logger.info("Unsupported assignment target") logger.info("Unsupported assignment target")
return return
var_name = target.id if isinstance(target, ast.Name) else target.value.id var_name = target.id if isinstance(target, ast.Name) else target.value.id
rval = stmt.value rval = stmt.value
if isinstance(target, ast.Attribute):
# struct field assignment # struct field assignment
field_name = target.attr field_name = target.attr
if var_name in local_sym_tab: if var_name in local_sym_tab:
struct_type = local_sym_tab[var_name].metadata struct_type = local_sym_tab[var_name].metadata
struct_info = structs_sym_tab[struct_type] struct_info = structs_sym_tab[struct_type]
if field_name in struct_info.fields: if field_name in struct_info.fields:
field_ptr = struct_info.gep( field_ptr = struct_info.gep(
builder, local_sym_tab[var_name].var, field_name builder, local_sym_tab[var_name].var, field_name
) )
val = eval_expr( val = eval_expr(
func, func,
module, module,
builder, builder,
rval, rval,
local_sym_tab, local_sym_tab,
map_sym_tab, map_sym_tab,
structs_sym_tab, structs_sym_tab,
) )
if isinstance(struct_info.field_type(field_name), ir.ArrayType) and val[ if isinstance(struct_info.field_type(field_name), ir.ArrayType) and val[
1 1
] == ir.PointerType(ir.IntType(8)): ] == ir.PointerType(ir.IntType(8)):
# TODO: Figure it out, not a priority rn # TODO: Figure it out, not a priority rn
# Special case for string assignment to char array # Special case for string assignment to char array
# str_len = struct_info["field_types"][field_idx].count # str_len = struct_info["field_types"][field_idx].count
# assign_string_to_array(builder, field_ptr, val[0], str_len) # assign_string_to_array(builder, field_ptr, val[0], str_len)
# print(f"Assigned to struct field {var_name}.{field_name}") # print(f"Assigned to struct field {var_name}.{field_name}")
pass pass
if val is None: if val is None:
logger.info("Failed to evaluate struct field assignment") logger.info("Failed to evaluate struct field assignment")
return
logger.info(field_ptr)
builder.store(val[0], field_ptr)
logger.info(f"Assigned to struct field {var_name}.{field_name}")
return return
elif isinstance(rval, ast.Constant): logger.info(field_ptr)
if isinstance(rval.value, bool): builder.store(val[0], field_ptr)
if rval.value: logger.info(f"Assigned to struct field {var_name}.{field_name}")
builder.store( return
ir.Constant(ir.IntType(1), 1), local_sym_tab[var_name].var
)
else:
builder.store(
ir.Constant(ir.IntType(1), 0), local_sym_tab[var_name].var
)
logger.info(f"Assigned constant {rval.value} to {var_name}")
elif isinstance(rval.value, int):
# Assume c_int64 for now
# var = builder.alloca(ir.IntType(64), name=var_name)
# var.align = 8
builder.store(
ir.Constant(ir.IntType(64), rval.value), local_sym_tab[var_name].var
)
logger.info(f"Assigned constant {rval.value} to {var_name}")
elif isinstance(rval.value, str):
str_val = rval.value.encode("utf-8") + b"\x00"
str_const = ir.Constant(
ir.ArrayType(ir.IntType(8), len(str_val)), bytearray(str_val)
)
global_str = ir.GlobalVariable(
module, str_const.type, name=f"{var_name}_str"
)
global_str.linkage = "internal"
global_str.global_constant = True
global_str.initializer = str_const
str_ptr = builder.bitcast(global_str, ir.PointerType(ir.IntType(8)))
builder.store(str_ptr, local_sym_tab[var_name].var)
logger.info(f"Assigned string constant '{rval.value}' to {var_name}")
else:
logger.info("Unsupported constant type")
elif isinstance(rval, ast.Call):
if isinstance(rval.func, ast.Name):
call_type = rval.func.id
logger.info(f"Assignment call type: {call_type}")
if (
call_type in num_types
and len(rval.args) == 1
and isinstance(rval.args[0], ast.Constant)
and isinstance(rval.args[0].value, int)
):
ir_type = ctypes_to_ir(call_type)
# var = builder.alloca(ir_type, name=var_name)
# var.align = ir_type.width // 8
builder.store(
ir.Constant(ir_type, rval.args[0].value),
local_sym_tab[var_name].var,
)
logger.info(
f"Assigned {call_type} constant {rval.args[0].value} to {var_name}"
)
elif HelperHandlerRegistry.has_handler(call_type):
# var = builder.alloca(ir.IntType(64), name=var_name)
# var.align = 8
val = handle_helper_call(
rval,
module,
builder,
func,
local_sym_tab,
map_sym_tab,
structs_sym_tab,
)
builder.store(val[0], local_sym_tab[var_name].var)
logger.info(f"Assigned constant {rval.func.id} to {var_name}")
elif call_type == "deref" and len(rval.args) == 1:
logger.info(f"Handling deref assignment {ast.dump(rval)}")
val = eval_expr(
func,
module,
builder,
rval,
local_sym_tab,
map_sym_tab,
structs_sym_tab,
)
if val is None:
logger.info("Failed to evaluate deref argument")
return
logger.info(f"Dereferenced value: {val}, storing in {var_name}")
builder.store(val[0], local_sym_tab[var_name].var)
logger.info(f"Dereferenced and assigned to {var_name}")
elif call_type in structs_sym_tab and len(rval.args) == 0:
struct_info = structs_sym_tab[call_type]
ir_type = struct_info.ir_type
# var = builder.alloca(ir_type, name=var_name)
# Null init
builder.store(ir.Constant(ir_type, None), local_sym_tab[var_name].var)
logger.info(f"Assigned struct {call_type} to {var_name}")
else:
logger.info(f"Unsupported assignment call type: {call_type}")
elif isinstance(rval.func, ast.Attribute):
logger.info(f"Assignment call attribute: {ast.dump(rval.func)}")
if isinstance(rval.func.value, ast.Name):
if rval.func.value.id in map_sym_tab:
map_name = rval.func.value.id
method_name = rval.func.attr
if HelperHandlerRegistry.has_handler(method_name):
val = handle_helper_call(
rval,
module,
builder,
func,
local_sym_tab,
map_sym_tab,
structs_sym_tab,
)
builder.store(val[0], local_sym_tab[var_name].var)
else:
# TODO: probably a struct access
logger.info(f"TODO STRUCT ACCESS {ast.dump(rval)}")
elif isinstance(rval.func.value, ast.Call) and isinstance(
rval.func.value.func, ast.Name
):
map_name = rval.func.value.func.id
method_name = rval.func.attr
if map_name in map_sym_tab:
if HelperHandlerRegistry.has_handler(method_name):
val = handle_helper_call(
rval,
module,
builder,
func,
local_sym_tab,
map_sym_tab,
structs_sym_tab,
)
# var = builder.alloca(ir.IntType(64), name=var_name)
# var.align = 8
builder.store(val[0], local_sym_tab[var_name].var)
else:
logger.info("Unsupported assignment call structure")
else:
logger.info("Unsupported assignment call function type")
elif isinstance(rval, ast.BinOp):
handle_binary_op(rval, builder, var_name, local_sym_tab)
else:
logger.info("Unsupported assignment value type")
def handle_cond( def handle_cond(