From 84500305dbaaf25b55a37da70f589c38addad599 Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyansh.chaturvedi@canonical.com>
Date: Tue, 30 Sep 2025 00:36:45 +0530
Subject: [PATCH] Move structs_pass under structs, create StructType

---
 pythonbpf/structs/struct_type.py        | 29 +++++++++++++++++++++++++
 pythonbpf/{ => structs}/structs_pass.py | 22 +++++++++----------
 2 files changed, 40 insertions(+), 11 deletions(-)
 create mode 100644 pythonbpf/structs/struct_type.py
 rename pythonbpf/{ => structs}/structs_pass.py (85%)

diff --git a/pythonbpf/structs/struct_type.py b/pythonbpf/structs/struct_type.py
new file mode 100644
index 0000000..129066f
--- /dev/null
+++ b/pythonbpf/structs/struct_type.py
@@ -0,0 +1,29 @@
+from llvmlite import ir
+
+
+class StructType:
+    def __init__(self, ir_type, fields, size):
+        self.ir_type = ir_type
+        self.fields = fields
+        self.size = size
+
+    def field_idx(self, field_name):
+        return self.fields.keys().index(field_name)
+
+    def field_type(self, field_name):
+        return self.fields[field_name]
+
+    def gep(self, builder, ptr, field_name):
+        idx = self.field_idx(field_name)
+        return builder.gep(ptr, [ir.Constant(ir.IntType(32), 0),
+                                 ir.Constant(ir.IntType(32), idx)],
+                           inbounds=True)
+
+    def field_size(self, field_name):
+        fld = self.fields[field_name]
+        if isinstance(fld, ir.ArrayType):
+            return fld.element.count * (fld.element.width // 8)
+        elif isinstance(fld, ir.IntType):
+            return fld.width // 8
+        elif isinstance(fld, ir.PointerType):
+            return 8
diff --git a/pythonbpf/structs_pass.py b/pythonbpf/structs/structs_pass.py
similarity index 85%
rename from pythonbpf/structs_pass.py
rename to pythonbpf/structs/structs_pass.py
index 4ba3869..6557e68 100644
--- a/pythonbpf/structs_pass.py
+++ b/pythonbpf/structs/structs_pass.py
@@ -34,32 +34,29 @@ def is_bpf_struct(cls_node):
 def process_bpf_struct(cls_node, module):
     """ Process a single BPF struct definition """
 
-    field_names, field_types = parse_struct_fields(cls_node)
-    total_size = calc_struct_size(field_types)
-    struct_type = ir.LiteralStructType(field_types)
-    logger.info(f"Created struct {cls_node.name} with fields {field_names}")
+    fields = parse_struct_fields(cls_node)
+    total_size = calc_struct_size(fields.values())
+    struct_type = ir.LiteralStructType(fields.values())
+    logger.info(f"Created struct {cls_node.name} with fields {fields.keys()}")
     return {
         "type": struct_type,
-        "fields": {name: idx for idx, name in enumerate(field_names)},
+        "fields": fields,
         "size": total_size,
-        "field_types": field_types,
     }
 
 
 def parse_struct_fields(cls_node):
     """ Parse fields of a struct class node """
-    field_names = []
-    field_types = []
+    fields = {}
 
     for item in cls_node.body:
         if isinstance(item, ast.AnnAssign) and \
            isinstance(item.target, ast.Name):
-            field_names.append(item.target.id)
-            field_types.append(get_type_from_ann(item.annotation))
+            fields[item.target.id] = get_type_from_ann(item.annotation)
         else:
             logger.error(f"Unsupported struct field: {ast.dump(item)}")
             raise TypeError(f"Unsupported field in {ast.dump(cls_node)}")
-    return field_names, field_types
+    return fields
 
 
 def get_type_from_ann(annotation):
@@ -67,10 +64,12 @@ def get_type_from_ann(annotation):
     if isinstance(annotation, ast.Call) and \
        isinstance(annotation.func, ast.Name):
         if annotation.func.id == "str":
+            # Char array
             # Assumes constant integer argument
             length = annotation.args[0].value
             return ir.ArrayType(ir.IntType(8), length)
     elif isinstance(annotation, ast.Name):
+        # Int type, written as c_int64, c_uint32, etc.
         return ctypes_to_ir(annotation.id)
 
     raise TypeError(f"Unsupported annotation type: {ast.dump(annotation)}")
@@ -87,6 +86,7 @@ def calc_struct_size(field_types):
             fsize = ftype.count * (ftype.element.width // 8)
             alignment = ftype.element.width // 8
         elif isinstance(ftype, ir.PointerType):
+            # We won't encounter this rn, but for the future
             fsize = 8
             alignment = 8
         else: