add assignment dict handling to class_handler.py

Author: varun-r-mallya

pythonbpf/codegen.py

@@ -55,11 +55,11 @@ def processor(source_code, filename, module):
     for func_node in bpf_chunks:
         logger.info(f"Found BPF function/struct: {func_node.name}")
 
-    vmlinux_proc(tree, module)
+    vmlinux_symtab = vmlinux_proc(tree, module)
     populate_global_symbol_table(tree, module)
     license_processing(tree, module)
     globals_processing(tree, module)
-
+    print("DEBUG:", vmlinux_symtab)
     structs_sym_tab = structs_proc(tree, module, bpf_chunks)
     map_sym_tab = maps_proc(tree, module, bpf_chunks)
     func_proc(tree, module, bpf_chunks, map_sym_tab, structs_sym_tab)

pythonbpf/vmlinux_parser/class_handler.py

@@ -2,9 +2,9 @@
 from functools import lru_cache
 import importlib
 
-from .assignment_info import AssignmentInfo
+from .assignment_info import AssignmentInfo, AssignmentType
 from .dependency_handler import DependencyHandler
-from .dependency_node import DependencyNode
+from .dependency_node import DependencyNode, Field
 import ctypes
 from typing import Optional, Any, Dict
 
@@ -103,12 +103,21 @@ def process_vmlinux_post_ast(
             else:
                 raise TypeError("Could not get required class and definition")
 
+            # Create a members dictionary for AssignmentInfo
+            members_dict: Dict[str, tuple[str, Field]] = {}
+
             logger.debug(f"Extracted fields for {current_symbol_name}: {field_table}")
             for elem in field_table.items():
                 elem_name, elem_temp_list = elem
                 [elem_type, elem_bitfield_size] = elem_temp_list
                 local_module_name = getattr(elem_type, "__module__", None)
                 new_dep_node.add_field(elem_name, elem_type, ready=False)
+
+                # Store field reference for struct assignment info
+                field_ref = new_dep_node.get_field(elem_name)
+                if field_ref:
+                    members_dict[elem_name] = (elem_name, field_ref)
+
                 if local_module_name == ctypes.__name__:
                     # TODO: need to process pointer to ctype and also CFUNCTYPES here recursively. Current processing is a single dereference
                     new_dep_node.set_field_bitfield_size(elem_name, elem_bitfield_size)
@@ -220,6 +229,7 @@ def process_vmlinux_post_ast(
                                     containing_type,
                                     llvm_handler,
                                     handler,
+                                    assignments,  # Pass assignments to recursive call
                                     processing_stack,
                                 )
                                 new_dep_node.set_field_ready(elem_name, True)
@@ -237,18 +247,33 @@ def process_vmlinux_post_ast(
                             else str(elem_type)
                         )
                         process_vmlinux_post_ast(
-                            elem_type, llvm_handler, handler, processing_stack
+                            elem_type,
+                            llvm_handler,
+                            handler,
+                            assignments,
+                            processing_stack,
                         )
                         new_dep_node.set_field_ready(elem_name, True)
                 else:
                     raise ValueError(
                         f"{elem_name} with type {elem_type} from module {module_name} not supported in recursive resolver"
                     )
 
+            # Add struct to assignments dictionary
+            assignments[current_symbol_name] = AssignmentInfo(
+                value_type=AssignmentType.STRUCT,
+                python_type=elem_type_class,
+                value=None,
+                pointer_level=None,
+                signature=None,
+                members=members_dict,
+            )
+            logger.info(f"Added struct assignment info for {current_symbol_name}")
+
     else:
         raise ImportError("UNSUPPORTED Module")
 
-    logging.info(
+    logger.info(
         f"{current_symbol_name} processed and handler readiness {handler.is_ready}"
     )
     return True