Move structs_pass under structs, create StructType

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

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)
+    fields = parse_struct_fields(cls_node)
-    total_size = calc_struct_size(field_types)
+    total_size = calc_struct_size(fields.values())
-    struct_type = ir.LiteralStructType(field_types)
+    struct_type = ir.LiteralStructType(fields.values())
-    logger.info(f"Created struct {cls_node.name} with fields {field_names}")
+    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 = []
+    fields = {}
-    field_types = []
 
     for item in cls_node.body:
         if isinstance(item, ast.AnnAssign) and \
            isinstance(item.target, ast.Name):
-            field_names.append(item.target.id)
+            fields[item.target.id] = get_type_from_ann(item.annotation)
-            field_types.append(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: