From 083ee21e38e0eb40d26461d730afff46c82b27bb Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyansh.chaturvedi@canonical.com>
Date: Mon, 29 Sep 2025 22:50:46 +0530
Subject: [PATCH] structs_pass cleanup

---
 pythonbpf/structs_pass.py | 44 ++++++++++++++++++++++++++-------------
 1 file changed, 29 insertions(+), 15 deletions(-)

diff --git a/pythonbpf/structs_pass.py b/pythonbpf/structs_pass.py
index 84394e3..382b0fa 100644
--- a/pythonbpf/structs_pass.py
+++ b/pythonbpf/structs_pass.py
@@ -6,6 +6,7 @@ structs_sym_tab = {}
 
 
 def structs_proc(tree, module, chunks):
+    """ Process all class definitions to find BPF structs """
     for cls_node in chunks:
         # Check if this class is a struct
         is_struct = False
@@ -21,22 +22,35 @@ def structs_proc(tree, module, chunks):
 
 
 def process_bpf_struct(cls_node, module):
-    struct_name = cls_node.name
+    """ Process a single BPF struct definition """
+
     field_names = []
     field_types = []
 
     for item in cls_node.body:
-        if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
-            print(f"Field: {item.target.id}, Type: "
-                  f"{ast.dump(item.annotation)}")
-            field_names.append(item.target.id)
-            if isinstance(item.annotation, ast.Call) and isinstance(item.annotation.func, ast.Name) and item.annotation.func.id == "str":
-                # This is a char array with fixed length
-                # TODO: For now assuming str is always called with constant
-                field_types.append(ir.ArrayType(
-                    ir.IntType(8), item.annotation.args[0].value))
-            else:
-                field_types.append(ctypes_to_ir(item.annotation.id))
+        #
+        # field syntax:
+        # class struct_example:
+        #     num: c_uint64
+        #
+        if isinstance(item, ast.AnnAssign):
+            if isinstance(item.target, ast.Name):
+                print(f"Field: {item.target.id}, Type: "
+                      f"{ast.dump(item.annotation)}")
+                field_names.append(item.target.id)
+                if isinstance(item.annotation, ast.Call):
+                    if isinstance(item.annotation.func, ast.Name):
+                        if item.annotation.func.id == "str":
+                            # This is a char array with fixed length
+                            # TODO: For now assume str is always with constant
+                            field_types.append(ir.ArrayType(
+                                ir.IntType(8), item.annotation.args[0].value))
+                        else:
+                            field_types.append(
+                                ctypes_to_ir(item.annotation.id))
+        else:
+            print(f"Unsupported struct field: {ast.dump(item)}")
+            return
 
     curr_offset = 0
     for ftype in field_types:
@@ -50,7 +64,7 @@ def process_bpf_struct(cls_node, module):
             fsize = 8
             alignment = 8
         else:
-            print(f"Unsupported field type in struct {struct_name}")
+            print(f"Unsupported field type in struct {cls_node.name}")
             return
         padding = (alignment - (curr_offset % alignment)) % alignment
         curr_offset += padding
@@ -59,10 +73,10 @@ def process_bpf_struct(cls_node, module):
     total_size = curr_offset + final_padding
 
     struct_type = ir.LiteralStructType(field_types)
-    structs_sym_tab[struct_name] = {
+    structs_sym_tab[cls_node.name] = {
         "type": struct_type,
         "fields": {name: idx for idx, name in enumerate(field_names)},
         "size": total_size,
         "field_types": field_types,
     }
-    print(f"Created struct {struct_name} with fields {field_names}")
+    print(f"Created struct {cls_node.name} with fields {field_names}")