From 99aacca94b88a2520bd8c774320f199c2ace8dc8 Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyanshchaturvedi18@gmail.com>
Date: Fri, 10 Oct 2025 13:48:40 +0530
Subject: [PATCH] WIP: allow pointer assignments to var

---
 pythonbpf/assign_pass.py              | 19 ++++++++++++++++++-
 pythonbpf/expr/__init__.py            |  4 ++--
 pythonbpf/expr/type_normalization.py  |  6 +++---
 pythonbpf/functions/functions_pass.py |  6 +++---
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/pythonbpf/assign_pass.py b/pythonbpf/assign_pass.py
index 32753ff..93c8003 100644
--- a/pythonbpf/assign_pass.py
+++ b/pythonbpf/assign_pass.py
@@ -1,7 +1,7 @@
 import ast
 import logging
 from llvmlite import ir
-from pythonbpf.expr import eval_expr
+from pythonbpf.expr import eval_expr, get_base_type_and_depth
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +38,9 @@ def handle_variable_assignment(
 
     val, val_type = val_result
     if val_type != var_type:
+        logger.info(f"val = {val}")
+        logger.info(f"var = {var_ptr}")
+        logger.info(f"truthy {var_type}")
         if isinstance(val_type, ir.IntType) and isinstance(var_type, ir.IntType):
             # Allow implicit int widening
             if val_type.width < var_type.width:
@@ -46,10 +49,24 @@ def handle_variable_assignment(
             elif val_type.width > var_type.width:
                 val = builder.trunc(val, var_type)
                 logger.info(f"Implicitly truncated int for variable {var_name}")
+        elif isinstance(val_type, ir.IntType) and isinstance(var_type, ir.PointerType):
+            ptr_target, ptr_depth = get_base_type_and_depth(var_type)
+            if ptr_target.width > val_type.width:
+                val = builder.sext(val, ptr_target)
+            elif ptr_target.width < val_type.width:
+                val = builder.trunc(val, ptr_target)
+
+            if ptr_depth > 1:
+                # NOTE: This is assignment to a PTR_TO_MAP_VALUE_OR_NULL
+                var_ptr_tmp = local_sym_tab[f"{var_name}_tmp"].var
+                builder.store(val, var_ptr_tmp)
+                val = var_ptr_tmp
         else:
             logger.error(
                 f"Type mismatch for variable {var_name}: {val_type} vs {var_type}"
             )
+            logger.error(f"var_type: {isinstance(var_type, ir.PointerType)}")
+            logger.error(f"val_type: {isinstance(val_type, ir.IntType)}")
             return False
 
     builder.store(val, var_ptr)
diff --git a/pythonbpf/expr/__init__.py b/pythonbpf/expr/__init__.py
index d58c543..577ee3c 100644
--- a/pythonbpf/expr/__init__.py
+++ b/pythonbpf/expr/__init__.py
@@ -1,4 +1,4 @@
 from .expr_pass import eval_expr, handle_expr
-from .type_normalization import convert_to_bool
+from .type_normalization import convert_to_bool, get_base_type_and_depth
 
-__all__ = ["eval_expr", "handle_expr", "convert_to_bool"]
+__all__ = ["eval_expr", "handle_expr", "convert_to_bool", "get_base_type_and_depth"]
diff --git a/pythonbpf/expr/type_normalization.py b/pythonbpf/expr/type_normalization.py
index 7a2fb57..34b4fb7 100644
--- a/pythonbpf/expr/type_normalization.py
+++ b/pythonbpf/expr/type_normalization.py
@@ -16,7 +16,7 @@ COMPARISON_OPS = {
 }
 
 
-def _get_base_type_and_depth(ir_type):
+def get_base_type_and_depth(ir_type):
     """Get the base type for pointer types."""
     cur_type = ir_type
     depth = 0
@@ -88,8 +88,8 @@ def _normalize_types(func, builder, lhs, rhs):
         logger.error(f"Type mismatch: {lhs.type} vs {rhs.type}")
         return None, None
     else:
-        lhs_base, lhs_depth = _get_base_type_and_depth(lhs.type)
-        rhs_base, rhs_depth = _get_base_type_and_depth(rhs.type)
+        lhs_base, lhs_depth = get_base_type_and_depth(lhs.type)
+        rhs_base, rhs_depth = get_base_type_and_depth(rhs.type)
         if lhs_base == rhs_base:
             if lhs_depth < rhs_depth:
                 rhs = _deref_to_depth(func, builder, rhs, rhs_depth - lhs_depth)
diff --git a/pythonbpf/functions/functions_pass.py b/pythonbpf/functions/functions_pass.py
index a1414ea..41d4048 100644
--- a/pythonbpf/functions/functions_pass.py
+++ b/pythonbpf/functions/functions_pass.py
@@ -461,8 +461,8 @@ def allocate_mem(
                     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")
+                    tmp_ir_type = ir.IntType(64)
+                    var_tmp = builder.alloca(tmp_ir_type, name=f"{var_name}_tmp")
                     double_alloc = True
                     # var.align = ir_type.width // 8
                     logger.info(
@@ -507,7 +507,7 @@ def allocate_mem(
                 local_sym_tab[var_name] = LocalSymbol(var, ir_type)
 
             if double_alloc:
-                local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, ir_type)
+                local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, tmp_ir_type)
     return local_sym_tab