Allow int** pointers to store binops of type int** op int

pythonbpf/assign_pass.py

@@ -37,6 +37,7 @@ def handle_variable_assignment(
         return False
 
     val, val_type = val_result
+    logger.info(f"Evaluated value for {var_name}: {val} of type {val_type}, {var_type}")
     if val_type != var_type:
         if isinstance(val_type, ir.IntType) and isinstance(var_type, ir.IntType):
             # Allow implicit int widening
@@ -46,6 +47,14 @@ def handle_variable_assignment(
             elif val_type.width > var_type.width:
                 val = builder.trunc(val, var_type)
                 logger.info(f"Implicitly truncated int for variable {var_name}")
+        elif isinstance(val_type, ir.IntType) and isinstance(var_type, ir.PointerType):
+            # NOTE: This is assignment to a PTR_TO_MAP_VALUE_OR_NULL
+            logger.info(
+                f"Creating temporary variable for pointer assignment to {var_name}"
+            )
+            var_ptr_tmp = local_sym_tab[f"{var_name}_tmp"].var
+            builder.store(val, var_ptr_tmp)
+            val = var_ptr_tmp
         else:
             logger.error(
                 f"Type mismatch for variable {var_name}: {val_type} vs {var_type}"

pythonbpf/binary_ops.py

@@ -3,34 +3,21 @@ from llvmlite import ir
 from logging import Logger
 import logging
 
+from pythonbpf.expr import get_base_type_and_depth, deref_to_depth
+
 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):
+def get_operand_value(func, 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
+            var_type = var.type
+            base_type, depth = get_base_type_and_depth(var_type)
+            logger.info(f"var is {var}, base_type is {base_type}, depth is {depth}")
+            val = deref_to_depth(func, builder, var, depth)
+            return val, [val], var
         raise ValueError(f"Undefined variable: {operand.id}")
     elif isinstance(operand, ast.Constant):
         if isinstance(operand.value, int):
@@ -38,7 +25,7 @@ def get_operand_value(operand, builder, local_sym_tab):
             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)
+        res = handle_binary_op_impl(func, operand, builder, local_sym_tab)
         return res, [res], None
     raise TypeError(f"Unsupported operand type: {type(operand)}")
 
@@ -53,10 +40,10 @@ def store_through_chain(value, chain, builder):
         value = ptr
 
 
-def handle_binary_op_impl(rval, builder, local_sym_tab):
+def handle_binary_op_impl(func, rval, builder, local_sym_tab):
     op = rval.op
-    left, lchain, _ = get_operand_value(rval.left, builder, local_sym_tab)
-    right, rchain, _ = get_operand_value(rval.right, builder, local_sym_tab)
+    left, lchain, _ = get_operand_value(func, rval.left, builder, local_sym_tab)
+    right, rchain, _ = get_operand_value(func, rval.right, builder, local_sym_tab)
     logger.info(f"left is {left}, right is {right}, op is {op}")
 
     logger.info(f"left chain: {lchain}, right chain: {rchain}")
@@ -83,8 +70,8 @@ def handle_binary_op_impl(rval, builder, local_sym_tab):
         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)
+def handle_binary_op(func, rval, builder, var_name, local_sym_tab):
+    result = handle_binary_op_impl(func, 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}"

pythonbpf/expr/__init__.py

@@ -1,4 +1,10 @@
 from .expr_pass import eval_expr, handle_expr
-from .type_normalization import convert_to_bool
+from .type_normalization import convert_to_bool, get_base_type_and_depth, deref_to_depth
 
-__all__ = ["eval_expr", "handle_expr", "convert_to_bool"]
+__all__ = [
+    "eval_expr",
+    "handle_expr",
+    "convert_to_bool",
+    "get_base_type_and_depth",
+    "deref_to_depth",
+]

pythonbpf/expr/expr_pass.py

@@ -26,7 +26,7 @@ def _handle_constant_expr(expr: ast.Constant):
     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")
+        logger.error(f"Unsupported constant type {ast.dump(expr)}")
         return None
 
 
@@ -402,7 +402,7 @@ def eval_expr(
     elif isinstance(expr, ast.BinOp):
         from pythonbpf.binary_ops import handle_binary_op
 
-        return handle_binary_op(expr, builder, None, local_sym_tab)
+        return handle_binary_op(func, 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

pythonbpf/expr/type_normalization.py

@@ -16,7 +16,7 @@ COMPARISON_OPS = {
 }
 
 
-def _get_base_type_and_depth(ir_type):
+def get_base_type_and_depth(ir_type):
     """Get the base type for pointer types."""
     cur_type = ir_type
     depth = 0
@@ -26,7 +26,7 @@ def _get_base_type_and_depth(ir_type):
     return cur_type, depth
 
 
-def _deref_to_depth(func, builder, val, target_depth):
+def deref_to_depth(func, builder, val, target_depth):
     """Dereference a pointer to a certain depth."""
 
     cur_val = val
@@ -88,13 +88,13 @@ def _normalize_types(func, builder, lhs, rhs):
         logger.error(f"Type mismatch: {lhs.type} vs {rhs.type}")
         return None, None
     else:
-        lhs_base, lhs_depth = _get_base_type_and_depth(lhs.type)
-        rhs_base, rhs_depth = _get_base_type_and_depth(rhs.type)
+        lhs_base, lhs_depth = get_base_type_and_depth(lhs.type)
+        rhs_base, rhs_depth = get_base_type_and_depth(rhs.type)
         if lhs_base == rhs_base:
             if lhs_depth < rhs_depth:
-                rhs = _deref_to_depth(func, builder, rhs, rhs_depth - lhs_depth)
+                rhs = deref_to_depth(func, builder, rhs, rhs_depth - lhs_depth)
             elif rhs_depth < lhs_depth:
-                lhs = _deref_to_depth(func, builder, lhs, lhs_depth - rhs_depth)
+                lhs = deref_to_depth(func, builder, lhs, lhs_depth - rhs_depth)
             return _normalize_types(func, builder, lhs, rhs)
 
 

pythonbpf/functions/functions_pass.py

@@ -386,6 +386,7 @@ def process_stmt(
 def allocate_mem(
     module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
 ):
+    double_alloc = False
     for stmt in body:
         has_metadata = False
         if isinstance(stmt, ast.If):
@@ -460,8 +461,9 @@ def allocate_mem(
                     var = builder.alloca(ir_type, name=var_name)
 
                     # declare an intermediate ptr type for map lookup
-                    ir_type = ir.IntType(64)
-                    var_tmp = builder.alloca(ir_type, name=f"{var_name}_tmp")
+                    tmp_ir_type = ir.IntType(64)
+                    var_tmp = builder.alloca(tmp_ir_type, name=f"{var_name}_tmp")
+                    double_alloc = True
                     # var.align = ir_type.width // 8
                     logger.info(
                         f"Pre-allocated variable {var_name} and {var_name}_tmp for map"
@@ -504,8 +506,8 @@ def allocate_mem(
             else:
                 local_sym_tab[var_name] = LocalSymbol(var, ir_type)
 
-            if var_tmp:
-                local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, ir_type)
+            if double_alloc:
+                local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, tmp_ir_type)
     return local_sym_tab
 
 

pythonbpf/helper/helper_utils.py

@@ -3,7 +3,7 @@ import logging
 from collections.abc import Callable
 
 from llvmlite import ir
-from pythonbpf.expr import eval_expr
+from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
 
 logger = logging.getLogger(__name__)
 
@@ -224,10 +224,27 @@ def _populate_fval(ftype, node, fmt_parts, exprs):
             raise NotImplementedError(
                 f"Unsupported integer width in f-string: {ftype.width}"
             )
-    elif ftype == ir.PointerType(ir.IntType(8)):
-        # NOTE: We assume i8* is a string
-        fmt_parts.append("%s")
-        exprs.append(node)
+    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}")
 
@@ -264,7 +281,20 @@ def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_ta
 
     if val:
         if isinstance(val.type, ir.PointerType):
-            val = builder.ptrtoint(val, ir.IntType(64))
+            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))