Fix local_sym_tab accesses in expr_pass

Signed-off-by: varun-r-mallya <varunrmallya@gmail.com>

demo/pybpf4.py

@@ -0,0 +1,56 @@
+import time
+
+from pythonbpf import bpf, map, section, bpfglobal, BPF
+from pythonbpf.helpers import pid
+from pythonbpf.maps import HashMap
+from pylibbpf import *
+from ctypes import c_void_p, c_int64, c_uint64, c_int32
+import matplotlib.pyplot as plt
+
+# This program attaches an eBPF tracepoint to sys_enter_clone,
+# counts per-PID clone syscalls, stores them in a hash map,
+# and then plots the distribution as a histogram using matplotlib.
+# It provides a quick view of process creation activity over 10 seconds.
+# Everything is done with Python only code and with the new pylibbpf library.
+# Run `sudo /path/to/python/binary/ pybpf4.py`
+
+@bpf
+@map
+def hist() -> HashMap:
+    return HashMap(key=c_int32, value=c_uint64, max_entries=4096)
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_clone")
+def hello(ctx: c_void_p) -> c_int64:
+    process_id = pid()
+    one = 1
+    prev = hist().lookup(process_id)
+    if prev:
+        previous_value = prev + 1
+        print(f"count: {previous_value} with {process_id}")
+        hist().update(process_id, previous_value)
+        return c_int64(0)
+    else:
+        hist().update(process_id, one)
+    return c_int64(0)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+b = BPF()
+b.load_and_attach()
+hist = BpfMap(b, hist)
+print("Recording")
+time.sleep(10)
+
+counts = list(hist.values())
+
+plt.hist(counts, bins=20)
+plt.xlabel("Clone calls per PID")
+plt.ylabel("Frequency")
+plt.title("Syscall clone counts")
+plt.show()

examples/c-form/ex7.bpf.c

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+struct trace_entry {
+  short unsigned int type;
+  unsigned char flags;
+  unsigned char preempt_count;
+  int pid;
+};
+
+struct trace_event_raw_sys_enter {
+  struct trace_entry ent;
+  long int id;
+  long unsigned int args[6];
+  char __data[0];
+};
+
+struct event {
+  __u32 pid;
+  __u32 uid;
+  __u64 ts;
+};
+
+struct {
+  __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
+  __uint(key_size, sizeof(int));
+  __uint(value_size, sizeof(int));
+} events SEC(".maps");
+
+SEC("tp/syscalls/sys_enter_setuid")
+int handle_setuid_entry(struct trace_event_raw_sys_enter *ctx) {
+  struct event data = {};
+
+  // Extract UID from the syscall arguments
+  data.uid = (unsigned int)ctx->args[0];
+  data.ts = bpf_ktime_get_ns();
+  data.pid = bpf_get_current_pid_tgid() >> 32;
+
+  bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, &data, sizeof(data));
+
+  return 0;
+}
+
+char LICENSE[] SEC("license") = "GPL";

examples/execve5.py

@@ -21,12 +21,14 @@ def events() -> PerfEventArray:
 @bpf
 @section("tracepoint/syscalls/sys_enter_clone")
 def hello(ctx: c_void_p) -> c_int32:
+    strobj = "Hi"
     dataobj = data_t()
     ts = ktime()
     process_id = pid()
     dataobj.pid = process_id
     dataobj.ts = ts
     print(f"clone called at {ts} by pid {process_id}")
+    events.output(dataobj)
     return c_int32(0)
 
 

pyproject.toml

@@ -16,7 +16,8 @@ requires-python = ">=3.8"
 
 dependencies = [
   "llvmlite",
-  "astpretty"
+  "astpretty",
+  "pylibbpf"
 ]
 
 [tool.setuptools.packages.find]

pythonbpf/__init__.py

@@ -1,2 +1,2 @@
 from .decorators import bpf, map, section, bpfglobal, struct
-from .codegen import compile_to_ir, compile
+from .codegen import compile_to_ir, compile, BPF

pythonbpf/binary_ops.py

@@ -15,6 +15,7 @@ def recursive_dereferencer(var, builder):
     else:
         raise TypeError(f"Unsupported type for dereferencing: {var.type}")
 
+
 def handle_binary_op(rval, module, builder, var_name, local_sym_tab, map_sym_tab, func):
     print(module)
     left = rval.left
@@ -24,7 +25,7 @@ def handle_binary_op(rval, module, builder, var_name, local_sym_tab, map_sym_tab
     # Handle left operand
     if isinstance(left, ast.Name):
         if left.id in local_sym_tab:
-            left = recursive_dereferencer(local_sym_tab[left.id], builder)
+            left = recursive_dereferencer(local_sym_tab[left.id][0], builder)
         else:
             raise SyntaxError(f"Undefined variable: {left.id}")
     elif isinstance(left, ast.Constant):
@@ -34,7 +35,7 @@ def handle_binary_op(rval, module, builder, var_name, local_sym_tab, map_sym_tab
 
     if isinstance(right, ast.Name):
         if right.id in local_sym_tab:
-            right = recursive_dereferencer(local_sym_tab[right.id], builder)
+            right = recursive_dereferencer(local_sym_tab[right.id][0], builder)
         else:
             raise SyntaxError(f"Undefined variable: {right.id}")
     elif isinstance(right, ast.Constant):
@@ -46,36 +47,36 @@ def handle_binary_op(rval, module, builder, var_name, local_sym_tab, map_sym_tab
 
     if isinstance(op, ast.Add):
         builder.store(builder.add(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.Sub):
         builder.store(builder.sub(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.Mult):
         builder.store(builder.mul(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.Div):
         builder.store(builder.sdiv(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.Mod):
         builder.store(builder.srem(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.LShift):
         builder.store(builder.shl(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.RShift):
         builder.store(builder.lshr(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.BitOr):
         builder.store(builder.or_(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.BitXor):
         builder.store(builder.xor(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.BitAnd):
         builder.store(builder.and_(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     elif isinstance(op, ast.FloorDiv):
         builder.store(builder.udiv(left, right),
-                      local_sym_tab[var_name])
+                      local_sym_tab[var_name][0])
     else:
         raise SyntaxError("Unsupported binary operation")

pythonbpf/bpf_helper_handler.py

@@ -3,7 +3,7 @@ from llvmlite import ir
 from .expr_pass import eval_expr
 
 
-def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
+def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_ktime_get_ns helper function call.
     """
@@ -16,7 +16,7 @@ def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab
     return result
 
 
-def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
+def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_map_lookup_elem helper function call.
     """
@@ -27,7 +27,7 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     if isinstance(key_arg, ast.Name):
         key_name = key_arg.id
         if local_sym_tab and key_name in local_sym_tab:
-            key_ptr = local_sym_tab[key_name]
+            key_ptr = local_sym_tab[key_name][0]
         else:
             raise ValueError(
                 f"Key variable {key_name} not found in local symbol table.")
@@ -63,7 +63,7 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     return result
 
 
-def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
+def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
     if not hasattr(func, "_fmt_counter"):
         func._fmt_counter = 0
 
@@ -85,6 +85,7 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None)
                     raise NotImplementedError(
                         "Only string and integer constants are supported in f-string.")
             elif isinstance(value, ast.FormattedValue):
+                print("Formatted value:", ast.dump(value))
                 # Assume int for now
                 fmt_parts.append("%lld")
                 if isinstance(value.value, ast.Name):
@@ -172,7 +173,7 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None)
                 ir.IntType(32), len(fmt_str))], tail=True)
 
 
-def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
+def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_map_update_elem helper function call.
     Expected call signature: map.update(key, value, flags=0)
@@ -189,7 +190,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     if isinstance(key_arg, ast.Name):
         key_name = key_arg.id
         if local_sym_tab and key_name in local_sym_tab:
-            key_ptr = local_sym_tab[key_name]
+            key_ptr = local_sym_tab[key_name][0]
         else:
             raise ValueError(
                 f"Key variable {key_name} not found in local symbol table.")
@@ -208,7 +209,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     if isinstance(value_arg, ast.Name):
         value_name = value_arg.id
         if local_sym_tab and value_name in local_sym_tab:
-            value_ptr = local_sym_tab[value_name]
+            value_ptr = local_sym_tab[value_name][0]
         else:
             raise ValueError(
                 f"Value variable {value_name} not found in local symbol table.")
@@ -231,7 +232,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
             flags_name = flags_arg.id
             if local_sym_tab and flags_name in local_sym_tab:
                 # Assume it's a stored integer value, load it
-                flags_ptr = local_sym_tab[flags_name]
+                flags_ptr = local_sym_tab[flags_name][0]
                 flags_val = builder.load(flags_ptr)
             else:
                 raise ValueError(
@@ -268,7 +269,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     return result
 
 
-def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
+def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_map_delete_elem helper function call.
     Expected call signature: map.delete(key)
@@ -284,7 +285,7 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     if isinstance(key_arg, ast.Name):
         key_name = key_arg.id
         if local_sym_tab and key_name in local_sym_tab:
-            key_ptr = local_sym_tab[key_name]
+            key_ptr = local_sym_tab[key_name][0]
         else:
             raise ValueError(
                 f"Key variable {key_name} not found in local symbol table.")
@@ -323,7 +324,7 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     return result
 
 
-def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
+def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
     """
@@ -340,6 +341,58 @@ def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local
     return pid
 
 
+def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
+    if len(call.args) != 1:
+        raise ValueError("Perf event output expects exactly one argument (data), got "
+                         f"{len(call.args)}")
+    data_arg = call.args[0]
+    ctx_ptr = func.args[0]  # First argument to the function is ctx
+
+    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][0]
+        else:
+            raise ValueError(
+                f"Data variable {data_name} not found in local symbol table.")
+        # Check is data_name is a struct
+        if local_var_metadata and data_name in local_var_metadata:
+            data_type = local_var_metadata[data_name]
+            if data_type in struct_sym_tab:
+                struct_info = struct_sym_tab[data_type]
+                size_val = ir.Constant(ir.IntType(64), struct_info["size"])
+            else:
+                raise ValueError(
+                    f"Struct type {data_type} for variable {data_name} not found in struct symbol table.")
+        else:
+            raise ValueError(
+                f"Metadata for variable {data_name} not found in local variable metadata.")
+
+        # BPF_F_CURRENT_CPU is -1 in 32 bit
+        flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
+
+        map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
+        data_void_ptr = builder.bitcast(data_ptr, ir.PointerType())
+        fn_type = ir.FunctionType(
+            ir.IntType(64),
+            [ir.PointerType(ir.IntType(8)), ir.PointerType(), ir.IntType(64),
+             ir.PointerType(), ir.IntType(64)],
+            var_arg=False
+        )
+        fn_ptr_type = ir.PointerType(fn_type)
+
+        # helper id
+        fn_addr = ir.Constant(ir.IntType(64), 25)
+        fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+
+        result = builder.call(
+            fn_ptr, [ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], tail=False)
+        return result
+    else:
+        raise NotImplementedError(
+            "Only simple object names are supported as data in perf event output.")
+
+
 helper_func_list = {
     "lookup": bpf_map_lookup_elem_emitter,
     "print": bpf_printk_emitter,
@@ -347,10 +400,11 @@ helper_func_list = {
     "update": bpf_map_update_elem_emitter,
     "delete": bpf_map_delete_elem_emitter,
     "pid": bpf_get_current_pid_tgid_emitter,
+    "output": bpf_perf_event_output_handler,
 }
 
 
-def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None):
+def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     if isinstance(call.func, ast.Name):
         func_name = call.func.id
         if func_name in helper_func_list:
@@ -367,14 +421,29 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
             if map_sym_tab and map_name in map_sym_tab:
                 map_ptr = map_sym_tab[map_name]
                 if method_name in helper_func_list:
+                    print(local_var_metadata)
                     return helper_func_list[method_name](
-                        call, map_ptr, module, builder, local_sym_tab)
+                        call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
                 else:
                     raise NotImplementedError(
                         f"Map method {method_name} is not implemented as a helper function.")
             else:
                 raise ValueError(
                     f"Map variable {map_name} not found in symbol tables.")
+        elif isinstance(call.func.value, ast.Name):
+            obj_name = call.func.value.id
+            method_name = call.func.attr
+            if map_sym_tab and obj_name in map_sym_tab:
+                map_ptr = map_sym_tab[obj_name]
+                if method_name in helper_func_list:
+                    return helper_func_list[method_name](
+                        call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
+                else:
+                    raise NotImplementedError(
+                        f"Map method {method_name} is not implemented as a helper function.")
+            else:
+                raise ValueError(
+                    f"Map variable {obj_name} not found in symbol tables.")
         else:
             raise NotImplementedError(
                 "Attribute not supported for map method calls.")

pythonbpf/codegen.py

@@ -9,6 +9,7 @@ import os
 import subprocess
 import inspect
 from pathlib import Path
+from pylibbpf import BpfProgram
 
 
 def find_bpf_chunks(tree):
@@ -116,3 +117,17 @@ def compile():
     ], check=True)
 
     print(f"Object written to {o_file}, {ll_file} can be removed")
+
+def BPF() -> BpfProgram:
+    caller_frame = inspect.stack()[1]
+    caller_file = Path(caller_frame.filename).resolve()
+    ll_file = Path("/tmp") / caller_file.with_suffix(".ll").name
+    o_file = Path("/tmp") / caller_file.with_suffix(".o").name
+    compile_to_ir(str(caller_file), str(ll_file))
+
+    subprocess.run([
+        "llc", "-march=bpf", "-filetype=obj", "-O2",
+        str(ll_file), "-o", str(o_file)
+    ], check=True)
+    
+    return BpfProgram(str(o_file))

pythonbpf/expr_pass.py

@@ -2,11 +2,11 @@ import ast
 from llvmlite import ir
 
 
-def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
+def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab=None, local_var_metadata=None):
     print(f"Evaluating expression: {expr}")
     if isinstance(expr, ast.Name):
         if expr.id in local_sym_tab:
-            var = local_sym_tab[expr.id]
+            var = local_sym_tab[expr.id][0]
             val = builder.load(var)
             return val
         else:
@@ -37,7 +37,7 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
                     return None
                 if isinstance(arg, ast.Name):
                     if arg.id in local_sym_tab:
-                        arg = local_sym_tab[arg.id]
+                        arg = local_sym_tab[arg.id][0]
                     else:
                         print(f"Undefined variable {arg.id}")
                         return None
@@ -50,22 +50,31 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
             # check for helpers
             if expr.func.id in helper_func_list:
                 return handle_helper_call(
-                    expr, module, builder, func, local_sym_tab, map_sym_tab)
+                    expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
         elif isinstance(expr.func, ast.Attribute):
+            print(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 method_name in helper_func_list:
                     return handle_helper_call(
-                        expr, module, builder, func, local_sym_tab, map_sym_tab)
+                        expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
+            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 method_name in helper_func_list:
+                        return handle_helper_call(
+                            expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
     print("Unsupported expression evaluation")
     return None
 
 
-def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
+def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata):
     """Handle expression statements in the function body."""
     print(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)
+        eval_expr(func, module, builder, call, local_sym_tab,
+                  map_sym_tab, structs_sym_tab, local_var_metadata)
     else:
         print("Unsupported expression type")

pythonbpf/functions_pass.py

@@ -51,13 +51,13 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
 
             if field_name in struct_info["fields"]:
                 field_idx = struct_info["fields"][field_name]
-                struct_ptr = local_sym_tab[var_name]
+                struct_ptr = local_sym_tab[var_name][0]
                 field_ptr = builder.gep(
                     struct_ptr, [ir.Constant(ir.IntType(32), 0),
                                  ir.Constant(ir.IntType(32), field_idx)],
                     inbounds=True)
                 val = eval_expr(func, module, builder, rval,
-                                local_sym_tab, map_sym_tab)
+                                local_sym_tab, map_sym_tab, structs_sym_tab)
                 if val is None:
                     print("Failed to evaluate struct field assignment")
                     return
@@ -68,19 +68,32 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
         if isinstance(rval.value, bool):
             if rval.value:
                 builder.store(ir.Constant(ir.IntType(1), 1),
-                              local_sym_tab[var_name])
+                              local_sym_tab[var_name][0])
             else:
                 builder.store(ir.Constant(ir.IntType(1), 0),
-                              local_sym_tab[var_name])
+                              local_sym_tab[var_name][0])
             print(f"Assigned constant {rval.value} to {var_name}")
         elif isinstance(rval.value, int):
             # Assume c_int64 for now
             # var = builder.alloca(ir.IntType(64), name=var_name)
             # var.align = 8
             builder.store(ir.Constant(ir.IntType(64), rval.value),
-                          local_sym_tab[var_name])
+                          local_sym_tab[var_name][0])
             # local_sym_tab[var_name] = var
             print(f"Assigned constant {rval.value} to {var_name}")
+        elif isinstance(rval.value, str):
+            str_val = rval.value.encode('utf-8') + b'\x00'
+            str_const = ir.Constant(ir.ArrayType(
+                ir.IntType(8), len(str_val)), bytearray(str_val))
+            global_str = ir.GlobalVariable(
+                module, str_const.type, name=f"{var_name}_str")
+            global_str.linkage = 'internal'
+            global_str.global_constant = True
+            global_str.initializer = str_const
+            str_ptr = builder.bitcast(
+                global_str, ir.PointerType(ir.IntType(8)))
+            builder.store(str_ptr, local_sym_tab[var_name][0])
+            print(f"Assigned string constant '{rval.value}' to {var_name}")
         else:
             print("Unsupported constant type")
     elif isinstance(rval, ast.Call):
@@ -92,7 +105,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                 # var = builder.alloca(ir_type, name=var_name)
                 # var.align = ir_type.width // 8
                 builder.store(ir.Constant(
-                    ir_type, rval.args[0].value), local_sym_tab[var_name])
+                    ir_type, rval.args[0].value), local_sym_tab[var_name][0])
                 print(f"Assigned {call_type} constant "
                       f"{rval.args[0].value} to {var_name}")
                 # local_sym_tab[var_name] = var
@@ -100,19 +113,19 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                 # var = builder.alloca(ir.IntType(64), name=var_name)
                 # var.align = 8
                 val = handle_helper_call(
-                    rval, module, builder, None, local_sym_tab, map_sym_tab)
-                builder.store(val, local_sym_tab[var_name])
+                    rval, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
+                builder.store(val, local_sym_tab[var_name][0])
                 # local_sym_tab[var_name] = var
                 print(f"Assigned constant {rval.func.id} to {var_name}")
             elif call_type == "deref" and len(rval.args) == 1:
                 print(f"Handling deref assignment {ast.dump(rval)}")
                 val = eval_expr(func, module, builder, rval,
-                                local_sym_tab, map_sym_tab)
+                                local_sym_tab, map_sym_tab, structs_sym_tab)
                 if val is None:
                     print("Failed to evaluate deref argument")
                     return
                 print(f"Dereferenced value: {val}, storing in {var_name}")
-                builder.store(val, local_sym_tab[var_name])
+                builder.store(val, local_sym_tab[var_name][0])
                 # local_sym_tab[var_name] = var
                 print(f"Dereferenced and assigned to {var_name}")
             elif call_type in structs_sym_tab and len(rval.args) == 0:
@@ -121,7 +134,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                 # var = builder.alloca(ir_type, name=var_name)
                 # Null init
                 builder.store(ir.Constant(ir_type, None),
-                              local_sym_tab[var_name])
+                              local_sym_tab[var_name][0])
                 local_var_metadata[var_name] = call_type
                 print(f"Assigned struct {call_type} to {var_name}")
                 # local_sym_tab[var_name] = var
@@ -139,10 +152,10 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                     map_ptr = map_sym_tab[map_name]
                     if method_name in helper_func_list:
                         val = handle_helper_call(
-                            rval, module, builder, func, local_sym_tab, map_sym_tab)
+                            rval, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
                         # var = builder.alloca(ir.IntType(64), name=var_name)
                         # var.align = 8
-                        builder.store(val, local_sym_tab[var_name])
+                        builder.store(val, local_sym_tab[var_name][0])
                         # local_sym_tab[var_name] = var
             else:
                 print("Unsupported assignment call structure")
@@ -166,7 +179,7 @@ def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
             return None
     elif isinstance(cond, ast.Name):
         if cond.id in local_sym_tab:
-            var = local_sym_tab[cond.id]
+            var = local_sym_tab[cond.id][0]
             val = builder.load(var)
             if val.type != ir.IntType(1):
                 # Convert nonzero values to true, zero to false
@@ -222,7 +235,7 @@ def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
         return None
 
 
-def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
+def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab, structs_sym_tab=None):
     """Handle if statements in the function body."""
     print("Handling if statement")
     start = builder.block.parent
@@ -243,7 +256,7 @@ def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
     builder.position_at_end(then_block)
     for s in stmt.body:
         process_stmt(func, module, builder, s,
-                     local_sym_tab, map_sym_tab, False)
+                     local_sym_tab, map_sym_tab, structs_sym_tab, False)
     if not builder.block.is_terminated:
         builder.branch(merge_block)
 
@@ -251,7 +264,7 @@ def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
         builder.position_at_end(else_block)
         for s in stmt.orelse:
             process_stmt(func, module, builder, s,
-                         local_sym_tab, map_sym_tab, False)
+                         local_sym_tab, map_sym_tab, structs_sym_tab, False)
         if not builder.block.is_terminated:
             builder.branch(merge_block)
 
@@ -261,14 +274,16 @@ def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
 def process_stmt(func, module, builder, stmt, local_sym_tab, map_sym_tab, structs_sym_tab, did_return, ret_type=ir.IntType(64)):
     print(f"Processing statement: {ast.dump(stmt)}")
     if isinstance(stmt, ast.Expr):
-        handle_expr(func, module, builder, stmt, local_sym_tab, map_sym_tab)
+        handle_expr(func, module, builder, stmt, local_sym_tab,
+                    map_sym_tab, structs_sym_tab, local_var_metadata)
     elif isinstance(stmt, ast.Assign):
         handle_assign(func, module, builder, stmt, map_sym_tab,
                       local_sym_tab, structs_sym_tab)
     elif isinstance(stmt, ast.AugAssign):
         raise SyntaxError("Augmented assignment not supported")
     elif isinstance(stmt, ast.If):
-        handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab)
+        handle_if(func, module, builder, stmt, map_sym_tab,
+                  local_sym_tab, structs_sym_tab)
     elif isinstance(stmt, ast.Return):
         if stmt.value is None:
             builder.ret(ir.Constant(ir.IntType(32), 0))
@@ -368,8 +383,14 @@ def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_t
                     var.align = ir_type.width // 8
                     print(
                         f"Pre-allocated variable {var_name} of type c_int64")
+                elif isinstance(rval.value, str):
+                    ir_type = ir.PointerType(ir.IntType(8))
+                    var = builder.alloca(ir_type, name=var_name)
+                    var.align = 8
+                    print(
+                        f"Pre-allocated variable {var_name} of type string")
                 else:
-                    print("Unsupported constant type")
+                    print(f"Unsupported constant type")
                     continue
             elif isinstance(rval, ast.BinOp):
                 # Assume c_int64 for now
@@ -381,7 +402,7 @@ def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_t
             else:
                 print("Unsupported assignment value type")
                 continue
-            local_sym_tab[var_name] = var
+            local_sym_tab[var_name] = (var, ir_type)
     return local_sym_tab
 
 

pythonbpf/maps.py

@@ -30,3 +30,6 @@ class PerfEventArray:
         self.key_type = key_size
         self.value_type = value_size
         self.entries = {}
+
+    def output(self, data):
+        pass  # Placeholder for output method

pythonbpf/structs_pass.py

@@ -31,9 +31,27 @@ def process_bpf_struct(cls_node, module):
             field_names.append(item.target.id)
             field_types.append(ctypes_to_ir(item.annotation.id))
 
+    curr_offset = 0
+    for ftype in field_types:
+        if isinstance(ftype, ir.IntType):
+            fsize = ftype.width // 8
+            alignment = fsize
+        elif isinstance(ftype, ir.PointerType):
+            fsize = 8
+            alignment = 8
+        else:
+            print(f"Unsupported field type in struct {struct_name}")
+            return
+        padding = (alignment - (curr_offset % alignment)) % alignment
+        curr_offset += padding
+        curr_offset += fsize
+    final_padding = (8 - (curr_offset % 8)) % 8
+    total_size = curr_offset + final_padding
+
     struct_type = ir.LiteralStructType(field_types)
     structs_sym_tab[struct_name] = {
         "type": struct_type,
-        "fields": {name: idx for idx, name in enumerate(field_names)}
+        "fields": {name: idx for idx, name in enumerate(field_names)},
+        "size": total_size
     }
     print(f"Created struct {struct_name} with fields {field_names}")