Add function to extract kernel-defined struct names from BTF file and update related interfaces.

Author: congwang-mk

src/btf_binary_parser.ml

@@ -38,6 +38,7 @@ external btf_type_by_id : btf_handle -> int -> (int * string * int * int * int)
 external btf_type_get_members : btf_handle -> int -> (string * int) array = "btf_type_get_members_stub"
 external btf_resolve_type : btf_handle -> int -> string = "btf_resolve_type_stub"
 external btf_extract_function_signatures : btf_handle -> string list -> (string * string) list = "btf_extract_function_signatures_stub"
+external btf_extract_kernel_struct_names : btf_handle -> string list = "btf_extract_kernel_struct_names_stub"
 external btf_free : btf_handle -> unit = "btf_free_stub"
 
 (** Parse BTF file and extract requested types using libbpf *)
@@ -180,4 +181,18 @@ let extract_kernel_function_signatures btf_path function_names =
   with
   | exn ->
       printf "Error extracting function signatures from BTF file %s: %s\n" btf_path (Printexc.to_string exn);
-      [] 
+      []
+
+(** Extract all kernel-defined struct names from BTF file.
+    @param btf_path Path to the binary BTF file
+    @return List of kernel struct names *)
+let extract_all_kernel_struct_names btf_path =
+  try
+    match btf_new_from_file btf_path with
+    | None -> []
+    | Some btf_handle ->
+        let struct_names = btf_extract_kernel_struct_names btf_handle in
+        btf_free btf_handle;
+        struct_names
+  with
+  | _ -> [] 

src/btf_binary_parser.mli

@@ -34,4 +34,9 @@ val parse_btf_file : string -> string list -> btf_type_info list
     @param btf_path Path to the binary BTF file
     @param function_names List of kernel function names to extract signatures for
     @return List of (function_name, signature) pairs *)
-val extract_kernel_function_signatures : string -> string list -> (string * string) list 
+val extract_kernel_function_signatures : string -> string list -> (string * string) list
+
+(** Extract all kernel-defined struct names from BTF file.
+    @param btf_path Path to the binary BTF file
+    @return List of kernel struct names *)
+val extract_all_kernel_struct_names : string -> string list 

src/btf_parser.ml

@@ -38,7 +38,7 @@ type program_template = {
 
 
 (** Check if a type name is a well-known kernel type *)
-let is_well_known_kernel_type = Kernel_types.is_well_known_ebpf_type
+let is_well_known_kernel_type ?btf_path = Kernel_types.is_well_known_ebpf_type ?btf_path
 
 (** Create hardcoded enum definitions for constants that can't be extracted from BTF *)
 let create_hardcoded_tc_action_enum () = {
@@ -82,7 +82,7 @@ let get_program_template prog_type btf_path =
           kind = bt.Btf_binary_parser.kind;
           size = bt.Btf_binary_parser.size;
           members = bt.Btf_binary_parser.members;
-          kernel_defined = is_well_known_kernel_type bt.Btf_binary_parser.name;
+          kernel_defined = is_well_known_kernel_type ?btf_path bt.Btf_binary_parser.name;
         }) binary_types
     | Some path -> failwith (sprintf "BTF file not found: %s" path)
     | None -> failwith "BTF file path is required. Use --btf-vmlinux-path option."
@@ -147,7 +147,7 @@ let get_tracepoint_program_template category_event btf_path =
           kind = bt.Btf_binary_parser.kind;
           size = bt.Btf_binary_parser.size;
           members = bt.Btf_binary_parser.members;
-          kernel_defined = is_well_known_kernel_type bt.Btf_binary_parser.name;
+          kernel_defined = is_well_known_kernel_type ?btf_path bt.Btf_binary_parser.name;
         }) binary_types
     | Some path -> failwith (sprintf "BTF file not found: %s" path)
     | None -> failwith "BTF file path is required for tracepoint extraction. Use --btf-vmlinux-path option."
@@ -310,7 +310,11 @@ let generate_kernelscript_source template project_name =
              let member_strings = List.map (fun (name, typ) ->
                sprintf "    %s: %s," name typ
              ) members in
-             sprintf "struct %s {\n%s\n}" type_info.name (String.concat "\n" member_strings)
+             (* Mark kernel-defined structs for eBPF-only usage *)
+             let kernel_marker = if type_info.kernel_defined then
+               "// @kernel_only - This struct is only for eBPF compilation, not userspace\n"
+             else "" in
+             sprintf "%sstruct %s {\n%s\n}" kernel_marker type_info.name (String.concat "\n" member_strings)
          | None ->
              sprintf "// %s type (placeholder)" type_info.name)
     | "enum" ->

src/btf_parser.mli

@@ -43,7 +43,7 @@ val get_kprobe_program_template : string -> string option -> program_template
 val get_tracepoint_program_template : string -> string option -> program_template
 
 (** Check if a type name is a well-known eBPF kernel type *)
-val is_well_known_kernel_type : string -> bool
+val is_well_known_kernel_type : ?btf_path:string -> string -> bool
 
 (** Extract struct_ops definitions from BTF and generate KernelScript code *)
 val extract_struct_ops_definitions : string option -> string list -> string list

src/btf_stubs.c

@@ -577,6 +577,40 @@ value btf_extract_function_signatures_stub(value btf_handle, value function_name
     CAMLreturn(result_list);
 }
 
+/* Extract all kernel-defined struct names from BTF */
+value btf_extract_kernel_struct_names_stub(value btf_handle) {
+    CAMLparam1(btf_handle);
+    CAMLlocal2(result, cons);
+    
+    struct btf *btf = btf_of_value(btf_handle);
+    if (!btf) {
+        CAMLreturn(Val_emptylist);
+    }
+    
+    result = Val_emptylist;
+    __u32 nr_types = btf__type_cnt(btf);
+    
+    /* Iterate through all BTF types */
+    for (__u32 i = 1; i < nr_types; i++) {
+        const struct btf_type *type = btf__type_by_id(btf, i);
+        if (!type) continue;
+        
+        /* Check if it's a struct type */
+        if (btf_kind(type) == BTF_KIND_STRUCT) {
+            const char *type_name = btf__name_by_offset(btf, type->name_off);
+            if (type_name && strlen(type_name) > 0) {
+                /* Create a new cons cell */
+                cons = caml_alloc(2, 0);
+                Store_field(cons, 0, caml_copy_string(type_name));
+                Store_field(cons, 1, result);
+                result = cons;
+            }
+        }
+    }
+    
+    CAMLreturn(result);
+}
+
 /* Free BTF handle */
 value btf_free_stub(value btf_handle) {
     CAMLparam1(btf_handle);

src/ir_generator.ml

@@ -2976,8 +2976,15 @@ let lower_multi_program ast symbol_table source_name =
         failwith "Only one main() function is allowed";
 
       (* Extract struct definitions from AST (single source of truth) *)
+      (* Filter out kernel-only structs that should not be in userspace code *)
       let struct_definitions = List.filter_map (function
-        | Ast.StructDecl struct_def -> Some struct_def
+        | Ast.StructDecl struct_def -> 
+            (* Check if the struct has @kernel_only marker in comments/attributes *)
+            let is_kernel_only = List.exists (function
+              | Ast.SimpleAttribute "kernel_only" -> true
+              | _ -> false
+            ) struct_def.struct_attributes in
+            if is_kernel_only then None else Some struct_def
         | _ -> None
       ) ast in
 

src/kernel_types.ml

@@ -14,11 +14,15 @@
  * limitations under the License.
  *)
 
-(** Shared module for eBPF kernel type definitions *)
+(** Dynamic kernel type detection using BTF parsing *)
 
-(** Well-known eBPF context types that are provided by kernel headers
-    These should be marked as kernel_defined = true **)
-let well_known_ebpf_types = [
+open Printf
+
+(** Cache for BTF-extracted kernel types to avoid re-parsing *)
+let kernel_types_cache : (string, string list) Hashtbl.t = Hashtbl.create 16
+
+(** Fallback well-known eBPF types for when BTF is not available *)
+let fallback_well_known_ebpf_types = [
   (* eBPF context structs - provided by linux/bpf.h *)
   "xdp_md";
   "__sk_buff";
@@ -44,6 +48,41 @@ let well_known_ebpf_types = [
   "bpf_sockopt";
 ]
 
-(** Check if a type name is a well-known eBPF kernel type *)
-let is_well_known_ebpf_type type_name =
-  List.mem type_name well_known_ebpf_types 
+(** Extract kernel types from BTF file with caching *)
+let get_kernel_types_from_btf btf_path =
+  match Hashtbl.find_opt kernel_types_cache btf_path with
+  | Some cached_types -> cached_types
+  | None ->
+      let kernel_types = 
+        try
+          Btf_binary_parser.extract_all_kernel_struct_names btf_path
+        with
+        | _ -> 
+            printf "Warning: Failed to extract kernel types from BTF, using fallback list\n";
+            fallback_well_known_ebpf_types
+      in
+      Hashtbl.add kernel_types_cache btf_path kernel_types;
+      kernel_types
+
+(** Check if a type name is a well-known eBPF kernel type.
+    Uses BTF if available, otherwise falls back to hardcoded list. *)
+let is_well_known_ebpf_type ?btf_path type_name =
+  match btf_path with
+  | Some path when Sys.file_exists path ->
+      let kernel_types = get_kernel_types_from_btf path in
+      List.mem type_name kernel_types
+  | _ ->
+      (* Fallback to hardcoded list when BTF is not available *)
+      List.mem type_name fallback_well_known_ebpf_types
+
+(** Clear the kernel types cache (useful for testing or when BTF file changes) *)
+let clear_cache () =
+  Hashtbl.clear kernel_types_cache
+
+(** Get all known kernel types for the given BTF file (for debugging/inspection) *)
+let get_all_kernel_types ?btf_path () =
+  match btf_path with
+  | Some path when Sys.file_exists path ->
+      get_kernel_types_from_btf path
+  | _ ->
+      fallback_well_known_ebpf_types 

src/main.ml

@@ -713,7 +713,7 @@ let compile_source input_file output_dir _verbose generate_makefile btf_vmlinux_
           Ast.StructDecl { 
             struct_name = btf_type.Btf_parser.name; 
             struct_fields = fields; 
-            struct_attributes = []; 
+            struct_attributes = if btf_type.Btf_parser.kernel_defined then [Ast.SimpleAttribute "kernel_only"] else []; 
             struct_pos = { filename = "btf"; line = 1; column = 1 }
           }
       | "enum" ->
@@ -975,7 +975,7 @@ let compile_source input_file output_dir _verbose generate_makefile btf_vmlinux_
 
     (* Generate userspace coordinator directly to output directory with tail call analysis *)
     Userspace_codegen.generate_userspace_code_from_ir 
-      ~config_declarations ~type_aliases ~tail_call_analysis ~kfunc_dependencies ~resolved_imports ~symbol_table updated_optimized_ir ~output_dir:actual_output_dir input_file;
+      ~config_declarations ~type_aliases ~tail_call_analysis ~kfunc_dependencies ~resolved_imports ~symbol_table ?btf_path:btf_vmlinux_path updated_optimized_ir ~output_dir:actual_output_dir input_file;
 
     (* Output directory already created earlier *)
 

src/userspace_codegen.ml

@@ -2887,7 +2887,7 @@ let generate_load_function_with_tail_calls _base_name all_usage tail_call_analys
   else ""
 
 (** Generate complete userspace program from IR *)
-let generate_complete_userspace_program_from_ir ?(config_declarations = []) ?(type_aliases = []) ?(tail_call_analysis = {Tail_call_analyzer.dependencies = []; prog_array_size = 0; index_mapping = Hashtbl.create 16; errors = []}) ?(kfunc_dependencies = {kfunc_definitions = []; private_functions = []; program_dependencies = []; module_name = ""}) ?(resolved_imports = []) ?symbol_table (userspace_prog : ir_userspace_program) (global_maps : ir_map_def list) (ir_multi_prog : ir_multi_program) source_filename =
+let generate_complete_userspace_program_from_ir ?(config_declarations = []) ?(type_aliases = []) ?(tail_call_analysis = {Tail_call_analyzer.dependencies = []; prog_array_size = 0; index_mapping = Hashtbl.create 16; errors = []}) ?(kfunc_dependencies = {kfunc_definitions = []; private_functions = []; program_dependencies = []; module_name = ""}) ?(resolved_imports = []) ?symbol_table ?btf_path (userspace_prog : ir_userspace_program) (global_maps : ir_map_def list) (ir_multi_prog : ir_multi_program) source_filename =
   (* Collect function usage information from all functions first to determine if we need BPF headers *)
   let all_usage = List.fold_left (fun acc_usage func ->
     let func_usage = collect_function_usage_from_ir_function ~global_variables:ir_multi_prog.global_variables func in
@@ -3027,7 +3027,7 @@ let generate_complete_userspace_program_from_ir ?(config_declarations = []) ?(ty
 
   (* Filter out kernel-defined structs that are provided by kernel headers *)
   let user_defined_ir_structs = List.filter (fun ir_struct ->
-    not (Kernel_types.is_well_known_ebpf_type ir_struct.struct_name) &&
+    not (Kernel_types.is_well_known_ebpf_type ?btf_path ir_struct.struct_name) &&
     not (Struct_ops_registry.is_known_struct_ops ir_struct.struct_name)
   ) non_config_ir_structs in
 
@@ -3233,10 +3233,10 @@ void cleanup_bpf_maps(void) {
 |} includes string_typedefs type_alias_definitions string_helpers enum_definitions structs_with_pinned skeleton_code all_fd_declarations map_operation_functions ringbuf_handlers ringbuf_dispatch_functions auto_bpf_init_code getopt_parsing_code bpf_helper_functions struct_ops_attach_functions functions
 
 (** Generate userspace C code from IR multi-program *)
-let generate_userspace_code_from_ir ?(config_declarations = []) ?(type_aliases = []) ?(tail_call_analysis = {Tail_call_analyzer.dependencies = []; prog_array_size = 0; index_mapping = Hashtbl.create 16; errors = []}) ?(kfunc_dependencies = {kfunc_definitions = []; private_functions = []; program_dependencies = []; module_name = ""}) ?(resolved_imports = []) ?symbol_table (ir_multi_prog : ir_multi_program) ?(output_dir = ".") source_filename =
+let generate_userspace_code_from_ir ?(config_declarations = []) ?(type_aliases = []) ?(tail_call_analysis = {Tail_call_analyzer.dependencies = []; prog_array_size = 0; index_mapping = Hashtbl.create 16; errors = []}) ?(kfunc_dependencies = {kfunc_definitions = []; private_functions = []; program_dependencies = []; module_name = ""}) ?(resolved_imports = []) ?symbol_table ?btf_path (ir_multi_prog : ir_multi_program) ?(output_dir = ".") source_filename =
   let content = match ir_multi_prog.userspace_program with
     | Some userspace_prog -> 
-        generate_complete_userspace_program_from_ir ~config_declarations ~type_aliases ~tail_call_analysis ~kfunc_dependencies ~resolved_imports ?symbol_table userspace_prog ir_multi_prog.global_maps ir_multi_prog source_filename
+        generate_complete_userspace_program_from_ir ~config_declarations ~type_aliases ~tail_call_analysis ~kfunc_dependencies ~resolved_imports ?symbol_table ?btf_path userspace_prog ir_multi_prog.global_maps ir_multi_prog source_filename
     | None -> 
         sprintf {|#include <stdio.h>