find global variable ir and field data from metadata

pythonbpf/allocation_pass.py

@@ -2,27 +2,15 @@ import ast
 import logging
 
 from llvmlite import ir
-from dataclasses import dataclass
-from typing import Any
+from .local_symbol import LocalSymbol
 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
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class LocalSymbol:
-    var: ir.AllocaInstr
-    ir_type: ir.Type
-    metadata: Any = None
-
-    def __iter__(self):
-        yield self.var
-        yield self.ir_type
-        yield self.metadata
-
-
 def create_targets_and_rvals(stmt):
     """Create lists of targets and right-hand values from an assignment statement."""
     if isinstance(stmt.targets[0], ast.Tuple):
@@ -60,21 +48,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 +226,46 @@ 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/expr_pass.py

@@ -72,20 +72,28 @@ def _handle_attribute_expr(
         if var_name in local_sym_tab:
             var_ptr, var_type, var_metadata = local_sym_tab[var_name]
             logger.info(f"Loading attribute {attr_name} from variable {var_name}")
-            logger.info(f"Variable type: {var_type}, Variable ptr: {var_ptr}")
+            logger.info(
+                f"Variable type: {var_type}, Variable ptr: {var_ptr}, Variable Metadata: {var_metadata}"
+            )
+            if (
+                hasattr(var_metadata, "__module__")
+                and var_metadata.__module__ == "vmlinux"
+            ):
+                # Try vmlinux handler when var_metadata is not a string, but has a module attribute.
+                # This has been done to keep everything separate in vmlinux struct handling.
+                vmlinux_result = VmlinuxHandlerRegistry.handle_attribute(
+                    expr, local_sym_tab, None, builder
+                )
+                if vmlinux_result is not None:
+                    return vmlinux_result
+                else:
+                    raise RuntimeError("Vmlinux struct did not process successfully")
             metadata = structs_sym_tab[var_metadata]
             if attr_name in metadata.fields:
                 gep = metadata.gep(builder, var_ptr, attr_name)
                 val = builder.load(gep)
                 field_type = metadata.field_type(attr_name)
                 return val, field_type
-
-        # Try vmlinux handler as fallback
-        vmlinux_result = VmlinuxHandlerRegistry.handle_attribute(
-            expr, local_sym_tab, None, builder
-        )
-        if vmlinux_result is not None:
-            return vmlinux_result
     return None
 
 

pythonbpf/expr/vmlinux_registry.py

@@ -52,3 +52,18 @@ class VmlinuxHandlerRegistry:
         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

@@ -21,6 +21,7 @@ from pythonbpf.allocation_pass import (
     handle_assign_allocation,
     allocate_temp_pool,
     create_targets_and_rvals,
+    LocalSymbol,
 )
 
 from .return_utils import handle_none_return, handle_xdp_return, is_xdp_name
@@ -347,16 +348,9 @@ def process_func_body(
                 resolved_type = VmlinuxHandlerRegistry.get_struct_type(
                     context_type_name
                 )
-                context_type = {"type": ir.PointerType(resolved_type), "ptr": True}
-            else:
-                try:
-                    resolved_type = ctypes_to_ir(context_type_name)
-                    context_type = {"type": ir.PointerType(resolved_type), "ptr": True}
-                except Exception:
-                    raise TypeError(f"Type '{context_type_name}' not declared")
-
-            local_sym_tab[context_name] = context_type
-            logger.info(f"Added argument '{context_name}' to local symbol table")
+                context_type = LocalSymbol(None, None, resolved_type)
+                local_sym_tab[context_name] = context_type
+                logger.info(f"Added argument '{context_name}' to local symbol table")
 
     # pre-allocate dynamic variables
     local_sym_tab = allocate_mem(

pythonbpf/local_symbol.py

@@ -0,0 +1,15 @@
+import llvmlite.ir as ir
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class LocalSymbol:
+    var: ir.AllocaInstr
+    ir_type: ir.Type
+    metadata: Any = None
+
+    def __iter__(self):
+        yield self.var
+        yield self.ir_type
+        yield self.metadata

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -1,6 +1,7 @@
 import logging
 from llvmlite import ir
 
+from pythonbpf.local_symbol import LocalSymbol
 from pythonbpf.vmlinux_parser.assignment_info import AssignmentType
 
 logger = logging.getLogger(__name__)
@@ -87,14 +88,37 @@ class VmlinuxHandler:
         self, struct_var_name, field_name, module, builder, local_sym_tab
     ):
         """Handle access to vmlinux struct fields"""
-        # Check if it's a variable of vmlinux struct type
         if struct_var_name in local_sym_tab:
-            var_info = local_sym_tab[struct_var_name]  # noqa: F841
-            # Need to check if this variable is a vmlinux struct
-            # This will depend on how you track vmlinux struct types in your symbol table
+            var_info: LocalSymbol = local_sym_tab[struct_var_name]
             logger.info(
                 f"Attempting to access field {field_name} of possible vmlinux struct {struct_var_name}"
             )
+            python_type: type = var_info.metadata
+            globvar_ir, field_data = self.get_field_type(
+                python_type.__name__, field_name
+            )
+
             # Return pointer to field and field type
             return None
-        return None
+        else:
+            raise RuntimeError("Variable accessed not found in symbol table")
+
+    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/c-form/struct_field_tests.bpf.c

@@ -19,17 +19,8 @@ There is no point of
 SEC("tp/syscalls/sys_enter_execve")
 int handle_setuid_entry(struct trace_event_raw_sys_enter *ctx) {
   // Access each argument separately with clear variable assignments
-  unsigned long arg0 = ctx->args[0];
-  bpf_printk("args[0]: %u", arg0);
-
-  unsigned long arg1 = ctx->args[1];
-  bpf_printk("args[1]: %u", arg1);
-
-  // Remove the duplicate access to args[1]
-
-  unsigned long arg2 = ctx->args[2];
-  bpf_printk("args[3]: %u", arg2);
-  bpf_printk("args[4]: %u", ctx->args[2]);
+  long int arg0 = ctx->id;
+  bpf_printk("args[0]: %d", arg0);
 
   return 0;
 }

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
@@ -71,4 +71,5 @@ def LICENSE() -> str:
     return "GPL"
 
 
+compile_to_ir("comprehensive.py", "comprehensive.ll")
 compile()