mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2026-03-25 22:51:28 +00:00
469 lines
14 KiB
Python
469 lines
14 KiB
Python
"""
|
|
Expression evaluation and LLVM IR generation.
|
|
|
|
This module handles the evaluation of Python expressions in BPF programs,
|
|
including variables, constants, function calls, comparisons, boolean
|
|
operations, and more.
|
|
"""
|
|
|
|
import ast
|
|
from llvmlite import ir
|
|
from logging import Logger
|
|
import logging
|
|
from typing import Dict
|
|
|
|
from pythonbpf.type_deducer import ctypes_to_ir, is_ctypes
|
|
from .type_normalization import convert_to_bool, handle_comparator
|
|
|
|
logger: Logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _handle_name_expr(expr: ast.Name, local_sym_tab: Dict, builder: ir.IRBuilder):
|
|
"""Handle ast.Name expressions."""
|
|
if expr.id in local_sym_tab:
|
|
var = local_sym_tab[expr.id].var
|
|
val = builder.load(var)
|
|
return val, local_sym_tab[expr.id].ir_type
|
|
else:
|
|
logger.info(f"Undefined variable {expr.id}")
|
|
return None
|
|
|
|
|
|
def _handle_constant_expr(expr: ast.Constant):
|
|
"""Handle ast.Constant expressions."""
|
|
if isinstance(expr.value, int) or isinstance(expr.value, bool):
|
|
return ir.Constant(ir.IntType(64), int(expr.value)), ir.IntType(64)
|
|
else:
|
|
logger.error("Unsupported constant type")
|
|
return None
|
|
|
|
|
|
def _handle_attribute_expr(
|
|
expr: ast.Attribute,
|
|
local_sym_tab: Dict,
|
|
structs_sym_tab: Dict,
|
|
builder: ir.IRBuilder,
|
|
):
|
|
"""Handle ast.Attribute expressions for struct field access."""
|
|
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, var_metadata = local_sym_tab[var_name]
|
|
logger.info(f"Loading attribute {attr_name} from variable {var_name}")
|
|
logger.info(f"Variable type: {var_type}, Variable ptr: {var_ptr}")
|
|
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
|
|
return None
|
|
|
|
|
|
def _handle_deref_call(expr: ast.Call, local_sym_tab: Dict, builder: ir.IRBuilder):
|
|
"""Handle deref function calls."""
|
|
logger.info(f"Handling deref {ast.dump(expr)}")
|
|
if len(expr.args) != 1:
|
|
logger.info("deref takes exactly one argument")
|
|
return None
|
|
|
|
arg = expr.args[0]
|
|
if (
|
|
isinstance(arg, ast.Call)
|
|
and isinstance(arg.func, ast.Name)
|
|
and arg.func.id == "deref"
|
|
):
|
|
logger.info("Multiple deref not supported")
|
|
return None
|
|
|
|
if isinstance(arg, ast.Name):
|
|
if arg.id in local_sym_tab:
|
|
arg_ptr = local_sym_tab[arg.id].var
|
|
else:
|
|
logger.info(f"Undefined variable {arg.id}")
|
|
return None
|
|
else:
|
|
logger.info("Unsupported argument type for deref")
|
|
return None
|
|
|
|
if arg_ptr is None:
|
|
logger.info("Failed to evaluate deref argument")
|
|
return None
|
|
|
|
# Load the value from pointer
|
|
val = builder.load(arg_ptr)
|
|
return val, local_sym_tab[arg.id].ir_type
|
|
|
|
|
|
def _handle_ctypes_call(
|
|
func,
|
|
module,
|
|
builder,
|
|
expr,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab=None,
|
|
):
|
|
"""Handle ctypes type constructor calls."""
|
|
if len(expr.args) != 1:
|
|
logger.info("ctypes constructor takes exactly one argument")
|
|
return None
|
|
|
|
arg = expr.args[0]
|
|
val = eval_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
arg,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
if val is None:
|
|
logger.info("Failed to evaluate argument to ctypes constructor")
|
|
return None
|
|
call_type = expr.func.id
|
|
expected_type = ctypes_to_ir(call_type)
|
|
|
|
if val[1] != expected_type:
|
|
# NOTE: We are only considering casting to and from int types for now
|
|
if isinstance(val[1], ir.IntType) and isinstance(expected_type, ir.IntType):
|
|
if val[1].width < expected_type.width:
|
|
val = (builder.sext(val[0], expected_type), expected_type)
|
|
else:
|
|
val = (builder.trunc(val[0], expected_type), expected_type)
|
|
else:
|
|
raise ValueError(f"Type mismatch: expected {expected_type}, got {val[1]}")
|
|
return val
|
|
|
|
|
|
def _handle_compare(
|
|
func, module, builder, cond, local_sym_tab, map_sym_tab, structs_sym_tab=None
|
|
):
|
|
"""Handle ast.Compare expressions."""
|
|
|
|
if len(cond.ops) != 1 or len(cond.comparators) != 1:
|
|
logger.error("Only single comparisons are supported")
|
|
return None
|
|
lhs = eval_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
cond.left,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
rhs = eval_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
cond.comparators[0],
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
|
|
if lhs is None or rhs is None:
|
|
logger.error("Failed to evaluate comparison operands")
|
|
return None
|
|
|
|
lhs, _ = lhs
|
|
rhs, _ = rhs
|
|
return handle_comparator(func, builder, cond.ops[0], lhs, rhs)
|
|
|
|
|
|
def _handle_unary_op(
|
|
func,
|
|
module,
|
|
builder,
|
|
expr: ast.UnaryOp,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab=None,
|
|
):
|
|
"""Handle ast.UnaryOp expressions."""
|
|
if not isinstance(expr.op, ast.Not):
|
|
logger.error("Only 'not' unary operator is supported")
|
|
return None
|
|
|
|
operand = eval_expr(
|
|
func, module, builder, expr.operand, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
if operand is None:
|
|
logger.error("Failed to evaluate operand for unary operation")
|
|
return None
|
|
|
|
operand_val, operand_type = operand
|
|
true_const = ir.Constant(ir.IntType(1), 1)
|
|
result = builder.xor(convert_to_bool(builder, operand_val), true_const)
|
|
return result, ir.IntType(1)
|
|
|
|
|
|
def _handle_and_op(func, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab):
|
|
"""Handle `and` boolean operations."""
|
|
|
|
logger.debug(f"Handling 'and' operator with {len(expr.values)} operands")
|
|
|
|
merge_block = func.append_basic_block(name="and.merge")
|
|
false_block = func.append_basic_block(name="and.false")
|
|
|
|
incoming_values = []
|
|
|
|
for i, value in enumerate(expr.values):
|
|
is_last = i == len(expr.values) - 1
|
|
|
|
# Evaluate current operand
|
|
operand_result = eval_expr(
|
|
func, None, builder, value, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
if operand_result is None:
|
|
logger.error(f"Failed to evaluate operand {i} in 'and' expression")
|
|
return None
|
|
|
|
operand_val, operand_type = operand_result
|
|
|
|
# Convert to boolean if needed
|
|
operand_bool = convert_to_bool(builder, operand_val)
|
|
current_block = builder.block
|
|
|
|
if is_last:
|
|
# Last operand: result is this value
|
|
builder.branch(merge_block)
|
|
incoming_values.append((operand_bool, current_block))
|
|
else:
|
|
# Not last: check if true, continue or short-circuit
|
|
next_check = func.append_basic_block(name=f"and.check_{i + 1}")
|
|
builder.cbranch(operand_bool, next_check, false_block)
|
|
builder.position_at_end(next_check)
|
|
|
|
# False block: short-circuit with false
|
|
builder.position_at_end(false_block)
|
|
builder.branch(merge_block)
|
|
false_value = ir.Constant(ir.IntType(1), 0)
|
|
incoming_values.append((false_value, false_block))
|
|
|
|
# Merge block: phi node
|
|
builder.position_at_end(merge_block)
|
|
phi = builder.phi(ir.IntType(1), name="and.result")
|
|
for val, block in incoming_values:
|
|
phi.add_incoming(val, block)
|
|
|
|
logger.debug(f"Generated 'and' with {len(incoming_values)} incoming values")
|
|
return phi, ir.IntType(1)
|
|
|
|
|
|
def _handle_or_op(func, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab):
|
|
"""Handle `or` boolean operations."""
|
|
|
|
logger.debug(f"Handling 'or' operator with {len(expr.values)} operands")
|
|
|
|
merge_block = func.append_basic_block(name="or.merge")
|
|
true_block = func.append_basic_block(name="or.true")
|
|
|
|
incoming_values = []
|
|
|
|
for i, value in enumerate(expr.values):
|
|
is_last = i == len(expr.values) - 1
|
|
|
|
# Evaluate current operand
|
|
operand_result = eval_expr(
|
|
func, None, builder, value, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
if operand_result is None:
|
|
logger.error(f"Failed to evaluate operand {i} in 'or' expression")
|
|
return None
|
|
|
|
operand_val, operand_type = operand_result
|
|
|
|
# Convert to boolean if needed
|
|
operand_bool = convert_to_bool(builder, operand_val)
|
|
current_block = builder.block
|
|
|
|
if is_last:
|
|
# Last operand: result is this value
|
|
builder.branch(merge_block)
|
|
incoming_values.append((operand_bool, current_block))
|
|
else:
|
|
# Not last: check if false, continue or short-circuit
|
|
next_check = func.append_basic_block(name=f"or.check_{i + 1}")
|
|
builder.cbranch(operand_bool, true_block, next_check)
|
|
builder.position_at_end(next_check)
|
|
|
|
# True block: short-circuit with true
|
|
builder.position_at_end(true_block)
|
|
builder.branch(merge_block)
|
|
true_value = ir.Constant(ir.IntType(1), 1)
|
|
incoming_values.append((true_value, true_block))
|
|
|
|
# Merge block: phi node
|
|
builder.position_at_end(merge_block)
|
|
phi = builder.phi(ir.IntType(1), name="or.result")
|
|
for val, block in incoming_values:
|
|
phi.add_incoming(val, block)
|
|
|
|
logger.debug(f"Generated 'or' with {len(incoming_values)} incoming values")
|
|
return phi, ir.IntType(1)
|
|
|
|
|
|
def _handle_boolean_op(
|
|
func,
|
|
module,
|
|
builder,
|
|
expr: ast.BoolOp,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab=None,
|
|
):
|
|
"""Handle `and` and `or` boolean operations."""
|
|
|
|
if isinstance(expr.op, ast.And):
|
|
return _handle_and_op(
|
|
func, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
elif isinstance(expr.op, ast.Or):
|
|
return _handle_or_op(
|
|
func, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
else:
|
|
logger.error(f"Unsupported boolean operator: {type(expr.op).__name__}")
|
|
return None
|
|
|
|
|
|
def eval_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
expr,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab=None,
|
|
):
|
|
"""
|
|
Evaluate an expression and return its LLVM IR value and type.
|
|
|
|
Args:
|
|
func: The LLVM IR function being built
|
|
module: The LLVM IR module
|
|
builder: LLVM IR builder
|
|
expr: The AST expression node to evaluate
|
|
local_sym_tab: Local symbol table
|
|
map_sym_tab: Map symbol table
|
|
structs_sym_tab: Struct symbol table
|
|
|
|
Returns:
|
|
A tuple of (value, type) or None if evaluation fails
|
|
"""
|
|
logger.info(f"Evaluating expression: {ast.dump(expr)}")
|
|
if isinstance(expr, ast.Name):
|
|
return _handle_name_expr(expr, local_sym_tab, builder)
|
|
elif isinstance(expr, ast.Constant):
|
|
return _handle_constant_expr(expr)
|
|
elif isinstance(expr, ast.Call):
|
|
if isinstance(expr.func, ast.Name) and expr.func.id == "deref":
|
|
return _handle_deref_call(expr, local_sym_tab, builder)
|
|
|
|
if isinstance(expr.func, ast.Name) and is_ctypes(expr.func.id):
|
|
return _handle_ctypes_call(
|
|
func,
|
|
module,
|
|
builder,
|
|
expr,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
|
|
# delayed import to avoid circular dependency
|
|
from pythonbpf.helper import HelperHandlerRegistry, handle_helper_call
|
|
|
|
if isinstance(expr.func, ast.Name) and HelperHandlerRegistry.has_handler(
|
|
expr.func.id
|
|
):
|
|
return handle_helper_call(
|
|
expr,
|
|
module,
|
|
builder,
|
|
func,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
elif isinstance(expr.func, ast.Attribute):
|
|
logger.info(f"Handling method call: {ast.dump(expr.func)}")
|
|
if isinstance(expr.func.value, ast.Call) and isinstance(
|
|
expr.func.value.func, ast.Name
|
|
):
|
|
method_name = expr.func.attr
|
|
if HelperHandlerRegistry.has_handler(method_name):
|
|
return handle_helper_call(
|
|
expr,
|
|
module,
|
|
builder,
|
|
func,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
elif isinstance(expr.func.value, ast.Name):
|
|
obj_name = expr.func.value.id
|
|
method_name = expr.func.attr
|
|
if obj_name in map_sym_tab:
|
|
if HelperHandlerRegistry.has_handler(method_name):
|
|
return handle_helper_call(
|
|
expr,
|
|
module,
|
|
builder,
|
|
func,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
elif isinstance(expr, ast.Attribute):
|
|
return _handle_attribute_expr(expr, local_sym_tab, structs_sym_tab, builder)
|
|
elif isinstance(expr, ast.BinOp):
|
|
from pythonbpf.binary_ops import handle_binary_op
|
|
|
|
return handle_binary_op(expr, builder, None, local_sym_tab)
|
|
elif isinstance(expr, ast.Compare):
|
|
return _handle_compare(
|
|
func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
elif isinstance(expr, ast.UnaryOp):
|
|
return _handle_unary_op(
|
|
func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
elif isinstance(expr, ast.BoolOp):
|
|
return _handle_boolean_op(
|
|
func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
logger.info("Unsupported expression evaluation")
|
|
return None
|
|
|
|
|
|
def handle_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
expr,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
):
|
|
"""Handle expression statements in the function body."""
|
|
logger.info(f"Handling expression: {ast.dump(expr)}")
|
|
call = expr.value
|
|
if isinstance(call, ast.Call):
|
|
eval_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
call,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
else:
|
|
logger.info("Unsupported expression type")
|