Store LocalSymbol in allocate_mem

Author: r41k0u

pythonbpf/functions_pass.py

@@ -93,19 +93,16 @@ def handle_assign(
     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][0])
+                builder.store(ir.Constant(ir.IntType(1), 1), local_sym_tab[var_name][0])
             else:
-                builder.store(ir.Constant(ir.IntType(1), 0),
-                              local_sym_tab[var_name][0])
+                builder.store(ir.Constant(ir.IntType(1), 0), 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][0]
+                ir.Constant(ir.IntType(64), rval.value), local_sym_tab[var_name][0]
             )
             # local_sym_tab[var_name] = var
             print(f"Assigned constant {rval.value} to {var_name}")
@@ -120,8 +117,7 @@ def handle_assign(
             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)))
+            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:
@@ -140,8 +136,7 @@ def handle_assign(
                 # 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][0]
+                    ir.Constant(ir_type, rval.args[0].value), local_sym_tab[var_name][0]
                 )
                 print(
                     f"Assigned {call_type} constant "
@@ -187,8 +182,7 @@ def handle_assign(
                 ir_type = struct_info.ir_type
                 # var = builder.alloca(ir_type, name=var_name)
                 # Null init
-                builder.store(ir.Constant(ir_type, None),
-                              local_sym_tab[var_name][0])
+                builder.store(ir.Constant(ir_type, None), 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
@@ -259,8 +253,7 @@ def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
             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)[0]
+        lhs = eval_expr(func, module, builder, cond.left, local_sym_tab, map_sym_tab)[0]
         if len(cond.ops) != 1 or len(cond.comparators) != 1:
             print("Unsupported complex comparison")
             return None
@@ -313,8 +306,7 @@ def handle_if(
     else:
         else_block = None
 
-    cond = handle_cond(func, module, builder, stmt.test,
-                       local_sym_tab, map_sym_tab)
+    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:
@@ -419,6 +411,7 @@ def allocate_mem(
     module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
 ):
     for stmt in body:
+        has_metadata = False
         if isinstance(stmt, ast.If):
             if stmt.body:
                 local_sym_tab = allocate_mem(
@@ -459,8 +452,7 @@ def allocate_mem(
                         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}")
+                        print(f"Pre-allocated variable {var_name} of type {call_type}")
                     elif HelperHandlerRegistry.has_handler(call_type):
                         # Assume return type is int64 for now
                         ir_type = ir.IntType(64)
@@ -477,7 +469,7 @@ def allocate_mem(
                         struct_info = structs_sym_tab[call_type]
                         ir_type = struct_info.ir_type
                         var = builder.alloca(ir_type, name=var_name)
-                        local_var_metadata[var_name] = call_type
+                        has_metadata = True
                         print(
                             f"Pre-allocated variable {var_name} "
                             f"for struct {call_type}"
@@ -519,7 +511,11 @@ def allocate_mem(
             else:
                 print("Unsupported assignment value type")
                 continue
-            local_sym_tab[var_name] = (var, ir_type)
+
+            if has_metadata:
+                local_sym_tab[var_name] = LocalSymbol(var, ir_type, call_type)
+            else:
+                local_sym_tab[var_name] = LocalSymbol(var, ir_type)
     return local_sym_tab
 
 
@@ -681,8 +677,7 @@ def _expr_type(e):
             if found_type is None:
                 found_type = t
             elif found_type != t:
-                raise ValueError("Conflicting return types:" f"{
-                                 found_type} vs {t}")
+                raise ValueError(f"Conflicting return types:{found_type} vs {t}")
     return found_type or "None"
 
 
@@ -719,8 +714,7 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l
         char = builder.load(src_ptr)
 
         # Store character in target
-        dst_ptr = builder.gep(
-            target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx])
+        dst_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx])
         builder.store(char, dst_ptr)
 
         # Increment counter
@@ -731,6 +725,5 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l
 
     # Ensure null termination
     last_idx = ir.Constant(ir.IntType(32), array_length - 1)
-    null_ptr = builder.gep(
-        target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx])
+    null_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx])
     builder.store(ir.Constant(ir.IntType(8), 0), null_ptr)