10 Commits

6 changed files with 84 additions and 118 deletions

View File

@ -25,7 +25,7 @@ def handle_binary_op(rval, module, builder, var_name, local_sym_tab, map_sym_tab
# Handle left operand
if isinstance(left, ast.Name):
if left.id in local_sym_tab:
left = recursive_dereferencer(local_sym_tab[left.id][0], builder)
left = recursive_dereferencer(local_sym_tab[left.id].var, builder)
else:
raise SyntaxError(f"Undefined variable: {left.id}")
elif isinstance(left, ast.Constant):
@ -35,7 +35,7 @@ def handle_binary_op(rval, module, builder, var_name, local_sym_tab, map_sym_tab
if isinstance(right, ast.Name):
if right.id in local_sym_tab:
right = recursive_dereferencer(local_sym_tab[right.id][0], builder)
right = recursive_dereferencer(local_sym_tab[right.id].var, builder)
else:
raise SyntaxError(f"Undefined variable: {right.id}")
elif isinstance(right, ast.Constant):
@ -46,26 +46,26 @@ def handle_binary_op(rval, module, builder, var_name, local_sym_tab, map_sym_tab
print(f"left is {left}, right is {right}, op is {op}")
if isinstance(op, ast.Add):
builder.store(builder.add(left, right), local_sym_tab[var_name][0])
builder.store(builder.add(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.Sub):
builder.store(builder.sub(left, right), local_sym_tab[var_name][0])
builder.store(builder.sub(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.Mult):
builder.store(builder.mul(left, right), local_sym_tab[var_name][0])
builder.store(builder.mul(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.Div):
builder.store(builder.sdiv(left, right), local_sym_tab[var_name][0])
builder.store(builder.sdiv(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.Mod):
builder.store(builder.srem(left, right), local_sym_tab[var_name][0])
builder.store(builder.srem(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.LShift):
builder.store(builder.shl(left, right), local_sym_tab[var_name][0])
builder.store(builder.shl(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.RShift):
builder.store(builder.lshr(left, right), local_sym_tab[var_name][0])
builder.store(builder.lshr(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.BitOr):
builder.store(builder.or_(left, right), local_sym_tab[var_name][0])
builder.store(builder.or_(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.BitXor):
builder.store(builder.xor(left, right), local_sym_tab[var_name][0])
builder.store(builder.xor(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.BitAnd):
builder.store(builder.and_(left, right), local_sym_tab[var_name][0])
builder.store(builder.and_(left, right), local_sym_tab[var_name].var)
elif isinstance(op, ast.FloorDiv):
builder.store(builder.udiv(left, right), local_sym_tab[var_name][0])
builder.store(builder.udiv(left, right), local_sym_tab[var_name].var)
else:
raise SyntaxError("Unsupported binary operation")

View File

@ -10,15 +10,13 @@ def eval_expr(
local_sym_tab,
map_sym_tab,
structs_sym_tab=None,
local_var_metadata=None,
):
print(f"Evaluating expression: {ast.dump(expr)}")
print(local_var_metadata)
if isinstance(expr, ast.Name):
if expr.id in local_sym_tab:
var = local_sym_tab[expr.id][0]
var = local_sym_tab[expr.id].var
val = builder.load(var)
return val, local_sym_tab[expr.id][1] # return value and type
return val, local_sym_tab[expr.id].ir_type # return value and type
else:
print(f"Undefined variable {expr.id}")
return None
@ -51,7 +49,7 @@ def eval_expr(
return None
if isinstance(arg, ast.Name):
if arg.id in local_sym_tab:
arg = local_sym_tab[arg.id][0]
arg = local_sym_tab[arg.id].var
else:
print(f"Undefined variable {arg.id}")
return None
@ -60,7 +58,7 @@ def eval_expr(
return None
# Since we are handling only name case, directly take type from sym tab
val = builder.load(arg)
return val, local_sym_tab[expr.args[0].id][1]
return val, local_sym_tab[expr.args[0].id].ir_type
# check for helpers
if HelperHandlerRegistry.has_handler(expr.func.id):
@ -72,7 +70,6 @@ def eval_expr(
local_sym_tab,
map_sym_tab,
structs_sym_tab,
local_var_metadata,
)
elif isinstance(expr.func, ast.Attribute):
print(f"Handling method call: {ast.dump(expr.func)}")
@ -89,7 +86,6 @@ def eval_expr(
local_sym_tab,
map_sym_tab,
structs_sym_tab,
local_var_metadata,
)
elif isinstance(expr.func.value, ast.Name):
obj_name = expr.func.value.id
@ -104,24 +100,21 @@ def eval_expr(
local_sym_tab,
map_sym_tab,
structs_sym_tab,
local_var_metadata,
)
elif isinstance(expr, ast.Attribute):
if isinstance(expr.value, ast.Name):
var_name = expr.value.id
attr_name = expr.attr
if var_name in local_sym_tab:
var_ptr, var_type = local_sym_tab[var_name]
print(f"Loading attribute " f"{attr_name} from variable {var_name}")
var_ptr, var_type, var_metadata = local_sym_tab[var_name]
print(f"Loading attribute {attr_name} from variable {var_name}")
print(f"Variable type: {var_type}, Variable ptr: {var_ptr}")
print(local_var_metadata)
if local_var_metadata and var_name in local_var_metadata:
metadata = structs_sym_tab[local_var_metadata[var_name]]
if attr_name in metadata.fields:
gep = metadata.gep(builder, var_ptr, attr_name)
val = builder.load(gep)
field_type = metadata.field_type(attr_name)
return val, field_type
metadata = structs_sym_tab[var_metadata]
if attr_name in metadata.fields:
gep = metadata.gep(builder, var_ptr, attr_name)
val = builder.load(gep)
field_type = metadata.field_type(attr_name)
return val, field_type
print("Unsupported expression evaluation")
return None
@ -134,11 +127,9 @@ def handle_expr(
local_sym_tab,
map_sym_tab,
structs_sym_tab,
local_var_metadata,
):
"""Handle expression statements in the function body."""
print(f"Handling expression: {ast.dump(expr)}")
print(local_var_metadata)
call = expr.value
if isinstance(call, ast.Call):
eval_expr(
@ -149,7 +140,6 @@ def handle_expr(
local_sym_tab,
map_sym_tab,
structs_sym_tab,
local_var_metadata,
)
else:
print("Unsupported expression type")

View File

@ -1,13 +1,27 @@
from llvmlite import ir
import ast
import logging
from typing import Any
from dataclasses import dataclass
from .helper import HelperHandlerRegistry, handle_helper_call
from .type_deducer import ctypes_to_ir
from .binary_ops import handle_binary_op
from .expr_pass import eval_expr, handle_expr
local_var_metadata: dict[str | Any, Any] = {}
logger = logging.getLogger(__name__)
@dataclass
class LocalSymbol:
var: ir.AllocaInstr
ir_type: ir.Type
metadata: Any = None
def __iter__(self):
yield self.var
yield self.ir_type
yield self.metadata
def get_probe_string(func_node):
@ -47,13 +61,12 @@ def handle_assign(
if isinstance(target, ast.Attribute):
# struct field assignment
field_name = target.attr
if var_name in local_sym_tab and var_name in local_var_metadata:
struct_type = local_var_metadata[var_name]
if var_name in local_sym_tab:
struct_type = local_sym_tab[var_name].metadata
struct_info = structs_sym_tab[struct_type]
if field_name in struct_info.fields:
field_ptr = struct_info.gep(
builder, local_sym_tab[var_name][0], field_name
builder, local_sym_tab[var_name].var, field_name
)
val = eval_expr(
func,
@ -83,18 +96,21 @@ def handle_assign(
elif isinstance(rval, ast.Constant):
if isinstance(rval.value, bool):
if rval.value:
builder.store(ir.Constant(ir.IntType(1), 1), local_sym_tab[var_name][0])
builder.store(
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][0])
builder.store(
ir.Constant(ir.IntType(1), 0), local_sym_tab[var_name].var
)
print(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][0]
ir.Constant(ir.IntType(64), rval.value), local_sym_tab[var_name].var
)
# local_sym_tab[var_name] = var
print(f"Assigned constant {rval.value} to {var_name}")
elif isinstance(rval.value, str):
str_val = rval.value.encode("utf-8") + b"\x00"
@ -108,7 +124,7 @@ def handle_assign(
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][0])
builder.store(str_ptr, local_sym_tab[var_name].var)
print(f"Assigned string constant '{rval.value}' to {var_name}")
else:
print("Unsupported constant type")
@ -126,13 +142,13 @@ def handle_assign(
# 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][0]
ir.Constant(ir_type, rval.args[0].value),
local_sym_tab[var_name].var,
)
print(
f"Assigned {call_type} constant "
f"{rval.args[0].value} to {var_name}"
)
# local_sym_tab[var_name] = var
elif HelperHandlerRegistry.has_handler(call_type):
# var = builder.alloca(ir.IntType(64), name=var_name)
# var.align = 8
@ -144,10 +160,8 @@ def handle_assign(
local_sym_tab,
map_sym_tab,
structs_sym_tab,
local_var_metadata,
)
builder.store(val[0], local_sym_tab[var_name][0])
# local_sym_tab[var_name] = var
builder.store(val[0], local_sym_tab[var_name].var)
print(f"Assigned constant {rval.func.id} to {var_name}")
elif call_type == "deref" and len(rval.args) == 1:
print(f"Handling deref assignment {ast.dump(rval)}")
@ -164,18 +178,15 @@ def handle_assign(
print("Failed to evaluate deref argument")
return
print(f"Dereferenced value: {val}, storing in {var_name}")
builder.store(val[0], local_sym_tab[var_name][0])
# local_sym_tab[var_name] = var
builder.store(val[0], local_sym_tab[var_name].var)
print(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][0])
local_var_metadata[var_name] = call_type
builder.store(ir.Constant(ir_type, None), local_sym_tab[var_name].var)
print(f"Assigned struct {call_type} to {var_name}")
# local_sym_tab[var_name] = var
else:
print(f"Unsupported assignment call type: {call_type}")
elif isinstance(rval.func, ast.Attribute):
@ -198,12 +209,10 @@ def handle_assign(
local_sym_tab,
map_sym_tab,
structs_sym_tab,
local_var_metadata,
)
# var = builder.alloca(ir.IntType(64), name=var_name)
# var.align = 8
builder.store(val[0], local_sym_tab[var_name][0])
# local_sym_tab[var_name] = var
builder.store(val[0], local_sym_tab[var_name].var)
else:
print("Unsupported assignment call structure")
else:
@ -227,7 +236,7 @@ def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
return None
elif isinstance(cond, ast.Name):
if cond.id in local_sym_tab:
var = local_sym_tab[cond.id][0]
var = local_sym_tab[cond.id].var
val = builder.load(var)
if val.type != ir.IntType(1):
# Convert nonzero values to true, zero to false
@ -342,7 +351,6 @@ def process_stmt(
):
print(f"Processing statement: {ast.dump(stmt)}")
if isinstance(stmt, ast.Expr):
print(local_var_metadata)
handle_expr(
func,
module,
@ -351,7 +359,6 @@ def process_stmt(
local_sym_tab,
map_sym_tab,
structs_sym_tab,
local_var_metadata,
)
elif isinstance(stmt, ast.Assign):
handle_assign(
@ -401,6 +408,7 @@ def allocate_mem(
module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
):
for stmt in body:
has_metadata = False
if isinstance(stmt, ast.If):
if stmt.body:
local_sym_tab = allocate_mem(
@ -458,7 +466,7 @@ def allocate_mem(
struct_info = structs_sym_tab[call_type]
ir_type = struct_info.ir_type
var = builder.alloca(ir_type, name=var_name)
local_var_metadata[var_name] = call_type
has_metadata = True
print(
f"Pre-allocated variable {var_name} "
f"for struct {call_type}"
@ -500,7 +508,11 @@ def allocate_mem(
else:
print("Unsupported assignment value type")
continue
local_sym_tab[var_name] = (var, ir_type)
if has_metadata:
local_sym_tab[var_name] = LocalSymbol(var, ir_type, call_type)
else:
local_sym_tab[var_name] = LocalSymbol(var, ir_type)
return local_sym_tab
@ -662,7 +674,7 @@ def infer_return_type(func_node: ast.FunctionDef):
if found_type is None:
found_type = t
elif found_type != t:
raise ValueError("Conflicting return types:" f"{found_type} vs {t}")
raise ValueError(f"Conflicting return types: {found_type} vs {t}")
return found_type or "None"

View File

@ -30,7 +30,6 @@ def bpf_ktime_get_ns_emitter(
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_ktime_get_ns helper function call.
@ -53,7 +52,6 @@ def bpf_map_lookup_elem_emitter(
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_map_lookup_elem helper function call.
@ -89,7 +87,6 @@ def bpf_printk_emitter(
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""Emit LLVM IR for bpf_printk helper function call."""
if not hasattr(func, "_fmt_counter"):
@ -107,7 +104,6 @@ def bpf_printk_emitter(
func,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
elif isinstance(call.args[0], ast.Constant) and isinstance(call.args[0].value, str):
# TODO: We are only supporting single arguments for now.
@ -138,7 +134,6 @@ def bpf_map_update_elem_emitter(
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_map_update_elem helper function call.
@ -190,7 +185,6 @@ def bpf_map_delete_elem_emitter(
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_map_delete_elem helper function call.
@ -228,7 +222,6 @@ def bpf_get_current_pid_tgid_emitter(
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""
Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
@ -255,7 +248,6 @@ def bpf_perf_event_output_handler(
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
if len(call.args) != 1:
raise ValueError(
@ -264,9 +256,7 @@ def bpf_perf_event_output_handler(
data_arg = call.args[0]
ctx_ptr = func.args[0] # First argument to the function is ctx
data_ptr, size_val = get_data_ptr_and_size(
data_arg, local_sym_tab, struct_sym_tab, local_var_metadata
)
data_ptr, size_val = get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab)
# BPF_F_CURRENT_CPU is -1 in 32 bit
flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
@ -304,7 +294,6 @@ def handle_helper_call(
local_sym_tab=None,
map_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""Process a BPF helper function call and emit the appropriate LLVM IR."""
@ -323,7 +312,6 @@ def handle_helper_call(
func,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
# Handle direct function calls (e.g., print(), ktime())

View File

@ -37,7 +37,7 @@ class HelperHandlerRegistry:
def get_var_ptr_from_name(var_name, local_sym_tab):
"""Get a pointer to a variable from the symbol table."""
if local_sym_tab and var_name in local_sym_tab:
return local_sym_tab[var_name][0]
return local_sym_tab[var_name].var
raise ValueError(f"Variable '{var_name}' not found in local symbol table")
@ -72,7 +72,7 @@ def get_flags_val(arg, builder, local_sym_tab):
if isinstance(arg, ast.Name):
if local_sym_tab and arg.id in local_sym_tab:
flags_ptr = local_sym_tab[arg.id][0]
flags_ptr = local_sym_tab[arg.id].var
return builder.load(flags_ptr)
else:
raise ValueError(f"Variable '{arg.id}' not found in local symbol table")
@ -100,7 +100,6 @@ def handle_fstring_print(
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""Handle f-string formatting for bpf_printk emitter."""
fmt_parts = []
@ -118,7 +117,6 @@ def handle_fstring_print(
exprs,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
else:
raise NotImplementedError(f"Unsupported f-string value type: {type(value)}")
@ -138,7 +136,6 @@ def handle_fstring_print(
builder,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
args.append(arg_value)
@ -158,9 +155,7 @@ def _process_constant_in_fstring(cst, fmt_parts, exprs):
)
def _process_fval(
fval, fmt_parts, exprs, local_sym_tab, struct_sym_tab, local_var_metadata
):
def _process_fval(fval, fmt_parts, exprs, local_sym_tab, struct_sym_tab):
"""Process formatted values in f-string."""
logger.debug(f"Processing formatted value: {ast.dump(fval)}")
@ -173,7 +168,6 @@ def _process_fval(
exprs,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
else:
raise NotImplementedError(
@ -184,13 +178,11 @@ def _process_fval(
def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab):
"""Process name nodes in formatted values."""
if local_sym_tab and name_node.id in local_sym_tab:
_, var_type = local_sym_tab[name_node.id]
_, var_type, tmp = local_sym_tab[name_node.id]
_populate_fval(var_type, name_node, fmt_parts, exprs)
def _process_attr_in_fval(
attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab, local_var_metadata
):
def _process_attr_in_fval(attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab):
"""Process attribute nodes in formatted values."""
if (
isinstance(attr_node.value, ast.Name)
@ -200,12 +192,7 @@ def _process_attr_in_fval(
var_name = attr_node.value.id
field_name = attr_node.attr
if not local_var_metadata or var_name not in local_var_metadata:
raise ValueError(
f"Metadata for '{var_name}' not found in local var metadata"
)
var_type = local_var_metadata[var_name]
var_type = local_sym_tab[var_name].metadata
if var_type not in struct_sym_tab:
raise ValueError(
f"Struct '{var_type}' for '{var_name}' not in symbol table"
@ -263,9 +250,7 @@ def _create_format_string_global(fmt_str, func, module, builder):
return builder.bitcast(fmt_gvar, ir.PointerType())
def _prepare_expr_args(
expr, func, module, builder, local_sym_tab, struct_sym_tab, local_var_metadata
):
def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_tab):
"""Evaluate and prepare an expression to use as an arg for bpf_printk."""
val, _ = eval_expr(
func,
@ -275,7 +260,6 @@ def _prepare_expr_args(
local_sym_tab,
None,
struct_sym_tab,
local_var_metadata,
)
if val:
@ -298,34 +282,26 @@ def _prepare_expr_args(
return ir.Constant(ir.IntType(64), 0)
def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab, local_var_metadata):
def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab):
"""Extract data pointer and size information for perf event output."""
if isinstance(data_arg, ast.Name):
data_name = data_arg.id
if local_sym_tab and data_name in local_sym_tab:
data_ptr = local_sym_tab[data_name][0]
data_ptr = local_sym_tab[data_name].var
else:
raise ValueError(
f"Data variable {data_name} not found in local symbol table."
)
# Check if data_name is a struct
if local_var_metadata and data_name in local_var_metadata:
data_type = local_var_metadata[data_name]
if data_type in struct_sym_tab:
struct_info = struct_sym_tab[data_type]
size_val = ir.Constant(ir.IntType(64), struct_info.size)
return data_ptr, size_val
else:
raise ValueError(
f"Struct {data_type} for {data_name} not in symbol table."
)
data_type = local_sym_tab[data_name].metadata
if data_type in struct_sym_tab:
struct_info = struct_sym_tab[data_type]
size_val = ir.Constant(ir.IntType(64), struct_info.size)
return data_ptr, size_val
else:
raise ValueError(
f"Metadata for variable {data_name} "
"not found in local variable metadata."
)
raise ValueError(f"Struct {data_type} for {data_name} not in symbol table.")
else:
raise NotImplementedError(
"Only simple object names are supported " "as data in perf event output."
"Only simple object names are supported as data in perf event output."
)

View File

@ -85,7 +85,7 @@ def create_bpf_map(module, map_name, map_params):
def create_map_debug_info(module, map_global, map_name, map_params):
"""Generate debug information metadata for BPF maps HASH and PERF_EVENT_ARRAY"""
"""Generate debug info metadata for BPF maps HASH and PERF_EVENT_ARRAY"""
generator = DebugInfoGenerator(module)
uint_type = generator.get_uint32_type()