From 047f361ea91024877ebc101d6a48eabcf0e5b855 Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyanshchaturvedi18@gmail.com>
Date: Fri, 10 Oct 2025 06:09:46 +0530
Subject: [PATCH] Allocate twice for map lookups

---
 pythonbpf/binary_ops.py               |  6 ++++--
 pythonbpf/functions/functions_pass.py | 12 +++++++++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/pythonbpf/binary_ops.py b/pythonbpf/binary_ops.py
index 7fa8969..77c184f 100644
--- a/pythonbpf/binary_ops.py
+++ b/pythonbpf/binary_ops.py
@@ -55,10 +55,12 @@ def store_through_chain(value, chain, builder):
 
 def handle_binary_op_impl(rval, builder, local_sym_tab):
     op = rval.op
-    left, _, _ = get_operand_value(rval.left, builder, local_sym_tab)
-    right, _, _ = get_operand_value(rval.right, builder, local_sym_tab)
+    left, lchain, _ = get_operand_value(rval.left, builder, local_sym_tab)
+    right, rchain, _ = get_operand_value(rval.right, builder, local_sym_tab)
     logger.info(f"left is {left}, right is {right}, op is {op}")
 
+    logger.info(f"left chain: {lchain}, right chain: {rchain}")
+
     # Map AST operation nodes to LLVM IR builder methods
     op_map = {
         ast.Add: builder.add,
diff --git a/pythonbpf/functions/functions_pass.py b/pythonbpf/functions/functions_pass.py
index d29de2f..6c2de1a 100644
--- a/pythonbpf/functions/functions_pass.py
+++ b/pythonbpf/functions/functions_pass.py
@@ -455,10 +455,17 @@ def allocate_mem(
                             f"Pre-allocated variable {var_name} for struct {call_type}"
                         )
                 elif isinstance(rval.func, ast.Attribute):
+                    # Map method call
                     ir_type = ir.PointerType(ir.IntType(64))
                     var = builder.alloca(ir_type, name=var_name)
+
+                    # declare an intermediate ptr type for map lookup
+                    ir_type = ir.IntType(64)
+                    var_tmp = builder.alloca(ir_type, name=f"{var_name}_tmp")
                     # var.align = ir_type.width // 8
-                    logger.info(f"Pre-allocated variable {var_name} for map")
+                    logger.info(
+                        f"Pre-allocated variable {var_name} and {var_name}_tmp for map"
+                    )
                 else:
                     logger.info("Unsupported assignment call function type")
                     continue
@@ -496,6 +503,9 @@ def allocate_mem(
                 local_sym_tab[var_name] = LocalSymbol(var, ir_type, call_type)
             else:
                 local_sym_tab[var_name] = LocalSymbol(var, ir_type)
+
+            if var_tmp:
+                local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, ir_type)
     return local_sym_tab