seperate expr handling logic to a different file to prevent circular import, add format strings

pythonbpf/bpf_helper_handler.py

@@ -1,8 +1,9 @@
 import ast
 from llvmlite import ir
+from .expr_pass import eval_expr
 
 
-def bpf_ktime_get_ns_emitter(call, module, builder, func):
+def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
     """
     Emit LLVM IR for bpf_ktime_get_ns helper function call.
     """
@@ -62,10 +63,87 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     return result
 
 
-def bpf_printk_emitter(call, module, builder, func):
+def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
     if not hasattr(func, "_fmt_counter"):
         func._fmt_counter = 0
 
+    if not call.args:
+        raise ValueError("print expects at least one argument")
+
+    if isinstance(call.args[0], ast.JoinedStr):
+        fmt_parts = []
+        exprs = []
+
+        for value in call.args[0].values:
+            if isinstance(value, ast.Constant):
+                if isinstance(value.value, str):
+                    fmt_parts.append(value.value)
+                elif isinstance(value.value, int):
+                    fmt_parts.append("%lld")
+                    exprs.append(ir.Constant(ir.IntType(64), value.value))
+                else:
+                    raise NotImplementedError(
+                        "Only string and integer constants are supported in f-string.")
+            elif isinstance(value, ast.FormattedValue):
+                # Assume int for now
+                fmt_parts.append("%d")
+                if isinstance(value.value, ast.Name):
+                    exprs.append(value.value)
+                else:
+                    raise NotImplementedError(
+                        "Only simple variable names are supported in formatted values.")
+            else:
+                raise NotImplementedError(
+                    "Unsupported value type in f-string.")
+
+        fmt_str = "".join(fmt_parts) + "\n" + "\0"
+        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
+
+        fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType())
+
+        args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]
+
+        # Only 3 args supported in bpf_printk
+        if len(exprs) > 3:
+            print(
+                "Warning: bpf_printk supports up to 3 arguments, extra arguments will be ignored.")
+
+        for expr in exprs[:3]:
+            val = eval_expr(func, module, builder, expr, local_sym_tab, None)
+            if val:
+                if isinstance(val.type, ir.PointerType):
+                    val = builder.ptrtoint(val, ir.IntType(64))
+                elif isinstance(val.type, ir.IntType):
+                    if val.type.width < 64:
+                        val = builder.sext(val, ir.IntType(64))
+                else:
+                    print(
+                        "Warning: Only integer and pointer types are supported in bpf_printk arguments. Others will be converted to 0.")
+                    val = ir.Constant(ir.IntType(64), 0)
+                args.append(val)
+            else:
+                print(
+                    "Warning: Failed to evaluate expression for bpf_printk argument. It will be converted to 0.")
+                args.append(ir.Constant(ir.IntType(64), 0))
+
+        fn_type = ir.FunctionType(ir.IntType(
+            64), [ir.PointerType(), ir.IntType(32)], var_arg=True)
+        fn_ptr_type = ir.PointerType(fn_type)
+        fn_addr = ir.Constant(ir.IntType(64), 6)
+        fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+        return builder.call(fn_ptr, args, tail=True)
+
     for arg in call.args:
         if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
             fmt_str = arg.value + "\n" + "\0"
@@ -93,6 +171,7 @@ def bpf_printk_emitter(call, module, builder, func):
             builder.call(fn_ptr, [fmt_ptr, ir.Constant(
                 ir.IntType(32), len(fmt_str))], tail=True)
 
+
 def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
     """
     Emit LLVM IR for bpf_map_update_elem helper function call.
@@ -101,11 +180,11 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     if not call.args or len(call.args) < 2 or len(call.args) > 3:
         raise ValueError("Map update expects 2 or 3 arguments (key, value, flags), got "
                          f"{len(call.args)}")
-    
+
     key_arg = call.args[0]
     value_arg = call.args[1]
     flags_arg = call.args[2] if len(call.args) > 2 else None
-    
+
     # Handle key
     if isinstance(key_arg, ast.Name):
         key_name = key_arg.id
@@ -124,7 +203,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     else:
         raise NotImplementedError(
             "Only simple variable names and integer constants are supported as keys in map update.")
-    
+
     # Handle value
     if isinstance(value_arg, ast.Name):
         value_name = value_arg.id
@@ -143,7 +222,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     else:
         raise NotImplementedError(
             "Only simple variable names and integer constants are supported as values in map update.")
-    
+
     # Handle flags argument (defaults to 0)
     if flags_arg is not None:
         if isinstance(flags_arg, ast.Constant) and isinstance(flags_arg.value, int):
@@ -162,7 +241,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
                 "Only integer constants and simple variable names are supported as flags in map update.")
     else:
         flags_val = 0
-    
+
     if key_ptr is None or value_ptr is None:
         raise ValueError("Key pointer or value pointer is None.")
 
@@ -173,20 +252,22 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
         var_arg=False
     )
     fn_ptr_type = ir.PointerType(fn_type)
-    
+
     # helper id
     fn_addr = ir.Constant(ir.IntType(64), 2)
     fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
-    
+
     if isinstance(flags_val, int):
         flags_const = ir.Constant(ir.IntType(64), flags_val)
     else:
         flags_const = flags_val
 
-    result = builder.call(fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False)
-    
+    result = builder.call(
+        fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False)
+
     return result
 
+
 helper_func_list = {
     "lookup": bpf_map_lookup_elem_emitter,
     "print": bpf_printk_emitter,
@@ -200,7 +281,7 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
         func_name = call.func.id
         if func_name in helper_func_list:
             # it is not a map method call
-            return helper_func_list[func_name](call, module, builder, func)
+            return helper_func_list[func_name](call, None, module, builder, func, local_sym_tab)
         else:
             raise NotImplementedError(
                 f"Function {func_name} is not implemented as a helper function.")