From ff7663803d66e31803595e6d07a80de95477811b Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyansh.chaturvedi@canonical.com>
Date: Tue, 9 Sep 2025 01:04:17 +0530
Subject: [PATCH] implement bpf map lookup emission

---
 pythonbpf/bpf_helper_handler.py | 26 ++++++++++++++++++++++++++
 pythonbpf/functions_pass.py     | 11 +++++++----
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/pythonbpf/bpf_helper_handler.py b/pythonbpf/bpf_helper_handler.py
index 7850891..94efcbb 100644
--- a/pythonbpf/bpf_helper_handler.py
+++ b/pythonbpf/bpf_helper_handler.py
@@ -6,6 +6,32 @@ def bpf_ktime_get_ns_emitter(call, module, builder, func):
     pass
 
 
+def bpf_map_lookup_elem_emitter(map_ptr, key_ptr, module, builder):
+    """
+    Emit LLVM IR for bpf_map_lookup_elem helper function call.
+    """
+    # Cast pointers to void*
+    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*)
+        var_arg=False
+    )
+    fn_ptr_type = ir.PointerType(fn_type)
+
+    # Helper ID 1 is bpf_map_lookup_elem
+    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):
     if not hasattr(func, "_fmt_counter"):
         func._fmt_counter = 0
diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py
index 576bae1..3187096 100644
--- a/pythonbpf/functions_pass.py
+++ b/pythonbpf/functions_pass.py
@@ -1,7 +1,7 @@
 from llvmlite import ir
 import ast
 
-from .bpf_helper_handler import bpf_printk_emitter, bpf_ktime_get_ns_emitter
+from .bpf_helper_handler import bpf_printk_emitter, bpf_ktime_get_ns_emitter, bpf_map_lookup_elem_emitter
 from .type_deducer import ctypes_to_ir
 
 
@@ -79,28 +79,31 @@ def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
                             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
-                                key_val = builder.load(key_var)
                             elif key_arg.id in map_sym_tab:
                                 key_var = map_sym_tab[key_arg.id]
                                 key_type = key_var.type.pointee
-                                key_val = builder.load(key_var)
                             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}")
-                            print(f"Key val: {key_val}")
                         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")