Files
python-bpf/pythonbpf/binary_ops.py
Pragyansh Chaturvedi 489244a015 Add store_through_chain
2025-10-10 02:56:11 +05:30

92 lines
3.0 KiB
Python

import ast
from llvmlite import ir
from logging import Logger
import logging
logger: Logger = logging.getLogger(__name__)
def deref_to_val(var, builder):
"""Dereference a variable to get its value and pointer chain."""
logger.info(f"Dereferencing {var}, type is {var.type}")
chain = [var]
cur = var
while isinstance(cur.type, ir.PointerType):
cur = builder.load(cur)
chain.append(cur)
if isinstance(cur.type, ir.IntType):
logger.info(f"dereference chain: {chain}")
return cur, chain
else:
raise TypeError(f"Unsupported type for dereferencing: {cur.type}")
def get_operand_value(operand, builder, local_sym_tab):
"""Extract the value from an operand, handling variables and constants."""
if isinstance(operand, ast.Name):
if operand.id in local_sym_tab:
var = local_sym_tab[operand.id].var
val, chain = deref_to_val(var, builder)
return val, chain, var
raise ValueError(f"Undefined variable: {operand.id}")
elif isinstance(operand, ast.Constant):
if isinstance(operand.value, int):
cst = ir.Constant(ir.IntType(64), operand.value)
return cst, [cst], None
raise TypeError(f"Unsupported constant type: {type(operand.value)}")
elif isinstance(operand, ast.BinOp):
res = handle_binary_op_impl(operand, builder, local_sym_tab)
return res, [res], None
raise TypeError(f"Unsupported operand type: {type(operand)}")
def store_through_chain(value, chain, builder):
"""Store a value through a pointer chain."""
if not chain or len(chain) < 2:
raise ValueError("Pointer chain must have at least two elements")
for ptr in reversed(chain[1:]):
builder.store(value, ptr)
value = ptr
def handle_binary_op_impl(rval, builder, local_sym_tab):
op = rval.op
left, _, _ = get_operand_value(rval.left, builder, local_sym_tab)
right, _, _ = get_operand_value(rval.right, builder, local_sym_tab)
logger.info(f"left is {left}, right is {right}, op is {op}")
# Map AST operation nodes to LLVM IR builder methods
op_map = {
ast.Add: builder.add,
ast.Sub: builder.sub,
ast.Mult: builder.mul,
ast.Div: builder.sdiv,
ast.Mod: builder.srem,
ast.LShift: builder.shl,
ast.RShift: builder.lshr,
ast.BitOr: builder.or_,
ast.BitXor: builder.xor,
ast.BitAnd: builder.and_,
ast.FloorDiv: builder.udiv,
}
if type(op) in op_map:
result = op_map[type(op)](left, right)
return result
else:
raise SyntaxError("Unsupported binary operation")
def handle_binary_op(rval, builder, var_name, local_sym_tab):
result = handle_binary_op_impl(rval, builder, local_sym_tab)
if var_name and var_name in local_sym_tab:
logger.info(
f"Storing result {result} into variable {local_sym_tab[var_name].var}"
)
builder.store(result, local_sym_tab[var_name].var)
return result, result.type