add alloc for only i64

Author: varun-r-mallya

pythonbpf/allocation_pass.py

@@ -1,10 +1,12 @@
 import ast
+import ctypes
 import logging
 
 from llvmlite import ir
 from dataclasses import dataclass
 from typing import Any
 from pythonbpf.helper import HelperHandlerRegistry
+from pythonbpf.vmlinux_parser.dependency_node import Field
 from .expr import VmlinuxHandlerRegistry
 from pythonbpf.type_deducer import ctypes_to_ir
 
@@ -60,21 +62,11 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
             continue
 
         var_name = target.id
-
         # Skip if already allocated
         if var_name in local_sym_tab:
             logger.debug(f"Variable {var_name} already allocated, skipping")
             continue
 
-        # When allocating a variable, check if it's a vmlinux struct type
-        if isinstance(
-            stmt.value, ast.Name
-        ) and VmlinuxHandlerRegistry.is_vmlinux_struct(stmt.value.id):
-            # Handle vmlinux struct allocation
-            # This requires more implementation
-            print(stmt.value)
-            pass
-
         # Determine type and allocate based on rval
         if isinstance(rval, ast.Call):
             _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab)
@@ -248,9 +240,40 @@ def _allocate_for_attribute(builder, var_name, rval, local_sym_tab, structs_sym_
         logger.error(f"Struct variable '{struct_var}' not found")
         return
 
-    struct_type = local_sym_tab[struct_var].metadata
+    struct_type: type = local_sym_tab[struct_var].metadata
     if not struct_type or struct_type not in structs_sym_tab:
-        logger.error(f"Struct type '{struct_type}' not found")
+        if VmlinuxHandlerRegistry.is_vmlinux_struct(struct_type.__name__):
+            # Handle vmlinux struct field access
+            vmlinux_struct_name = struct_type.__name__
+            if not VmlinuxHandlerRegistry.has_field(vmlinux_struct_name, field_name):
+                logger.error(f"Field '{field_name}' not found in vmlinux struct '{vmlinux_struct_name}'")
+                return
+
+            field_type: tuple[ir.GlobalVariable, Field] = VmlinuxHandlerRegistry.get_field_type(vmlinux_struct_name, field_name)
+            field_ir, field = field_type
+            #TODO: For now, we only support integer type allocations.
+
+            # loaded_value = builder.load(field_ir, align=8)
+            # #TODO: fatal flaw that this always assumes first argument of function to be the context of what this gets.
+            # base_ptr = builder.function.args[0]
+            # gep_result = builder.gep(
+            #     base_ptr,
+            #     [loaded_value],
+            #     inbounds=False,  # Not using inbounds GEP
+            # )
+            # print("DEBB", loaded_value, base_ptr, gep_result)
+            # Use i64 for allocation since that's what the global variable contains
+
+            actual_ir_type = ir.IntType(64)
+
+            # Allocate with the actual IR type, not the GlobalVariable
+            var = _allocate_with_type(builder, var_name, actual_ir_type)
+            local_sym_tab[var_name] = LocalSymbol(var, actual_ir_type, field)
+
+            logger.info(f"Pre-allocated {var_name} from vmlinux struct {vmlinux_struct_name}.{field_name}")
+            return
+        else:
+            logger.error(f"Struct type '{struct_type}' not found")
         return
 
     struct_info = structs_sym_tab[struct_type]

pythonbpf/expr/vmlinux_registry.py

@@ -52,3 +52,18 @@ def get_struct_type(cls, name):
         if cls._handler is None:
             return None
         return cls._handler.get_vmlinux_struct_type(name)
+
+    @classmethod
+    def has_field(cls, vmlinux_struct_name, field_name):
+        """Check if a vmlinux struct has a specific field"""
+        if cls._handler is None:
+            return False
+        return cls._handler.has_field(vmlinux_struct_name, field_name)
+
+    @classmethod
+    def get_field_type(cls, vmlinux_struct_name, field_name):
+        """Get the type of a field in a vmlinux struct"""
+        if cls._handler is None:
+            return None
+        assert isinstance(cls._handler, VmlinuxHandler)
+        return cls._handler.get_field_type(vmlinux_struct_name, field_name)

pythonbpf/functions/functions_pass.py

@@ -20,7 +20,7 @@
 from pythonbpf.allocation_pass import (
     handle_assign_allocation,
     allocate_temp_pool,
-    create_targets_and_rvals,
+    create_targets_and_rvals, LocalSymbol,
 )
 
 from .return_utils import handle_none_return, handle_xdp_return, is_xdp_name
@@ -347,11 +347,12 @@ def process_func_body(
                 resolved_type = VmlinuxHandlerRegistry.get_struct_type(
                     context_type_name
                 )
-                context_type = {"type": ir.PointerType(resolved_type), "ptr": True}
+                context_type = LocalSymbol(None, ir.PointerType(resolved_type), resolved_type)
             else:
                 try:
                     resolved_type = ctypes_to_ir(context_type_name)
-                    context_type = {"type": ir.PointerType(resolved_type), "ptr": True}
+                    logger.error("THIS SHOULD NOT HAPPEN. I THINK. PROBABLY.")
+                    context_type = LocalSymbol(None, ir.PointerType(resolved_type), resolved_type)
                 except Exception:
                     raise TypeError(f"Type '{context_type_name}' not declared")
 

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -98,3 +98,23 @@ def handle_vmlinux_struct_field(
             # Return pointer to field and field type
             return None
         return None
+
+    def has_field(self, struct_name, field_name):
+        """Check if a vmlinux struct has a specific field"""
+        if self.is_vmlinux_struct(struct_name):
+            python_type = self.vmlinux_symtab[struct_name]["python_type"]
+            return hasattr(python_type, field_name)
+        return False
+
+    def get_field_type(self, vmlinux_struct_name, field_name):
+        """Get the type of a field in a vmlinux struct"""
+        if self.is_vmlinux_struct(vmlinux_struct_name):
+            python_type = self.vmlinux_symtab[vmlinux_struct_name]["python_type"]
+            if hasattr(python_type, field_name):
+                return self.vmlinux_symtab[vmlinux_struct_name]["members"][field_name]
+            else:
+                raise ValueError(
+                    f"Field {field_name} not found in vmlinux struct {vmlinux_struct_name}"
+                )
+        else:
+            raise ValueError(f"{vmlinux_struct_name} is not a vmlinux struct")

tests/c-form/Makefile

@@ -1,5 +1,5 @@
 BPF_CLANG := clang
-CFLAGS := -O2 -emit-llvm -target bpf -c
+CFLAGS := -O0 -emit-llvm -target bpf -c
 
 SRC := $(wildcard *.bpf.c)
 LL := $(SRC:.bpf.c=.bpf.ll)

tests/failing_tests/vmlinux/struct_field_access.py

@@ -17,7 +17,7 @@ def hello_world(ctx: struct_trace_event_raw_sys_enter) -> c_int64:
     a = 2 + TASK_COMM_LEN + TASK_COMM_LEN
     b = ctx.id
     print(f"Hello, World{TASK_COMM_LEN} and {a}")
-    print(f"This is context field {b}")
+    # print(f"This is context field {b}")
     return c_int64(TASK_COMM_LEN + 2)
 
 

tests/passing_tests/assign/comprehensive.py

@@ -1,4 +1,4 @@
-from pythonbpf import bpf, map, section, bpfglobal, compile, struct
+from pythonbpf import bpf, map, section, bpfglobal, compile, struct, compile_to_ir
 from ctypes import c_void_p, c_int64, c_int32, c_uint64
 from pythonbpf.maps import HashMap
 from pythonbpf.helper import ktime
@@ -70,5 +70,5 @@ def hello_world(ctx: c_void_p) -> c_int64:
 def LICENSE() -> str:
     return "GPL"
 
-
-compile()
+compile_to_ir("comprehensive.py", "comprehensive.ll")
+# compile()