fix struct usage in functions_pass

Author: r41k0u

pythonbpf/functions_pass.py

@@ -49,21 +49,17 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
             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][0]
-                field_ptr = builder.gep(
-                    struct_ptr, [ir.Constant(ir.IntType(32), 0),
-                                 ir.Constant(ir.IntType(32), field_idx)],
-                    inbounds=True)
+            if field_name in struct_info.fields:
+                field_ptr = struct_info.gep(
+                    builder, local_sym_tab[var_name][0], field_name)
                 val = eval_expr(func, module, builder, rval,
                                 local_sym_tab, map_sym_tab, structs_sym_tab)
-                if isinstance(struct_info["field_types"][field_idx], ir.ArrayType) and val[1] == ir.PointerType(ir.IntType(8)):
+                if isinstance(struct_info.field_type(field_name), ir.ArrayType) and val[1] == ir.PointerType(ir.IntType(8)):
                     # TODO: Figure it out, not a priority rn
                     # Special case for string assignment to char array
-                    #str_len = struct_info["field_types"][field_idx].count
-                    #assign_string_to_array(builder, field_ptr, val[0], str_len)
-                    #print(f"Assigned to struct field {var_name}.{field_name}")
+                    # str_len = struct_info["field_types"][field_idx].count
+                    # assign_string_to_array(builder, field_ptr, val[0], str_len)
+                    # print(f"Assigned to struct field {var_name}.{field_name}")
                     pass
                 if val is None:
                     print("Failed to evaluate struct field assignment")
@@ -138,7 +134,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                 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"]
+                ir_type = struct_info.ir_type
                 # var = builder.alloca(ir_type, name=var_name)
                 # Null init
                 builder.store(ir.Constant(ir_type, None),
@@ -364,7 +360,7 @@ def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_t
                             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"]
+                        ir_type = struct_info.ir_type
                         var = builder.alloca(ir_type, name=var_name)
                         local_var_metadata[var_name] = call_type
                         print(
@@ -548,6 +544,8 @@ def _expr_type(e):
     return found_type or "None"
 
 # For string assignment to fixed-size arrays
+
+
 def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_length):
     """
     Copy a string (i8*) to a fixed-size array ([N x i8]*)
@@ -556,36 +554,39 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l
     entry_block = builder.block
     copy_block = builder.append_basic_block("copy_char")
     end_block = builder.append_basic_block("copy_end")
-    
+
     # Create loop counter
     i = builder.alloca(ir.IntType(32))
     builder.store(ir.Constant(ir.IntType(32), 0), i)
-    
+
     # Start the loop
     builder.branch(copy_block)
-    
+
     # Copy loop
     builder.position_at_end(copy_block)
     idx = builder.load(i)
-    in_bounds = builder.icmp_unsigned('<', idx, ir.Constant(ir.IntType(32), array_length))
+    in_bounds = builder.icmp_unsigned(
+        '<', idx, ir.Constant(ir.IntType(32), array_length))
     builder.cbranch(in_bounds, copy_block, end_block)
-    
+
     with builder.if_then(in_bounds):
         # Load character from source
         src_ptr = builder.gep(source_string_ptr, [idx])
         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
         next_idx = builder.add(idx, ir.Constant(ir.IntType(32), 1))
         builder.store(next_idx, i)
-    
+
     builder.position_at_end(end_block)
-    
+
     # 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)