update release for pylibbpf

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

Makefile

@@ -1,6 +1,6 @@
 compile:
 	chmod +x ./tools/compile.py
-	./tools/compile.py ./examples/execve2.py
+	./tools/compile.py ./examples/execve3.py
 
 install: 
 	pip install -e .

README.md

@@ -1,15 +1,76 @@
 # Python-BPF
-This is an LLVM IR generator for eBPF programs in Python. We use `llvmlite` to generate LLVM IR code from pure Python code. This is then compiled to LLVM object files, which can be loaded into the kernel for execution.
+<p align="center">
+<a href="https://www.python.org/downloads/release/python-3080/"><img src="https://img.shields.io/badge/python-3.8-blue.svg"></a>
+<a href="https://pypi.org/project/pythonbpf"><img src="https://badge.fury.io/py/pythonbpf.svg"></a>
+</p>
+
+This is an LLVM IR generator for eBPF programs in Python. We use llvmlite to generate LLVM IR from pure Python. This is then compiled to LLVM object files, which can be loaded into the kernel for execution. We do not rely on BCC to do our compilation.
+
+# DO NOT USE IN PRODUCTION. IN DEVELOPMENT.
+
+## Video Demo
+[Video demo for code under demo/](https://youtu.be/eMyLW8iWbks)
+
+## Slide Deck
+[Slide deck explaining the project](https://docs.google.com/presentation/d/1DsWDIVrpJhM4RgOETO9VWqUtEHo3-c7XIWmNpi6sTSo/edit?usp=sharing)
+
+##  Installation
+- Have `clang` installed.
+- `pip install pythonbpf`
+
+## Usage
+```python
+# pythonbpf_example.py
+from pythonbpf import bpf, map, bpfglobal, section, compile
+from pythonbpf.helpers import bpf_ktime_get_ns
+from pythonbpf.maps import HashMap
+
+from ctypes import c_void_p, c_int64, c_int32, c_uint64
+
+@bpf
+@map
+def last() -> HashMap:
+    return HashMap(key=c_uint64, value=c_uint64, max_entries=1)
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def hello(ctx: c_void_p) -> c_int32:
+    print("entered")
+    return c_int32(0)
+
+@bpf
+@section("tracepoint/syscalls/sys_exit_execve")
+def hello_again(ctx: c_void_p) -> c_int64:
+    print("exited")
+    key = 0
+    tsp = last().lookup(key)
+    print(tsp)
+    ts = bpf_ktime_get_ns()
+    return c_int64(0)
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+def some_normal_function():
+    print("normal function")
+
+# compiles and dumps object file in the same directory
+compile()
+```
+- Run `python pythonbpf_example.py` to get the compiled object file that can be then loaded into the kernel.
 
 ## Development
-Step 0. Make a virtual environment and activate it using `python3 -m venv .venv && source .venv/bin/activate`.  
-Step 1. Run `make install` to install the required dependencies.  
-Step 2. Run `make` to see the compilation output of the example.  
-Step 3. Run `check.sh` to check if generated object file passes through the verifier inside the examples directory.  
-Step 4. Run `make` in the `examples/c-form` directory to modify the example C BPF program to check the actual LLVM IR generated by clang.  
+- Make a virtual environment and activate it using `python3 -m venv .venv && source .venv/bin/activate`.  
+- Run `make install` to install the required dependencies.  
+- Run `make` to see the compilation output of the example.  
+- Run `check.sh` to check if generated object file passes through the verifier inside the examples directory.  
+- Run `make` in the `examples/c-form` directory to modify the example C BPF program to check the actual LLVM IR generated by clang.
 
 ### Development Notes
-Run ` ./check.sh run execve2.o;` in examples folder
+- Run ` ./check.sh check execve2.o;` in examples folder to check if the object code passes the verifier.
+- Run ` ./check.sh run execve2.o;` in examples folder to run the object code using `bpftool`.
 
 ## Authors
 - [@r41k0u](https://github.com/r41k0u)

TODO.md

@@ -0,0 +1,9 @@
+## Short term
+
+- Implement enough functionality to port the BCC tutorial examples in PythonBPF
+
+
+## Long term
+
+- Refactor the codebase to be better than a hackathon project
+- Port to C++ and use actual LLVM?

demo/bcc.py

@@ -0,0 +1,46 @@
+from __future__ import print_function
+from bcc import BPF
+from bcc.utils import printb
+
+# load BPF program
+b = BPF(text="""
+#include <uapi/linux/ptrace.h>
+
+BPF_HASH(last);
+
+int do_trace(struct pt_regs *ctx) {
+    u64 ts, *tsp, delta, key = 0;
+
+    // attempt to read stored timestamp
+    tsp = last.lookup(&key);
+    if (tsp != NULL) {
+        delta = bpf_ktime_get_ns() - *tsp;
+        if (delta < 1000000000) {
+            // output if time is less than 1 second
+            bpf_trace_printk("%d\\n", delta / 1000000);
+        }
+        last.delete(&key);
+    }
+
+    // update stored timestamp
+    ts = bpf_ktime_get_ns();
+    last.update(&key, &ts);
+    return 0;
+}
+""")
+
+b.attach_kprobe(event=b.get_syscall_fnname("sync"), fn_name="do_trace")
+print("Tracing for quick sync's... Ctrl-C to end")
+
+# TODO
+# format output
+start = 0
+while 1:
+    try:
+        (task, pid, cpu, flags, ts, ms) = b.trace_fields()
+        if start == 0:
+            start = ts
+        ts = ts - start
+        printb(b"At time %.2f s: multiple syncs detected, last %s ms ago" % (ts, ms))
+    except KeyboardInterrupt:
+        exit()

demo/pybpf0.py

@@ -0,0 +1,23 @@
+from pythonbpf import bpf, section, bpfglobal, compile
+
+from ctypes import c_void_p, c_int64
+
+# Instructions to how to run this program
+# 1. Install PythonBPF: pip install pythonbpf
+# 2. Run the program: python demo/pybpf0.py
+# 3. Run the program with sudo: sudo examples/check.sh run demo/pybpf0.o
+# 4. Start up any program and watch the output
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def hello_world(ctx: c_void_p) -> c_int64:
+    print("Hello, World!")
+    return c_int64(0)
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+compile()

demo/pybpf1.py

@@ -0,0 +1,41 @@
+from pythonbpf import bpf, map, section, bpfglobal, compile
+from pythonbpf.helpers import XDP_PASS
+from pythonbpf.maps import HashMap
+
+from ctypes import c_void_p, c_int64
+
+# Instructions to how to run this program
+# 1. Install PythonBPF: pip install pythonbpf
+# 2. Run the program: python demo/pybpf1.py
+# 3. Run the program with sudo: sudo examples/check.sh run demo/pybpf1.o
+# 4. Attach object file to any network device with something like ./check.sh xdp ../demo/pybpf1.o tailscale0
+# 5. send traffic through the device and observe effects
+
+@bpf
+@map
+def count() -> HashMap:
+    return HashMap(key=c_int64, value=c_int64, max_entries=1)
+
+
+@bpf
+@section("xdp")
+def hello_world(ctx: c_void_p) -> c_int64:
+    key = 0
+    one = 1
+    prev = count().lookup(key)
+    if prev:
+        prevval = prev + 1
+        print(f"count: {prevval}")
+        count().update(key, prevval)
+        return XDP_PASS
+    else:
+        count().update(key, one)
+
+    return XDP_PASS
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+compile()

demo/pybpf2.py

@@ -0,0 +1,43 @@
+from pythonbpf import bpf, map, section, bpfglobal, compile
+from pythonbpf.helpers import ktime
+from pythonbpf.maps import HashMap
+
+from ctypes import c_void_p, c_int64, c_uint64
+
+# Instructions to how to run this program
+# 1. Install PythonBPF: pip install pythonbpf
+# 2. Run the program: python demo/pybpf2.py
+# 3. Run the program with sudo: sudo examples/check.sh run demo/pybpf2.o
+# 4. Start a Python repl and `import os` and then keep entering `os.sync()` to see reponses.
+
+@bpf
+@map
+def last() -> HashMap:
+    return HashMap(key=c_uint64, value=c_uint64, max_entries=3)
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_sync")
+def do_trace(ctx: c_void_p) -> c_int64:
+    key = 0
+    tsp = last().lookup(key)
+    if tsp:
+        kt = ktime()
+        delta = (kt - tsp)
+        if delta < 1000000000:
+            time_ms = (delta // 1000000)
+            print(f"sync called within last second, last {time_ms} ms ago")
+        last().delete(key)
+    else:
+        kt = ktime()
+        last().update(key, kt)
+    return c_int64(0)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

demo/pybpf3.py

@@ -0,0 +1,52 @@
+from pythonbpf import bpf, map, section, bpfglobal, compile
+from pythonbpf.helpers import ktime
+from pythonbpf.maps import HashMap
+
+from ctypes import c_void_p, c_int64, c_uint64
+
+# Instructions to how to run this program
+# 1. Install PythonBPF: pip install pythonbpf
+# 2. Run the program: python demo/pybpf3.py
+# 3. Run the program with sudo: sudo examples/check.sh run demo/pybpf3.o
+# 4. Start up any program and watch the output
+
+@bpf
+@map
+def last() -> HashMap:
+    return HashMap(key=c_uint64, value=c_uint64, max_entries=3)
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def do_trace(ctx: c_void_p) -> c_int64:
+    key = 0
+    tsp = last().lookup(key)
+    if tsp:
+        kt = ktime()
+        delta = (kt - tsp)
+        if delta < 1000000000:
+            time_ms = (delta // 1000000)
+            print(f"Execve syscall entered within last second, last {time_ms} ms ago")
+        last().delete(key)
+    else:
+        kt = ktime()
+        last().update(key, kt)
+    return c_int64(0)
+
+@bpf
+@section("tracepoint/syscalls/sys_exit_execve")
+def do_exit(ctx: c_void_p) -> c_int64:
+    va = 8
+    nm = 5 ^ va
+    al = 6 & 3
+    ru = (nm + al)
+    print(f"this is a variable {ru}")
+    return c_int64(0)
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

examples/c-form/ex2.bpf.c

@@ -3,17 +3,10 @@
 #define u64 unsigned long long
 #define u32 unsigned int
 
-struct {
-    __uint(type, BPF_MAP_TYPE_HASH);
-    __uint(max_entries, 1);
-    __type(key, u32);
-    __type(value, u64);
-} last SEC(".maps");
-
-SEC("tracepoint/syscalls/sys_enter_execve")
-int hello(struct pt_regs *ctx) {
+SEC("xdp")
+int hello(struct xdp_md *ctx) {
     bpf_printk("Hello, World!\n");
-    return 0;
+    return XDP_PASS;
 }
 
 char LICENSE[] SEC("license") = "GPL";

examples/c-form/ex3-2.bpf.c

@@ -8,7 +8,7 @@ struct {
     __uint(type, BPF_MAP_TYPE_HASH);
     __type(key, u64);
     __type(value, u64);
-    __uint(max_entries, 1);
+    __uint(max_entries, 4);
 } last SEC(".maps");
 
 // Handler for syscall entry

examples/c-form/ex4.bpf.c

@@ -33,7 +33,7 @@ long hello_again(void *ctx) {
         u64 delta = bpf_ktime_get_ns() - *tsp;
         if (delta < 1000000000) {
             // output if time is less than 1 second
-            bpf_trace_printk("%d\\n", delta / 1000000);
+            bpf_printk("execve called within last second");
         }
         bpf_map_delete_elem(&last, &key);
     }

examples/c-form/ex5.bpf.c

@@ -0,0 +1,25 @@
+#define __TARGET_ARCH_arm64
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <bpf/bpf_core_read.h>
+
+// Map: key = struct request*, value = u64 timestamp
+struct {
+    __uint(type, BPF_MAP_TYPE_HASH);
+    __type(key, struct request *);
+    __type(value, u64);
+    __uint(max_entries, 1024);
+} start SEC(".maps");
+
+// Attach to kprobe for blk_start_request
+SEC("kprobe/blk_start_request")
+int BPF_KPROBE(trace_start, struct request *req)
+{
+    u64 ts = bpf_ktime_get_ns();
+    bpf_map_update_elem(&start, &req, &ts, BPF_ANY);
+    return 0;
+}
+
+char LICENSE[] SEC("license") = "GPL";

examples/c-form/ex6.bpf.c

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+
+
+#define TASK_COMM_LEN 16
+
+// Define output data structure
+struct data_t {
+    __u32 pid;
+    __u64 ts;
+//    char comm[TASK_COMM_LEN];
+};
+
+// Define a perf event output map
+struct {
+    __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
+    __uint(key_size, sizeof(__u32));
+    __uint(value_size, sizeof(__u32));
+} events SEC(".maps");
+
+SEC("tracepoint/syscalls/sys_enter_clone")
+int hello(struct pt_regs *ctx)
+{
+    struct data_t data = {};
+    
+    // Get PID (lower 32 bits of the 64-bit value returned)
+    data.pid = bpf_get_current_pid_tgid() & 0xFFFFFFFF;
+    
+    // Get timestamp
+    data.ts = bpf_ktime_get_ns();
+    
+    // Get current process name
+//    bpf_get_current_comm(&data.comm, sizeof(data.comm));
+    
+    // Submit data to userspace via perf event
+    bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, 
+                         &data, sizeof(data));
+    
+    return 0;
+}
+
+char LICENSE[] SEC("license") = "GPL";

examples/check.sh

@@ -22,11 +22,23 @@ case "$1" in
         sudo rm -f "$PIN_PATH"
         echo "[+] Stopped"
         ;;
+    xdp)
+        echo "[*] Loading and running $FILE"
+        sudo bpftool net detach xdp dev $3
+        sudo bpftool prog load "$FILE" "$PIN_PATH" type xdp
+        sudo bpftool net attach xdp pinned "$PIN_PATH" dev $3
+        echo "[+] Program loaded. Press Ctrl+C to stop"
+        sudo cat /sys/kernel/debug/tracing/trace_pipe
+        sudo bpftool net detach xdp dev $3
+        sudo rm -rf "$PIN_PATH"
+        echo "[+] Stopped"
+        ;;
     *)
         echo "Usage: $0 <check|run|stop> <file.o>"
         echo "Examples:"
         echo "  $0 check program.bpf.o"
         echo "  $0 run program.bpf.o"
+        echo "  $0 xdp program.bpf.o wlp6s0"
         echo "  $0 stop"
         exit 1
         ;;

examples/execve2.py

@@ -1,13 +1,14 @@
 from pythonbpf.decorators import bpf, map, section, bpfglobal
 from ctypes import c_void_p, c_int64, c_int32, c_uint64
-from pythonbpf.helpers import bpf_ktime_get_ns
+from pythonbpf.helpers import ktime
 from pythonbpf.maps import HashMap
 
 
 @bpf
 @map
 def last() -> HashMap:
-    return HashMap(key_type=c_uint64, value_type=c_uint64, max_entries=1)
+    return HashMap(key=c_uint64, value=c_uint64, max_entries=1)
+
 
 @bpf
 @section("tracepoint/syscalls/sys_enter_execve")
@@ -24,9 +25,10 @@ def hello_again(ctx: c_void_p) -> c_int64:
     key = 0
     tsp = last().lookup(key)
     print(tsp)
-    ts = bpf_ktime_get_ns()
+    ts = ktime()
     return c_int64(0)
 
+
 @bpf
 @bpfglobal
 def LICENSE() -> str:

examples/execve3.py

@@ -1,13 +1,14 @@
-from pythonbpf.decorators import bpf, map, section, bpfglobal
-from ctypes import c_void_p, c_int64, c_int32, c_uint64
-from pythonbpf.helpers import bpf_ktime_get_ns
+from pythonbpf import bpf, map, section, bpfglobal, compile
+from pythonbpf.helpers import ktime, deref
 from pythonbpf.maps import HashMap
 
+from ctypes import c_void_p, c_int64, c_int32, c_uint64
+
 
 @bpf
 @map
 def last() -> HashMap:
-    return HashMap(key_type=c_uint64, value_type=c_uint64, max_entries=1)
+    return HashMap(key=c_uint64, value=c_uint64, max_entries=3)
 
 
 @bpf
@@ -23,14 +24,25 @@ def hello(ctx: c_void_p) -> c_int32:
 def hello_again(ctx: c_void_p) -> c_int64:
     print("exited")
     key = 0
+    delta = 0
+    dddelta = 0
     tsp = last().lookup(key)
-    if tsp:
-        delta = (bpf_ktime_get_ns() - tsp.value)
-        if delta < 1000000000:
+    if True:
+        delta = ktime()
+        ddelta = deref(delta)
+        ttsp = deref(deref(tsp))
+        dddelta = ddelta - ttsp
+        if dddelta < 1000000000:
             print("execve called within last second")
         last().delete(key)
-    ts = bpf_ktime_get_ns()
+    ts = ktime()
     last().update(key, ts)
+
+    va = 8
+    nm = 5 + va
+    al = 6 & 3
+    print(f"this is a variable {nm}")
+
     return c_int64(0)
 
 
@@ -38,3 +50,6 @@ def hello_again(ctx: c_void_p) -> c_int64:
 @bpfglobal
 def LICENSE() -> str:
     return "GPL"
+
+
+compile()

examples/execve4.py

@@ -0,0 +1,28 @@
+from pythonbpf import bpf, map, section, bpfglobal, compile
+from pythonbpf.helpers import ktime, deref
+from pythonbpf.maps import HashMap
+
+from ctypes import c_void_p, c_int64, c_int32, c_uint64
+
+
+@bpf
+@map
+def last() -> HashMap:
+    return HashMap(key=c_uint64, value=c_uint64, max_entries=3)
+
+
+@bpf
+@section("blk_start_request")
+def trace_start(ctx: c_void_p) -> c_int32:
+    ts = ktime()
+    print("req started")
+    return c_int32(0)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

examples/execve5.py

@@ -0,0 +1,39 @@
+from pythonbpf import bpf, map, struct, section, bpfglobal, compile
+from pythonbpf.helpers import ktime, pid
+from pythonbpf.maps import PerfEventArray
+
+from ctypes import c_void_p, c_int64, c_int32, c_uint64
+
+
+@bpf
+@struct
+class data_t:
+    pid: c_uint64
+    ts: c_uint64
+
+
+@bpf
+@map
+def events() -> PerfEventArray:
+    return PerfEventArray(key_size=c_int32, value_size=c_int32)
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_clone")
+def hello(ctx: c_void_p) -> c_int32:
+    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}")
+    return c_int32(0)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pythonbpf"
-version = "0.1.0"
+version = "0.1.2"
 description = "Reduced Python frontend for eBPF"
 authors = [
   { name = "r41k0u", email="pragyanshchaturvedi18@gmail.com" },

pythonbpf/__init__.py

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

pythonbpf/binary_ops.py

@@ -0,0 +1,81 @@
+import ast
+from llvmlite import ir
+
+
+def recursive_dereferencer(var, builder):
+    """ dereference until primitive type comes out"""
+    if var.type == ir.PointerType(ir.PointerType(ir.IntType(64))):
+        a = builder.load(var)
+        return recursive_dereferencer(a, builder)
+    elif var.type == ir.PointerType(ir.IntType(64)):
+        a = builder.load(var)
+        return recursive_dereferencer(a, builder)
+    elif var.type == ir.IntType(64):
+        return var
+    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
+    right = rval.right
+    op = rval.op
+
+    # Handle left operand
+    if isinstance(left, ast.Name):
+        if left.id in local_sym_tab:
+            left = recursive_dereferencer(local_sym_tab[left.id], builder)
+        else:
+            raise SyntaxError(f"Undefined variable: {left.id}")
+    elif isinstance(left, ast.Constant):
+        left = ir.Constant(ir.IntType(64), left.value)
+    else:
+        raise SyntaxError("Unsupported left operand type")
+
+    if isinstance(right, ast.Name):
+        if right.id in local_sym_tab:
+            right = recursive_dereferencer(local_sym_tab[right.id], builder)
+        else:
+            raise SyntaxError(f"Undefined variable: {right.id}")
+    elif isinstance(right, ast.Constant):
+        right = ir.Constant(ir.IntType(64), right.value)
+    else:
+        raise SyntaxError("Unsupported right operand type")
+
+    print(f"left is {left}, right is {right}, op is {op}")
+
+    if isinstance(op, ast.Add):
+        builder.store(builder.add(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.Sub):
+        builder.store(builder.sub(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.Mult):
+        builder.store(builder.mul(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.Div):
+        builder.store(builder.sdiv(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.Mod):
+        builder.store(builder.srem(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.LShift):
+        builder.store(builder.shl(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.RShift):
+        builder.store(builder.lshr(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.BitOr):
+        builder.store(builder.or_(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.BitXor):
+        builder.store(builder.xor(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.BitAnd):
+        builder.store(builder.and_(left, right),
+                      local_sym_tab[var_name])
+    elif isinstance(op, ast.FloorDiv):
+        builder.store(builder.udiv(left, right),
+                      local_sym_tab[var_name])
+    else:
+        raise SyntaxError("Unsupported binary operation")

pythonbpf/bpf_helper_handler.py

@@ -1,20 +1,52 @@
 import ast
 from llvmlite import ir
+from .expr_pass import eval_expr
 
 
-def bpf_ktime_get_ns_emitter(call, module, builder, func):
-    pass
+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.
+    """
+    # func is an arg to just have a uniform signature with other emitters
+    helper_id = ir.Constant(ir.IntType(64), 5)
+    fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
+    result = builder.call(fn_ptr, [], tail=False)
+    return result
 
 
-def bpf_map_lookup_elem_emitter(map_ptr, key_ptr, module, builder):
+def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
     """
     Emit LLVM IR for bpf_map_lookup_elem helper function call.
     """
-    # Cast pointers to void*
+    if call.args and len(call.args) != 1:
+        raise ValueError("Map lookup expects exactly one argument, got "
+                         f"{len(call.args)}")
+    key_arg = call.args[0]
+    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]
+        else:
+            raise ValueError(
+                f"Key variable {key_name} not found in local symbol table.")
+    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
+        # handle constant integer keys
+        key_val = key_arg.value
+        key_type = ir.IntType(64)
+        key_ptr = builder.alloca(key_type)
+        key_ptr.align = key_type // 8
+        builder.store(ir.Constant(key_type, key_val), key_ptr)
+    else:
+        raise NotImplementedError(
+            "Only simple variable names are supported as keys in map lookup.")
+
+    if key_ptr is None:
+        raise ValueError("Key pointer is None.")
+
     map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
 
-    # Define function type for bpf_map_lookup_elem
-    # The function takes two void* arguments and returns void*
     fn_type = ir.FunctionType(
         ir.PointerType(),  # Return type: void*
         [ir.PointerType(), ir.PointerType()],  # Args: (void*, void*)
@@ -26,16 +58,92 @@ def bpf_map_lookup_elem_emitter(map_ptr, key_ptr, module, builder):
     fn_addr = ir.Constant(ir.IntType(64), 1)
     fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
 
-    # Call the helper function
     result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
 
     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("%lld")
+                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"
@@ -62,3 +170,211 @@ 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.
+    Expected call signature: map.update(key, value, flags=0)
+    """
+    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
+        if local_sym_tab and key_name in local_sym_tab:
+            key_ptr = local_sym_tab[key_name]
+        else:
+            raise ValueError(
+                f"Key variable {key_name} not found in local symbol table.")
+    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
+        # Handle constant integer keys
+        key_val = key_arg.value
+        key_type = ir.IntType(64)
+        key_ptr = builder.alloca(key_type)
+        key_ptr.align = key_type.width // 8
+        builder.store(ir.Constant(key_type, key_val), key_ptr)
+    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
+        if local_sym_tab and value_name in local_sym_tab:
+            value_ptr = local_sym_tab[value_name]
+        else:
+            raise ValueError(
+                f"Value variable {value_name} not found in local symbol table.")
+    elif isinstance(value_arg, ast.Constant) and isinstance(value_arg.value, int):
+        # Handle constant integers
+        value_val = value_arg.value
+        value_type = ir.IntType(64)
+        value_ptr = builder.alloca(value_type)
+        value_ptr.align = value_type.width // 8
+        builder.store(ir.Constant(value_type, value_val), value_ptr)
+    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):
+            flags_val = flags_arg.value
+        elif isinstance(flags_arg, ast.Name):
+            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_val = builder.load(flags_ptr)
+            else:
+                raise ValueError(
+                    f"Flags variable {flags_name} not found in local symbol table.")
+        else:
+            raise NotImplementedError(
+                "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.")
+
+    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
+    fn_type = ir.FunctionType(
+        ir.IntType(64),
+        [ir.PointerType(), ir.PointerType(), ir.PointerType(), ir.IntType(64)],
+        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)
+
+    return result
+
+
+def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
+    """
+    Emit LLVM IR for bpf_map_delete_elem helper function call.
+    Expected call signature: map.delete(key)
+    """
+    # Check for correct number of arguments
+    if not call.args or len(call.args) != 1:
+        raise ValueError("Map delete expects exactly 1 argument (key), got "
+                         f"{len(call.args)}")
+
+    key_arg = call.args[0]
+
+    # Handle key argument
+    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]
+        else:
+            raise ValueError(
+                f"Key variable {key_name} not found in local symbol table.")
+    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
+        # Handle constant integer keys
+        key_val = key_arg.value
+        key_type = ir.IntType(64)
+        key_ptr = builder.alloca(key_type)
+        key_ptr.align = key_type.width // 8
+        builder.store(ir.Constant(key_type, key_val), key_ptr)
+    else:
+        raise NotImplementedError(
+            "Only simple variable names and integer constants are supported as keys in map delete.")
+
+    if key_ptr is None:
+        raise ValueError("Key pointer is None.")
+
+    # Cast map pointer to void*
+    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
+
+    # Define function type for bpf_map_delete_elem
+    fn_type = ir.FunctionType(
+        ir.IntType(64),  # Return type: int64 (status code)
+        [ir.PointerType(), ir.PointerType()],  # Args: (void*, void*)
+        var_arg=False
+    )
+    fn_ptr_type = ir.PointerType(fn_type)
+
+    # Helper ID 3 is bpf_map_delete_elem
+    fn_addr = ir.Constant(ir.IntType(64), 3)
+    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+
+    # Call the helper function
+    result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
+
+    return result
+
+
+def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
+    """
+    Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
+    """
+    # func is an arg to just have a uniform signature with other emitters
+    helper_id = ir.Constant(ir.IntType(64), 14)
+    fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
+    result = builder.call(fn_ptr, [], tail=False)
+
+    # Extract the lower 32 bits (PID) using bitwise AND with 0xFFFFFFFF
+    mask = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
+    pid = builder.and_(result, mask)
+    return pid
+
+
+helper_func_list = {
+    "lookup": bpf_map_lookup_elem_emitter,
+    "print": bpf_printk_emitter,
+    "ktime": bpf_ktime_get_ns_emitter,
+    "update": bpf_map_update_elem_emitter,
+    "delete": bpf_map_delete_elem_emitter,
+    "pid": bpf_get_current_pid_tgid_emitter,
+}
+
+
+def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None):
+    if isinstance(call.func, ast.Name):
+        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, None, module, builder, func, local_sym_tab)
+        else:
+            raise NotImplementedError(
+                f"Function {func_name} is not implemented as a helper function.")
+    elif isinstance(call.func, ast.Attribute):
+        # likely a map method call
+        if isinstance(call.func.value, ast.Call) and isinstance(call.func.value.func, ast.Name):
+            map_name = call.func.value.func.id
+            method_name = call.func.attr
+            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:
+                    return helper_func_list[method_name](
+                        call, map_ptr, module, builder, local_sym_tab)
+                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.")
+        else:
+            raise NotImplementedError(
+                "Attribute not supported for map method calls.")

pythonbpf/codegen.py

@@ -3,15 +3,19 @@ from llvmlite import ir
 from .license_pass import license_processing
 from .functions_pass import func_proc
 from .maps_pass import maps_proc
+from .structs_pass import structs_proc
 from .globals_pass import globals_processing
 import os
+import subprocess
+import inspect
+from pathlib import Path
 
 
 def find_bpf_chunks(tree):
     """Find all functions decorated with @bpf in the AST."""
     bpf_functions = []
     for node in ast.walk(tree):
-        if isinstance(node, ast.FunctionDef):
+        if isinstance(node, ast.FunctionDef) or isinstance(node, ast.ClassDef):
             for decorator in node.decorator_list:
                 if isinstance(decorator, ast.Name) and decorator.id == "bpf":
                     bpf_functions.append(node)
@@ -25,10 +29,11 @@ def processor(source_code, filename, module):
 
     bpf_chunks = find_bpf_chunks(tree)
     for func_node in bpf_chunks:
-        print(f"Found BPF function: {func_node.name}")
+        print(f"Found BPF function/struct: {func_node.name}")
 
+    structs_sym_tab = structs_proc(tree, module, bpf_chunks)
     map_sym_tab = maps_proc(tree, module, bpf_chunks)
-    func_proc(tree, module, bpf_chunks, map_sym_tab)
+    func_proc(tree, module, bpf_chunks, map_sym_tab, structs_sym_tab)
 
     license_processing(tree, module)
     globals_processing(tree, module)
@@ -47,7 +52,7 @@ def compile_to_ir(filename: str, output: str):
             "filename": filename,
             "directory": os.path.dirname(filename)
         })
-        
+
         module._debug_compile_unit = module.add_debug_info("DICompileUnit", {       # type: ignore
             "language": 29,  # DW_LANG_C11
             "file": module._file_metadata,      # type: ignore
@@ -59,7 +64,8 @@ def compile_to_ir(filename: str, output: str):
             "nameTableKind": 0
         }, is_distinct=True)
 
-        module.add_named_metadata("llvm.dbg.cu", module._debug_compile_unit)        # type: ignore
+        module.add_named_metadata(
+            "llvm.dbg.cu", module._debug_compile_unit)        # type: ignore
 
     processor(source, filename, module)
 
@@ -92,3 +98,21 @@ def compile_to_ir(filename: str, output: str):
         f.write("\n")
 
     return output
+
+
+def compile():
+    # Look one level up the stack to the caller of this function
+    caller_frame = inspect.stack()[1]
+    caller_file = Path(caller_frame.filename).resolve()
+
+    ll_file = Path("/tmp") / caller_file.with_suffix(".ll").name
+    o_file = caller_file.with_suffix(".o")
+
+    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)
+
+    print(f"Object written to {o_file}, {ll_file} can be removed")

pythonbpf/decorators.py

@@ -9,11 +9,19 @@ def bpfglobal(func):
     func._is_bpfglobal = True
     return func
 
+
 def map(func):
     """Decorator to mark a function as a BPF map."""
     func._is_map = True
     return func
 
+
+def struct(cls):
+    """Decorator to mark a class as a BPF struct."""
+    cls._is_struct = True
+    return cls
+
+
 def section(name: str):
     def wrapper(fn):
         fn._section = name

pythonbpf/expr_pass.py

@@ -0,0 +1,71 @@
+import ast
+from llvmlite import ir
+
+
+def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
+    print(f"Evaluating expression: {expr}")
+    if isinstance(expr, ast.Name):
+        if expr.id in local_sym_tab:
+            var = local_sym_tab[expr.id]
+            val = builder.load(var)
+            return val
+        else:
+            print(f"Undefined variable {expr.id}")
+            return None
+    elif isinstance(expr, ast.Constant):
+        if isinstance(expr.value, int):
+            return ir.Constant(ir.IntType(64), expr.value)
+        elif isinstance(expr.value, bool):
+            return ir.Constant(ir.IntType(1), int(expr.value))
+        else:
+            print("Unsupported constant type")
+            return None
+    elif isinstance(expr, ast.Call):
+        # delayed import to avoid circular dependency
+        from .bpf_helper_handler import helper_func_list, handle_helper_call
+
+        if isinstance(expr.func, ast.Name):
+            # check deref
+            if expr.func.id == "deref":
+                print(f"Handling deref {ast.dump(expr)}")
+                if len(expr.args) != 1:
+                    print("deref takes exactly one argument")
+                    return None
+                arg = expr.args[0]
+                if isinstance(arg, ast.Call) and isinstance(arg.func, ast.Name) and arg.func.id == "deref":
+                    print("Multiple deref not supported")
+                    return None
+                if isinstance(arg, ast.Name):
+                    if arg.id in local_sym_tab:
+                        arg = local_sym_tab[arg.id]
+                    else:
+                        print(f"Undefined variable {arg.id}")
+                        return None
+                if arg is None:
+                    print("Failed to evaluate deref argument")
+                    return None
+                val = builder.load(arg)
+                return val
+
+            # 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)
+        elif isinstance(expr.func, ast.Attribute):
+            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)
+    print("Unsupported expression evaluation")
+    return None
+
+
+def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
+    """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)
+    else:
+        print("Unsupported expression type")

pythonbpf/functions_pass.py

@@ -1,8 +1,13 @@
 from llvmlite import ir
 import ast
 
-from .bpf_helper_handler import bpf_printk_emitter, bpf_ktime_get_ns_emitter, bpf_map_lookup_elem_emitter
+
+from .bpf_helper_handler import helper_func_list, handle_helper_call
 from .type_deducer import ctypes_to_ir
+from .binary_ops import handle_binary_op
+from .expr_pass import eval_expr, handle_expr
+
+local_var_metadata = {}
 
 
 def get_probe_string(func_node):
@@ -22,7 +27,7 @@ def get_probe_string(func_node):
     return "helper"
 
 
-def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
+def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, structs_sym_tab):
     """Handle assignment statements in the function body."""
     if len(stmt.targets) != 1:
         print("Unsupported multiassignment")
@@ -31,122 +36,377 @@ def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
     num_types = ("c_int32", "c_int64", "c_uint32", "c_uint64")
 
     target = stmt.targets[0]
-    if not isinstance(target, ast.Name):
+    print(f"Handling assignment to {ast.dump(target)}")
+    if not isinstance(target, ast.Name) and not isinstance(target, ast.Attribute):
         print("Unsupported assignment target")
         return
-    var_name = target.id
+    var_name = target.id if isinstance(target, ast.Name) else target.value.id
     rval = stmt.value
-    if isinstance(rval, ast.Constant):
-        if isinstance(rval.value, int):
-            # Assume c_int64 for now
-            # TODO: make symtab for this
-            var = builder.alloca(ir.IntType(64), name=var_name)
-            var.align = 8
-            builder.store(ir.Constant(ir.IntType(64), rval.value), var)
-            local_sym_tab[var_name] = var
+    if isinstance(target, ast.Attribute):
+        # struct field assignment
+        field_name = target.attr
+        if var_name in local_sym_tab and var_name in local_var_metadata:
+            struct_type = local_var_metadata[var_name]
+            struct_info = structs_sym_tab[struct_type]
+
+            if field_name in struct_info["fields"]:
+                field_idx = struct_info["fields"][field_name]
+                struct_ptr = local_sym_tab[var_name]
+                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)
+                if val is None:
+                    print("Failed to evaluate struct field assignment")
+                    return
+                builder.store(val, field_ptr)
+                print(f"Assigned to struct field {var_name}.{field_name}")
+                return
+    elif isinstance(rval, ast.Constant):
+        if isinstance(rval.value, bool):
+            if rval.value:
+                builder.store(ir.Constant(ir.IntType(1), 1),
+                              local_sym_tab[var_name])
+            else:
+                builder.store(ir.Constant(ir.IntType(1), 0),
+                              local_sym_tab[var_name])
             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] = var
+            print(f"Assigned constant {rval.value} to {var_name}")
+        else:
+            print("Unsupported constant type")
     elif isinstance(rval, ast.Call):
         if isinstance(rval.func, ast.Name):
             call_type = rval.func.id
             print(f"Assignment call type: {call_type}")
             if call_type in num_types and len(rval.args) == 1 and isinstance(rval.args[0], ast.Constant) and isinstance(rval.args[0].value, int):
                 ir_type = ctypes_to_ir(call_type)
-                var = builder.alloca(ir_type, name=var_name)
-                var.align = ir_type.width // 8
-                builder.store(ir.Constant(ir_type, rval.args[0].value), var)
+                # 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])
                 print(f"Assigned {call_type} constant "
                       f"{rval.args[0].value} to {var_name}")
-                local_sym_tab[var_name] = var
+                # local_sym_tab[var_name] = var
+            elif call_type in helper_func_list:
+                # 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])
+                # 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)
+                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])
+                # 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:
+                struct_info = structs_sym_tab[call_type]
+                ir_type = struct_info["type"]
+                # var = builder.alloca(ir_type, name=var_name)
+                # Null init
+                builder.store(ir.Constant(ir_type, None),
+                              local_sym_tab[var_name])
+                local_var_metadata[var_name] = call_type
+                print(f"Assigned struct {call_type} to {var_name}")
+                # local_sym_tab[var_name] = var
             else:
                 print(f"Unsupported assignment call type: {call_type}")
         elif isinstance(rval.func, ast.Attribute):
-            if isinstance(rval.func.attr, str) and rval.func.attr == "lookup":
-                # Get map name and check symtab
-                # maps are called as funcs
-                if isinstance(rval.func.value, ast.Call) and isinstance(rval.func.value.func, ast.Name):
-                    map_name = rval.func.value.func.id
-                    if map_name in map_sym_tab:
-                        map_global = map_sym_tab[map_name]
-                        print(f"Found map {map_name} in symtab for lookup")
-                        if len(rval.args) != 1:
-                            print("Unsupported lookup with != 1 arg")
-                            return
-                        key_arg = rval.args[0]
-                        print(f"Lookup key arg type: {type(key_arg)}")
-                        # TODO: implement a parse_arg ffs as this can be a fucking expr
-                        if isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
-                            key_val = key_arg.value
-                            key_type = ir.IntType(64)
-                            print(f"Key type: {key_type}")
-                            print(f"Key val: {key_val}")
-                            key_var = builder.alloca(key_type)
-                            key_var.align = key_type // 8
-                            builder.store(ir.Constant(
-                                key_type, key_val), key_var)
-                        elif isinstance(key_arg, ast.Name):
-                            # Check in local symtab first
-                            if key_arg.id in local_sym_tab:
-                                key_var = local_sym_tab[key_arg.id]
-                                key_type = key_var.type.pointee
-                            elif key_arg.id in map_sym_tab:
-                                key_var = map_sym_tab[key_arg.id]
-                                key_type = key_var.type.pointee
-                            else:
-                                print("Key variable "
-                                      f"{key_arg.id} not found in symtabs")
-                                return
-                            print(f"Found key variable {key_arg.id} in symtab")
-                            print(f"Key type: {key_type}")
-                        else:
-                            print("Unsupported lookup key arg")
-                            return
-
-                        # TODO: generate call to bpf_map_lookup_elem
-                        result_ptr = bpf_map_lookup_elem_emitter(
-                            map_global, key_var, module, builder)
-
-                    else:
-                        print(f"Map {map_name} not found in symbol table")
+            print(f"Assignment call attribute: {ast.dump(rval.func)}")
+            if isinstance(rval.func.value, ast.Name):
+                # TODO: probably a struct access
+                print(f"TODO STRUCT ACCESS {ast.dump(rval)}")
+            elif isinstance(rval.func.value, ast.Call) and isinstance(rval.func.value.func, ast.Name):
+                map_name = rval.func.value.func.id
+                method_name = rval.func.attr
+                if map_name in map_sym_tab:
+                    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)
+                        # var = builder.alloca(ir.IntType(64), name=var_name)
+                        # var.align = 8
+                        builder.store(val, local_sym_tab[var_name])
+                        # local_sym_tab[var_name] = var
             else:
-                print("Unsupported assignment from method call")
+                print("Unsupported assignment call structure")
+        else:
+            print("Unsupported assignment call function type")
+    elif isinstance(rval, ast.BinOp):
+        handle_binary_op(rval, module, builder, var_name,
+                         local_sym_tab, map_sym_tab, func)
+    else:
+        print("Unsupported assignment value type")
 
 
-def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab):
+def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
+    if isinstance(cond, ast.Constant):
+        if isinstance(cond.value, bool):
+            return ir.Constant(ir.IntType(1), int(cond.value))
+        elif isinstance(cond.value, int):
+            return ir.Constant(ir.IntType(1), int(bool(cond.value)))
+        else:
+            print("Unsupported constant type in condition")
+            return None
+    elif isinstance(cond, ast.Name):
+        if cond.id in local_sym_tab:
+            var = local_sym_tab[cond.id]
+            val = builder.load(var)
+            if val.type != ir.IntType(1):
+                # Convert nonzero values to true, zero to false
+                if isinstance(val.type, ir.PointerType):
+                    # For pointer types, compare with null pointer
+                    zero = ir.Constant(val.type, None)
+                else:
+                    # For integer types, compare with zero
+                    zero = ir.Constant(val.type, 0)
+                val = builder.icmp_signed("!=", val, zero)
+            return val
+        else:
+            print(f"Undefined variable {cond.id} in condition")
+            return None
+    elif isinstance(cond, ast.Compare):
+        lhs = eval_expr(func, module, builder, cond.left,
+                        local_sym_tab, map_sym_tab)
+        if len(cond.ops) != 1 or len(cond.comparators) != 1:
+            print("Unsupported complex comparison")
+            return None
+        rhs = eval_expr(func, module, builder,
+                        cond.comparators[0], local_sym_tab, map_sym_tab)
+        op = cond.ops[0]
+
+        if lhs.type != rhs.type:
+            if isinstance(lhs.type, ir.IntType) and isinstance(rhs.type, ir.IntType):
+                # Extend the smaller type to the larger type
+                if lhs.type.width < rhs.type.width:
+                    lhs = builder.sext(lhs, rhs.type)
+                elif lhs.type.width > rhs.type.width:
+                    rhs = builder.sext(rhs, lhs.type)
+            else:
+                print("Type mismatch in comparison")
+                return None
+
+        if isinstance(op, ast.Eq):
+            return builder.icmp_signed("==", lhs, rhs)
+        elif isinstance(op, ast.NotEq):
+            return builder.icmp_signed("!=", lhs, rhs)
+        elif isinstance(op, ast.Lt):
+            return builder.icmp_signed("<", lhs, rhs)
+        elif isinstance(op, ast.LtE):
+            return builder.icmp_signed("<=", lhs, rhs)
+        elif isinstance(op, ast.Gt):
+            return builder.icmp_signed(">", lhs, rhs)
+        elif isinstance(op, ast.GtE):
+            return builder.icmp_signed(">=", lhs, rhs)
+        else:
+            print("Unsupported comparison operator")
+            return None
+    else:
+        print("Unsupported condition expression")
+        return None
+
+
+def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
+    """Handle if statements in the function body."""
+    print("Handling if statement")
+    start = builder.block.parent
+    then_block = func.append_basic_block(name="if.then")
+    merge_block = func.append_basic_block(name="if.end")
+    if stmt.orelse:
+        else_block = func.append_basic_block(name="if.else")
+    else:
+        else_block = None
+
+    cond = handle_cond(func, module, builder, stmt.test,
+                       local_sym_tab, map_sym_tab)
+    if else_block:
+        builder.cbranch(cond, then_block, else_block)
+    else:
+        builder.cbranch(cond, then_block, merge_block)
+
+    builder.position_at_end(then_block)
+    for s in stmt.body:
+        process_stmt(func, module, builder, s,
+                     local_sym_tab, map_sym_tab, False)
+    if not builder.block.is_terminated:
+        builder.branch(merge_block)
+
+    if else_block:
+        builder.position_at_end(else_block)
+        for s in stmt.orelse:
+            process_stmt(func, module, builder, s,
+                         local_sym_tab, map_sym_tab, False)
+        if not builder.block.is_terminated:
+            builder.branch(merge_block)
+
+    builder.position_at_end(merge_block)
+
+
+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)
+    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)
+    elif isinstance(stmt, ast.Return):
+        if stmt.value is None:
+            builder.ret(ir.Constant(ir.IntType(32), 0))
+            did_return = True
+        elif isinstance(stmt.value, ast.Call) and isinstance(stmt.value.func, ast.Name) and len(stmt.value.args) == 1 and isinstance(stmt.value.args[0], ast.Constant) and isinstance(stmt.value.args[0].value, int):
+            call_type = stmt.value.func.id
+            if ctypes_to_ir(call_type) != ret_type:
+                raise ValueError("Return type mismatch: expected"
+                                 f"{ctypes_to_ir(call_type)}, got {call_type}")
+            else:
+                builder.ret(ir.Constant(
+                    ret_type, stmt.value.args[0].value))
+                did_return = True
+        elif isinstance(stmt.value, ast.Name):
+            if stmt.value.id == "XDP_PASS":
+                builder.ret(ir.Constant(ret_type, 2))
+                did_return = True
+            elif stmt.value.id == "XDP_DROP":
+                builder.ret(ir.Constant(ret_type, 1))
+                did_return = True
+            else:
+                raise ValueError("Failed to evaluate return expression")
+        else:
+            raise ValueError("Unsupported return value")
+    return did_return
+
+
+def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab):
+    for stmt in body:
+        if isinstance(stmt, ast.If):
+            if stmt.body:
+                local_sym_tab = allocate_mem(
+                    module, builder, stmt.body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab)
+            if stmt.orelse:
+                local_sym_tab = allocate_mem(
+                    module, builder, stmt.orelse, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab)
+        elif isinstance(stmt, ast.Assign):
+            if len(stmt.targets) != 1:
+                print("Unsupported multiassignment")
+                continue
+            target = stmt.targets[0]
+            if not isinstance(target, ast.Name):
+                print("Unsupported assignment target")
+                continue
+            var_name = target.id
+            rval = stmt.value
+            if isinstance(rval, ast.Call):
+                if isinstance(rval.func, ast.Name):
+                    call_type = rval.func.id
+                    if call_type in ("c_int32", "c_int64", "c_uint32", "c_uint64"):
+                        ir_type = ctypes_to_ir(call_type)
+                        var = builder.alloca(ir_type, name=var_name)
+                        var.align = ir_type.width // 8
+                        print(
+                            f"Pre-allocated variable {var_name} of type {call_type}")
+                    elif call_type in helper_func_list:
+                        # Assume return type is int64 for now
+                        ir_type = ir.IntType(64)
+                        var = builder.alloca(ir_type, name=var_name)
+                        var.align = ir_type.width // 8
+                        print(
+                            f"Pre-allocated variable {var_name} for helper")
+                    elif call_type == "deref" and len(rval.args) == 1:
+                        # Assume return type is int64 for now
+                        ir_type = ir.IntType(64)
+                        var = builder.alloca(ir_type, name=var_name)
+                        var.align = ir_type.width // 8
+                        print(
+                            f"Pre-allocated variable {var_name} for deref")
+                    elif call_type in structs_sym_tab:
+                        struct_info = structs_sym_tab[call_type]
+                        ir_type = struct_info["type"]
+                        var = builder.alloca(ir_type, name=var_name)
+                        local_var_metadata[var_name] = call_type
+                        print(
+                            f"Pre-allocated variable {var_name} for struct {call_type}")
+                elif isinstance(rval.func, ast.Attribute):
+                    ir_type = ir.PointerType(ir.IntType(64))
+                    var = builder.alloca(ir_type, name=var_name)
+                    # var.align = ir_type.width // 8
+                    print(
+                        f"Pre-allocated variable {var_name} for map")
+                else:
+                    print("Unsupported assignment call function type")
+                    continue
+            elif isinstance(rval, ast.Constant):
+                if isinstance(rval.value, bool):
+                    ir_type = ir.IntType(1)
+                    var = builder.alloca(ir_type, name=var_name)
+                    var.align = 1
+                    print(
+                        f"Pre-allocated variable {var_name} of type c_bool")
+                elif isinstance(rval.value, int):
+                    # Assume c_int64 for now
+                    ir_type = ir.IntType(64)
+                    var = builder.alloca(ir_type, name=var_name)
+                    var.align = ir_type.width // 8
+                    print(
+                        f"Pre-allocated variable {var_name} of type c_int64")
+                else:
+                    print("Unsupported constant type")
+                    continue
+            elif isinstance(rval, ast.BinOp):
+                # Assume c_int64 for now
+                ir_type = ir.IntType(64)
+                var = builder.alloca(ir_type, name=var_name)
+                var.align = ir_type.width // 8
+                print(
+                    f"Pre-allocated variable {var_name} of type c_int64")
+            else:
+                print("Unsupported assignment value type")
+                continue
+            local_sym_tab[var_name] = var
+    return local_sym_tab
+
+
+def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab, structs_sym_tab):
     """Process the body of a bpf function"""
     # TODO: A lot.  We just have print -> bpf_trace_printk for now
     did_return = False
 
     local_sym_tab = {}
 
+    # pre-allocate dynamic variables
+    local_sym_tab = allocate_mem(
+        module, builder, func_node.body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab)
+
+    print(f"Local symbol table: {local_sym_tab.keys()}")
+
     for stmt in func_node.body:
-        if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
-            call = stmt.value
-            if isinstance(call.func, ast.Name) and call.func.id == "print":
-                bpf_printk_emitter(call, module, builder, func)
-            if isinstance(call.func, ast.Name) and call.func.id == "bpf_ktime_get_ns":
-                bpf_ktime_get_ns_emitter(call, module, builder, func)
-        elif isinstance(stmt, ast.Assign):
-            handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab)
-        elif isinstance(stmt, ast.Return):
-            if stmt.value is None:
-                builder.ret(ir.Constant(ir.IntType(32), 0))
-                did_return = True
-            elif isinstance(stmt.value, ast.Call) and isinstance(stmt.value.func, ast.Name) and len(stmt.value.args) == 1 and isinstance(stmt.value.args[0], ast.Constant) and isinstance(stmt.value.args[0].value, int):
-                call_type = stmt.value.func.id
-                if ctypes_to_ir(call_type) != ret_type:
-                    raise ValueError("Return type mismatch: expected"
-                                     f"{ctypes_to_ir(call_type)}, got {call_type}")
-                else:
-                    builder.ret(ir.Constant(
-                        ret_type, stmt.value.args[0].value))
-                    did_return = True
-            else:
-                print("Unsupported return value")
+        did_return = process_stmt(func, module, builder, stmt, local_sym_tab,
+                                  map_sym_tab, structs_sym_tab, did_return, ret_type)
+
     if not did_return:
         builder.ret(ir.Constant(ir.IntType(32), 0))
 
 
-def process_bpf_chunk(func_node, module, return_type, map_sym_tab):
+def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_tab):
     """Process a single BPF chunk (function) and emit corresponding LLVM IR."""
 
     func_name = func_node.name
@@ -179,19 +439,17 @@ def process_bpf_chunk(func_node, module, return_type, map_sym_tab):
     block = func.append_basic_block(name="entry")
     builder = ir.IRBuilder(block)
 
-    process_func_body(module, builder, func_node, func, ret_type, map_sym_tab)
+    process_func_body(module, builder, func_node, func,
+                      ret_type, map_sym_tab, structs_sym_tab)
 
     return func
 
 
-def func_proc(tree, module, chunks, map_sym_tab):
+def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab):
     for func_node in chunks:
         is_global = False
         for decorator in func_node.decorator_list:
-            if isinstance(decorator, ast.Name) and decorator.id == "map":
-                is_global = True
-                break
-            elif isinstance(decorator, ast.Name) and decorator.id == "bpfglobal":
+            if isinstance(decorator, ast.Name) and decorator.id in ("map", "bpfglobal", "struct"):
                 is_global = True
                 break
         if is_global:
@@ -200,7 +458,7 @@ def func_proc(tree, module, chunks, map_sym_tab):
         print(f"Found probe_string of {func_node.name}: {func_type}")
 
         process_bpf_chunk(func_node, module, ctypes_to_ir(
-            infer_return_type(func_node)), map_sym_tab)
+            infer_return_type(func_node)), map_sym_tab, structs_sym_tab)
 
 
 def infer_return_type(func_node: ast.FunctionDef):

pythonbpf/helpers.py

@@ -1,4 +1,15 @@
 import ctypes
 
-def bpf_ktime_get_ns():
+def ktime():
     return ctypes.c_int64(0)
+
+def pid():
+    return ctypes.c_int32(0)
+
+def deref(ptr):
+    "dereference a pointer"
+    result = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_void_p)).contents.value
+    return result if result is not None else 0
+
+XDP_DROP = ctypes.c_int64(1)
+XDP_PASS = ctypes.c_int64(2)

pythonbpf/maps.py

@@ -1,7 +1,7 @@
 class HashMap:
-    def __init__(self, key_type, value_type, max_entries):
-        self.key_type = key_type
-        self.value_type = value_type
+    def __init__(self, key, value, max_entries):
+        self.key = key
+        self.value = value
         self.max_entries = max_entries
         self.entries = {}
 
@@ -10,15 +10,23 @@ class HashMap:
             return self.entries[key]
         else:
             return None
-            
+
     def delete(self, key):
         if key in self.entries:
             del self.entries[key]
         else:
             raise KeyError(f"Key {key} not found in map")
-            
-    def update(self, key, value):
+
+    # TODO: define the flags that can be added
+    def update(self, key, value, flags=None):
         if key in self.entries:
             self.entries[key] = value
         else:
             raise KeyError(f"Key {key} not found in map")
+
+
+class PerfEventArray:
+    def __init__(self, key_size, value_size):
+        self.key_type = key_size
+        self.value_type = value_size
+        self.entries = {}

pythonbpf/maps_pass.py

@@ -21,22 +21,28 @@ def maps_proc(tree, module, chunks):
     return map_sym_tab
 
 
+BPF_MAP_MAPPINGS = {
+    "HASH": 1,                # BPF_MAP_TYPE_HASH
+    "PERF_EVENT_ARRAY": 4,    # BPF_MAP_TYPE_PERF_EVENT_ARRAY
+}
+
+
 def create_bpf_map(module, map_name, map_params):
     """Create a BPF map in the module with the given parameters and debug info"""
 
+    map_type_str = map_params.get("type", "HASH")
+    map_type = BPF_MAP_MAPPINGS.get(map_type_str)
+
     # Create the anonymous struct type for BPF map
-    map_struct_type = ir.LiteralStructType([
-        ir.PointerType(),
-        ir.PointerType(),
-        ir.PointerType(),
-        ir.PointerType()
-    ])
+    map_struct_type = ir.LiteralStructType(
+        [ir.PointerType() for _ in range(len(map_params))])
 
     # Create the global variable
     map_global = ir.GlobalVariable(module, map_struct_type, name=map_name)
     map_global.linkage = 'dso_local'
     map_global.global_constant = False
-    map_global.initializer = ir.Constant(map_struct_type, None)     # type: ignore
+    map_global.initializer = ir.Constant(
+        map_struct_type, None)     # type: ignore
     map_global.section = ".maps"
     map_global.align = 8        # type: ignore
 
@@ -47,6 +53,7 @@ def create_bpf_map(module, map_name, map_params):
     map_sym_tab[map_name] = map_global
     return map_global
 
+
 def create_map_debug_info(module, map_global, map_name, map_params):
     """Generate debug information metadata for BPF map"""
     file_metadata = module._file_metadata
@@ -67,7 +74,8 @@ def create_map_debug_info(module, map_global, map_name, map_params):
     })
 
     # Create array type for map type field (array of 1 unsigned int)
-    array_subrange = module.add_debug_info("DISubrange", {"count": 1})
+    array_subrange = module.add_debug_info(
+        "DISubrange", {"count": BPF_MAP_MAPPINGS[map_params.get("type", "HASH")]})
     array_type = module.add_debug_info("DICompositeType", {
         "tag": dc.DW_TAG_array_type,
         "baseType": uint_type,
@@ -82,68 +90,75 @@ def create_map_debug_info(module, map_global, map_name, map_params):
         "size": 64
     })
 
-    max_entries_ptr = module.add_debug_info("DIDerivedType", {
-        "tag": dc.DW_TAG_pointer_type,
-        "baseType": array_type,
-        "size": 64
-    })
-
     key_ptr = module.add_debug_info("DIDerivedType", {
         "tag": dc.DW_TAG_pointer_type,
-        "baseType": uint_type,  # Adjust based on actual key type
+        # Adjust based on actual key type
+        "baseType": array_type if "key_size" in map_params else uint_type,
         "size": 64
     })
 
     value_ptr = module.add_debug_info("DIDerivedType", {
         "tag": dc.DW_TAG_pointer_type,
-        "baseType": ulong_type,  # Adjust based on actual value type
+        # Adjust based on actual value type
+        "baseType": array_type if "value_size" in map_params else ulong_type,
         "size": 64
     })
 
+    elements_arr = []
+
     # Create struct members
     # scope field does not appear for some reason
-    type_member = module.add_debug_info("DIDerivedType", {
-        "tag": dc.DW_TAG_member,
-        "name": "type",
-        "file": file_metadata,
-        "baseType": type_ptr,
-        "size": 64,
-        "offset": 0
-    })
+    cnt = 0
+    for elem in map_params:
+        if elem == "max_entries":
+            continue
+        if elem == "type":
+            ptr = type_ptr
+        elif "key" in elem:
+            ptr = key_ptr
+        else:
+            ptr = value_ptr
+        member = module.add_debug_info("DIDerivedType", {
+            "tag": dc.DW_TAG_member,
+            "name": elem,
+            "file": file_metadata,
+            "baseType": ptr,
+            "size": 64,
+            "offset": cnt * 64
+        })
+        elements_arr.append(member)
+        cnt += 1
 
-    max_entries_member = module.add_debug_info("DIDerivedType", {
-        "tag": dc.DW_TAG_member,
-        "name": "max_entries",
-        "file": file_metadata,
-        "baseType": max_entries_ptr,
-        "size": 64,
-        "offset": 64
-    })
-
-    key_member = module.add_debug_info("DIDerivedType", {
-        "tag": dc.DW_TAG_member,
-        "name": "key",
-        "file": file_metadata,
-        "baseType": key_ptr,
-        "size": 64,
-        "offset": 128
-    })
-
-    value_member = module.add_debug_info("DIDerivedType", {
-        "tag": dc.DW_TAG_member,
-        "name": "value",
-        "file": file_metadata,
-        "baseType": value_ptr,
-        "size": 64,
-        "offset": 192
-    })
+    if "max_entries" in map_params:
+        array_subrange_max_entries = module.add_debug_info(
+            "DISubrange", {"count": map_params["max_entries"]})
+        array_type_max_entries = module.add_debug_info("DICompositeType", {
+            "tag": dc.DW_TAG_array_type,
+            "baseType": uint_type,
+            "size": 32,
+            "elements": [array_subrange_max_entries]
+        })
+        max_entries_ptr = module.add_debug_info("DIDerivedType", {
+            "tag": dc.DW_TAG_pointer_type,
+            "baseType": array_type_max_entries,
+            "size": 64
+        })
+        max_entries_member = module.add_debug_info("DIDerivedType", {
+            "tag": dc.DW_TAG_member,
+            "name": "max_entries",
+            "file": file_metadata,
+            "baseType": max_entries_ptr,
+            "size": 64,
+            "offset": cnt * 64
+        })
+        elements_arr.append(max_entries_member)
 
     # Create the struct type
     struct_type = module.add_debug_info("DICompositeType", {
         "tag": dc.DW_TAG_structure_type,
         "file": file_metadata,
-        "size": 256,  # 4 * 64-bit pointers
-        "elements": [type_member, max_entries_member, key_member, value_member]
+        "size": 64 * len(elements_arr),  # 4 * 64-bit pointers
+        "elements": elements_arr,
     }, is_distinct=True)
 
     # Create global variable debug info
@@ -170,23 +185,23 @@ def create_map_debug_info(module, map_global, map_name, map_params):
 
 def process_hash_map(map_name, rval, module):
     print(f"Creating HashMap map: {map_name}")
-    map_params: dict[str, object] = {"map_type": "HASH"}
+    map_params: dict[str, object] = {"type": "HASH"}
 
     # Assuming order: key_type, value_type, max_entries
     if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name):
-        map_params["key_type"] = rval.args[0].id
+        map_params["key"] = rval.args[0].id
     if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name):
-        map_params["value_type"] = rval.args[1].id
+        map_params["value"] = rval.args[1].id
     if len(rval.args) >= 3 and isinstance(rval.args[2], ast.Constant):
         const_val = rval.args[2].value
         if isinstance(const_val, (int, str)):  # safe check
             map_params["max_entries"] = const_val
 
     for keyword in rval.keywords:
-        if keyword.arg == "key_type" and isinstance(keyword.value, ast.Name):
-            map_params["key_type"] = keyword.value.id
-        elif keyword.arg == "value_type" and isinstance(keyword.value, ast.Name):
-            map_params["value_type"] = keyword.value.id
+        if keyword.arg == "key" and isinstance(keyword.value, ast.Name):
+            map_params["key"] = keyword.value.id
+        elif keyword.arg == "value" and isinstance(keyword.value, ast.Name):
+            map_params["value"] = keyword.value.id
         elif keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant):
             const_val = keyword.value.value
             if isinstance(const_val, (int, str)):
@@ -196,11 +211,34 @@ def process_hash_map(map_name, rval, module):
     return create_bpf_map(module, map_name, map_params)
 
 
+def process_perf_event_map(map_name, rval, module):
+    print(f"Creating PerfEventArray map: {map_name}")
+    map_params = {"type": "PERF_EVENT_ARRAY"}
+
+    if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name):
+        map_params["key_size"] = rval.args[0].id
+    if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name):
+        map_params["value_size"] = rval.args[1].id
+
+    for keyword in rval.keywords:
+        if keyword.arg == "key_size" and isinstance(keyword.value, ast.Name):
+            map_params["key_size"] = keyword.value.id
+        elif keyword.arg == "value_size" and isinstance(keyword.value, ast.Name):
+            map_params["value_size"] = keyword.value.id
+
+    print(f"Map parameters: {map_params}")
+    return create_bpf_map(module, map_name, map_params)
+
+
 def process_bpf_map(func_node, module):
     """Process a BPF map (a function decorated with @map)"""
     map_name = func_node.name
     print(f"Processing BPF map: {map_name}")
 
+    BPF_MAP_TYPES = {"HashMap": process_hash_map,            # BPF_MAP_TYPE_HASH
+                     "PerfEventArray": process_perf_event_map,   # BPF_MAP_TYPE_PERF_EVENT_ARRAY
+                     }
+
     # For now, assume single return statement
     return_stmt = None
     for stmt in func_node.body:
@@ -213,7 +251,12 @@ def process_bpf_map(func_node, module):
     rval = return_stmt.value
 
     # Handle only HashMap maps
-    if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name) and rval.func.id == "HashMap":
-        process_hash_map(map_name, rval, module)
+    if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name):
+        if rval.func.id in BPF_MAP_TYPES:
+            handler = BPF_MAP_TYPES[rval.func.id]
+            handler(map_name, rval, module)
+        else:
+            print(f"Unknown map type {rval.func.id}, defaulting to HashMap")
+            process_hash_map(map_name, rval, module)
     else:
         raise ValueError("Function under @map must return a map")

pythonbpf/structs_pass.py

@@ -0,0 +1,39 @@
+import ast
+from llvmlite import ir
+from .type_deducer import ctypes_to_ir
+from . import dwarf_constants as dc
+
+structs_sym_tab = {}
+
+
+def structs_proc(tree, module, chunks):
+    for cls_node in chunks:
+        # Check if this class is a struct
+        is_struct = False
+        for decorator in cls_node.decorator_list:
+            if isinstance(decorator, ast.Name) and decorator.id == "struct":
+                is_struct = True
+                break
+        if is_struct:
+            print(f"Found BPF struct: {cls_node.name}")
+            process_bpf_struct(cls_node, module)
+            continue
+    return structs_sym_tab
+
+
+def process_bpf_struct(cls_node, module):
+    struct_name = cls_node.name
+    field_names = []
+    field_types = []
+
+    for item in cls_node.body:
+        if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
+            field_names.append(item.target.id)
+            field_types.append(ctypes_to_ir(item.annotation.id))
+
+    struct_type = ir.LiteralStructType(field_types)
+    structs_sym_tab[struct_name] = {
+        "type": struct_type,
+        "fields": {name: idx for idx, name in enumerate(field_names)}
+    }
+    print(f"Created struct {struct_name} with fields {field_names}")