add assignments table and offset handler

Author: varun-r-mallya

pythonbpf/vmlinux_parser/class_handler.py

@@ -71,7 +71,9 @@ def process_vmlinux_post_ast(
                     if len(field_elem) == 2:
                         field_name, field_type = field_elem
                     elif len(field_elem) == 3:
-                        raise NotImplementedError("Bitfields are not supported in the current version")
+                        raise NotImplementedError(
+                            "Bitfields are not supported in the current version"
+                        )
                         field_name, field_type, bitfield_size = field_elem
                     field_table[field_name] = [field_type, bitfield_size]
             elif hasattr(class_obj, "__annotations__"):
@@ -145,7 +147,8 @@ def process_vmlinux_post_ast(
                             process_vmlinux_post_ast(
                                 containing_type, llvm_handler, handler, processing_stack
                             )
-                            new_dep_node.set_field_ready(elem_name, True)
+                            size_of_containing_type  = (handler[containing_type.__name__]).__sizeof__()
+                            new_dep_node.set_field_ready(elem_name, True, size_of_containing_type)
                         elif containing_type.__module__ == ctypes.__name__:
                             logger.debug(f"Processing ctype internal{containing_type}")
                             new_dep_node.set_field_ready(elem_name, True)
@@ -162,7 +165,8 @@ def process_vmlinux_post_ast(
                         process_vmlinux_post_ast(
                             elem_type, llvm_handler, handler, processing_stack
                         )
-                        new_dep_node.set_field_ready(elem_name, True)
+                        size_of_containing_type = (handler[elem_type.__name__]).__sizeof__()
+                        new_dep_node.set_field_ready(elem_name, True, size_of_containing_type)
                 else:
                     raise ValueError(
                         f"{elem_name} with type {elem_type} from module {module_name} not supported in recursive resolver"

pythonbpf/vmlinux_parser/dependency_node.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
 from typing import Dict, Any, Optional
+import ctypes
 
 
 # TODO: FIX THE FUCKING TYPE NAME CONVENTION.
@@ -140,11 +141,14 @@ def add_field(
             type_size=type_size,
             ctype_complex_type=ctype_complex_type,
             bitfield_size=bitfield_size,
-            offset=offset
+            offset=offset,
         )
         # Invalidate readiness cache
         self._ready_cache = None
 
+    def __sizeof__(self):
+        return self.current_offset
+
     def get_field(self, name: str) -> Field:
         """Get a field by name."""
         return self.fields[name]
@@ -211,20 +215,53 @@ def set_field_bitfield_size(
         # Invalidate readiness cache
         self._ready_cache = None
 
-    def set_field_ready(self, name: str, is_ready: bool = False) -> None:
+    def set_field_ready(self, name: str, is_ready: bool = False, size_of_containing_type: Optional[int] = None) -> None:
         """Mark a field as ready or not ready."""
         if name not in self.fields:
             raise KeyError(f"Field '{name}' does not exist in node '{self.name}'")
 
         self.fields[name].set_ready(is_ready)
         self.fields[name].set_offset(self.current_offset)
-        self.current_offset += self._calculate_size(name)
-
+        self.current_offset += self._calculate_size(name, size_of_containing_type)
         # Invalidate readiness cache
         self._ready_cache = None
 
-    def _calculate_size(self, name: str) -> int:
-        pass
+    def _calculate_size(self, name: str, size_of_containing_type: Optional[int] = None) -> int:
+        processing_field = self.fields[name]
+        # size_of_field will be in bytes
+        if processing_field.type.__module__ == ctypes.__name__:
+            size_of_field = ctypes.sizeof(processing_field.type)
+            return size_of_field
+        elif processing_field.type.__module__ == "vmlinux":
+            size_of_field: int = 0
+            if processing_field.ctype_complex_type is not None:
+                if issubclass(processing_field.ctype_complex_type, ctypes.Array):
+                    if processing_field.containing_type.__module__ == ctypes.__name__:
+                        size_of_field = (
+                            ctypes.sizeof(processing_field.containing_type)
+                            * processing_field.type_size
+                        )
+                        return size_of_field
+                    elif processing_field.containing_type.__module__ == "vmlinux":
+                        size_of_field = (
+                            size_of_containing_type
+                            * processing_field.type_size
+                        )
+                        return size_of_field
+                elif issubclass(processing_field.ctype_complex_type, ctypes._Pointer):
+                    return ctypes.sizeof(ctypes.pointer())
+                else:
+                    raise NotImplementedError(
+                        "This subclass of ctype not supported yet"
+                    )
+            else:
+                # search up pre-created stuff and get size
+                return size_of_containing_type
+
+        else:
+            raise ModuleNotFoundError("Module is not supported for the operation")
+        raise RuntimeError("control should not reach here")
+
     @property
     def is_ready(self) -> bool:
         """Check if the node is ready (all fields are ready)."""

pythonbpf/vmlinux_parser/import_detector.py

@@ -129,7 +129,13 @@ def vmlinux_proc(tree: ast.AST, module):
                 )
 
     IRGenerator(module, handler)
+    return assignments
 
 
 def process_vmlinux_assign(node, module, assignments: Dict[str, type]):
-    raise NotImplementedError("Assignment handling has not been implemented yet")
+    # Check if this is a simple assignment with a constant value
+    if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
+        target_name = node.targets[0].id
+        if isinstance(node.value, ast.Constant):
+            assignments[target_name] = node.value.value
+            logger.info(f"Added assignment: {target_name} = {node.value.value}")

pythonbpf/vmlinux_parser/ir_gen/ir_generation.py

@@ -30,6 +30,7 @@ def struct_processor(self, struct):
             # this part cannot yet resolve circular dependencies. Gets stuck on an infinite loop during that.
             self.generated.append(struct.name)
 
-
-    def struct_name_generator(self, ):
+    def struct_name_generator(
+        self,
+    ) -> None:
         pass

tests/failing_tests/xdp_pass.py

@@ -1,8 +1,9 @@
 from pythonbpf import bpf, map, section, bpfglobal, compile_to_ir
 from pythonbpf.maps import HashMap
 from pythonbpf.helper import XDP_PASS
+from vmlinux import TASK_COMM_LEN  # noqa: F401
+from vmlinux import struct_trace_event_raw_sys_enter  # noqa: F401
 # from vmlinux import struct_request
-from vmlinux import struct_trace_event_raw_sys_enter
 from vmlinux import struct_xdp_md
 # from vmlinux import struct_trace_event_raw_sys_enter  # noqa: F401
 # from vmlinux import struct_ring_buffer_per_cpu  # noqa: F401