python-bpf/pythonbpf/helper/helper_utils.py

import ast
import logging
from collections.abc import Callable

from llvmlite import ir
from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
from pythonbpf.binary_ops import get_operand_value

logger = logging.getLogger(__name__)


class HelperHandlerRegistry:
    """Registry for BPF helpers"""

    _handlers: dict[str, Callable] = {}

    @classmethod
    def register(cls, helper_name):
        """Decorator to register a handler function for a helper"""

        def decorator(func):
            cls._handlers[helper_name] = func
            return func

        return decorator

    @classmethod
    def get_handler(cls, helper_name):
        """Get the handler function for a helper"""
        return cls._handlers.get(helper_name)

    @classmethod
    def has_handler(cls, helper_name):
        """Check if a handler function is registered for a helper"""
        return helper_name in cls._handlers


class ScratchPoolManager:
    """Manage the temporary helper variables in local_sym_tab"""

    def __init__(self):
        self._counter = 0

    @property
    def counter(self):
        return self._counter

    def reset(self):
        self._counter = 0
        logger.debug("Scratch pool counter reset to 0")

    def get_next_temp(self, local_sym_tab):
        temp_name = f"__helper_temp_{self._counter}"
        self._counter += 1

        if temp_name not in local_sym_tab:
            raise ValueError(
                f"Scratch pool exhausted or inadequate: {temp_name}. "
                f"Current counter: {self._counter}"
            )

        return local_sym_tab[temp_name].var, temp_name


_temp_pool_manager = ScratchPoolManager()  # Singleton instance


def reset_scratch_pool():
    """Reset the scratch pool counter"""
    _temp_pool_manager.reset()


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].var
    raise ValueError(f"Variable '{var_name}' not found in local symbol table")


def create_int_constant_ptr(value, builder, local_sym_tab, int_width=64):
    """Create a pointer to an integer constant."""

    # Default to 64-bit integer
    ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab)
    logger.info(f"Using temp variable '{temp_name}' for int constant {value}")
    const_val = ir.Constant(ir.IntType(int_width), value)
    builder.store(const_val, ptr)
    return ptr


def get_or_create_ptr_from_arg(
    func, module, arg, builder, local_sym_tab, map_sym_tab, struct_sym_tab=None
):
    """Extract or create pointer from the call arguments."""

    if isinstance(arg, ast.Name):
        ptr = get_var_ptr_from_name(arg.id, local_sym_tab)
    elif isinstance(arg, ast.Constant) and isinstance(arg.value, int):
        ptr = create_int_constant_ptr(arg.value, builder, local_sym_tab)
    else:
        # Evaluate the expression and store the result in a temp variable
        val = get_operand_value(
            func, module, arg, builder, local_sym_tab, map_sym_tab, struct_sym_tab
        )
        if val is None:
            raise ValueError("Failed to evaluate expression for helper arg.")

        # NOTE: We assume the result is an int64 for now
        # if isinstance(arg, ast.Attribute):
        # return val
        ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab)
        logger.info(f"Using temp variable '{temp_name}' for expression result")
        builder.store(val, ptr)

    return ptr


def get_flags_val(arg, builder, local_sym_tab):
    """Extract or create flags value from the call arguments."""
    if not arg:
        return 0

    if isinstance(arg, ast.Name):
        if local_sym_tab and arg.id in local_sym_tab:
            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")
    elif isinstance(arg, ast.Constant) and isinstance(arg.value, int):
        return arg.value

    raise NotImplementedError(
        "Only var names or int consts are supported as map helpers flags."
    )


def simple_string_print(string_value, module, builder, func):
    """Prepare arguments for bpf_printk from a simple string value"""
    fmt_str = string_value + "\n\0"
    fmt_ptr = _create_format_string_global(fmt_str, func, module, builder)

    args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]
    return args


def handle_fstring_print(
    joined_str,
    module,
    builder,
    func,
    local_sym_tab=None,
    struct_sym_tab=None,
):
    """Handle f-string formatting for bpf_printk emitter."""
    fmt_parts = []
    exprs = []

    for value in joined_str.values:
        logger.debug(f"Processing f-string value: {ast.dump(value)}")

        if isinstance(value, ast.Constant):
            _process_constant_in_fstring(value, fmt_parts, exprs)
        elif isinstance(value, ast.FormattedValue):
            _process_fval(
                value,
                fmt_parts,
                exprs,
                local_sym_tab,
                struct_sym_tab,
            )
        else:
            raise NotImplementedError(f"Unsupported f-string value type: {type(value)}")

    fmt_str = "".join(fmt_parts)
    args = simple_string_print(fmt_str, module, builder, func)

    # NOTE: Process expressions (limited to 3 due to BPF constraints)
    if len(exprs) > 3:
        logger.warning("bpf_printk supports up to 3 args, extra args will be ignored.")

    for expr in exprs[:3]:
        arg_value = _prepare_expr_args(
            expr,
            func,
            module,
            builder,
            local_sym_tab,
            struct_sym_tab,
        )
        args.append(arg_value)

    return args


def _process_constant_in_fstring(cst, fmt_parts, exprs):
    """Process constant values in f-string."""
    if isinstance(cst.value, str):
        fmt_parts.append(cst.value)
    elif isinstance(cst.value, int):
        fmt_parts.append("%lld")
        exprs.append(ir.Constant(ir.IntType(64), cst.value))
    else:
        raise NotImplementedError(
            f"Unsupported constant type in f-string: {type(cst.value)}"
        )


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)}")

    if isinstance(fval.value, ast.Name):
        _process_name_in_fval(fval.value, fmt_parts, exprs, local_sym_tab)
    elif isinstance(fval.value, ast.Attribute):
        _process_attr_in_fval(
            fval.value,
            fmt_parts,
            exprs,
            local_sym_tab,
            struct_sym_tab,
        )
    else:
        raise NotImplementedError(
            f"Unsupported formatted value in f-string: {type(fval.value)}"
        )


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, 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):
    """Process attribute nodes in formatted values."""
    if (
        isinstance(attr_node.value, ast.Name)
        and local_sym_tab
        and attr_node.value.id in local_sym_tab
    ):
        var_name = attr_node.value.id
        field_name = attr_node.attr

        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"
            )

        struct_info = struct_sym_tab[var_type]
        if field_name not in struct_info.fields:
            raise ValueError(f"Field '{field_name}' not found in struct '{var_type}'")

        field_type = struct_info.field_type(field_name)
        _populate_fval(field_type, attr_node, fmt_parts, exprs)
    else:
        raise NotImplementedError(
            "Only simple attribute on local vars is supported in f-strings."
        )


def _populate_fval(ftype, node, fmt_parts, exprs):
    """Populate format parts and expressions based on field type."""
    if isinstance(ftype, ir.IntType):
        # TODO: We print as signed integers only for now
        if ftype.width == 64:
            fmt_parts.append("%lld")
            exprs.append(node)
        elif ftype.width == 32:
            fmt_parts.append("%d")
            exprs.append(node)
        else:
            raise NotImplementedError(
                f"Unsupported integer width in f-string: {ftype.width}"
            )
    elif isinstance(ftype, ir.PointerType):
        target, depth = get_base_type_and_depth(ftype)
        if isinstance(target, ir.IntType):
            if target.width == 64:
                fmt_parts.append("%lld")
                exprs.append(node)
            elif target.width == 32:
                fmt_parts.append("%d")
                exprs.append(node)
            elif target.width == 8 and depth == 1:
                # NOTE: Assume i8* is a string
                fmt_parts.append("%s")
                exprs.append(node)
            else:
                raise NotImplementedError(
                    f"Unsupported pointer target type in f-string: {target}"
                )
        else:
            raise NotImplementedError(
                f"Unsupported pointer target type in f-string: {target}"
            )
    else:
        raise NotImplementedError(f"Unsupported field type in f-string: {ftype}")


def _create_format_string_global(fmt_str, func, module, builder):
    """Create a global variable for the format string."""
    fmt_name = f"{func.name}____fmt{func._fmt_counter}"
    func._fmt_counter += 1

    fmt_gvar = ir.GlobalVariable(
        module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name
    )
    fmt_gvar.global_constant = True
    fmt_gvar.initializer = ir.Constant(
        ir.ArrayType(ir.IntType(8), len(fmt_str)), bytearray(fmt_str.encode("utf8"))
    )
    fmt_gvar.linkage = "internal"
    fmt_gvar.align = 1

    return builder.bitcast(fmt_gvar, ir.PointerType())


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,
        module,
        builder,
        expr,
        local_sym_tab,
        None,
        struct_sym_tab,
    )

    if val:
        if isinstance(val.type, ir.PointerType):
            target, depth = get_base_type_and_depth(val.type)
            if isinstance(target, ir.IntType):
                if target.width >= 32:
                    val = deref_to_depth(func, builder, val, depth)
                    val = builder.sext(val, ir.IntType(64))
                elif target.width == 8 and depth == 1:
                    # NOTE: i8* is string, no need to deref
                    pass

            else:
                logger.warning(
                    "Only int and ptr supported in bpf_printk args. Others default to 0."
                )
                val = ir.Constant(ir.IntType(64), 0)
        elif isinstance(val.type, ir.IntType):
            if val.type.width < 64:
                val = builder.sext(val, ir.IntType(64))
        else:
            logger.warning(
                "Only int and ptr supported in bpf_printk args. Others default to 0."
            )
            val = ir.Constant(ir.IntType(64), 0)
        return val
    else:
        logger.warning(
            "Failed to evaluate expression for bpf_printk argument. "
            "It will be converted to 0."
        )
        return ir.Constant(ir.IntType(64), 0)


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].var
        else:
            raise ValueError(
                f"Data variable {data_name} not found in local symbol table."
            )

        # Check if data_name is a struct
        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"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."
        )