add alloc for only i64

This commit is contained in:
2025-10-24 02:06:39 +05:30
parent 36a1a0903e
commit 64674cf646
7 changed files with 79 additions and 20 deletions

View File

@ -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]

View File

@ -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)

View File

@ -20,7 +20,7 @@ from pythonbpf.assign_pass import (
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")

View File

@ -98,3 +98,23 @@ class VmlinuxHandler:
# 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")

View File

@ -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)

View File

@ -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)

View File

@ -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()