Add support for enum64 handling in BTF parsing and struct_ops header generation.

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

examples/sched_ext_ops.kh

@@ -43,3 +43,20 @@ struct sched_ext_ops {
     hotplug_seq: u64,
     name: u32,
 }
+
+// Related kernel enums
+enum scx_public_consts {
+    SCX_OPS_NAME_LEN = 128,
+    SCX_SLICE_DFL = 20000000,
+    SCX_SLICE_INF = 18446744073709551615,
+}
+
+enum scx_dsq_id_flags {
+    SCX_DSQ_FLAG_BUILTIN = 9223372036854775808,
+    SCX_DSQ_FLAG_LOCAL_ON = 4611686018427387904,
+    SCX_DSQ_INVALID = 9223372036854775808,
+    SCX_DSQ_GLOBAL = 9223372036854775809,
+    SCX_DSQ_LOCAL = 9223372036854775810,
+    SCX_DSQ_LOCAL_ON = 13835058055282163712,
+    SCX_DSQ_LOCAL_CPU_MASK = 4294967295,
+}

examples/sched_ext_simple.ks

@@ -8,23 +8,6 @@ extern scx_bpf_select_cpu_dfl(p: *u8, prev_cpu: i32, wake_flags: u64, direct: *b
 extern scx_bpf_dsq_insert(p: *u8, dsq_id: u64, slice: u64, enq_flags: u64) -> void
 extern scx_bpf_consume(dsq_id: u64, cpu: i32, flags: u64) -> i32
 
-// Kernel enums (extracted from BTF)
-enum scx_public_consts {
-        SCX_OPS_NAME_LEN = 128,
-        SCX_SLICE_DFL = 20000000,
-        SCX_SLICE_INF = 18446744073709551615,
-}
-
-enum scx_dsq_id_flags {
-        SCX_DSQ_FLAG_BUILTIN = 9223372036854775808,
-        SCX_DSQ_FLAG_LOCAL_ON = 4611686018427387904,
-        SCX_DSQ_INVALID = 9223372036854775808,
-        SCX_DSQ_GLOBAL = 9223372036854775809,
-        SCX_DSQ_LOCAL = 9223372036854775810,
-        SCX_DSQ_LOCAL_ON = 13835058055282163712,
-        SCX_DSQ_LOCAL_CPU_MASK = 4294967295,
-}
-
 // Simple FIFO scheduler implementation
 @struct_ops("sched_ext_ops")
 impl simple_fifo_scheduler {

src/btf_binary_parser.ml

@@ -35,7 +35,7 @@ external btf_new_from_file : string -> btf_handle option = "btf_new_from_file_st
 external btf_get_nr_types : btf_handle -> int = "btf_get_nr_types_stub"
 external btf_type_by_id : btf_handle -> int -> (int * string * int * int * int) = "btf_type_by_id_stub"
 
-external btf_type_get_members : btf_handle -> int -> (string * int) array = "btf_type_get_members_stub"
+external btf_type_get_members : btf_handle -> int -> (string * string) 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_and_enum_names : btf_handle -> string list = "btf_extract_kernel_struct_and_enum_names_stub"
@@ -59,12 +59,13 @@ let parse_btf_file btf_path target_types =
         try
           let member_array = btf_type_get_members btf_handle union_type_id in
           let member_list = Array.to_list member_array in
-          List.fold_left (fun acc (field_name, field_type_id) ->
+          List.fold_left (fun acc (field_name, field_type_id_str) ->
             if field_name = "" then
               (* Skip anonymous members within the union to avoid infinite recursion *)
               acc
             else
               try
+                let field_type_id = int_of_string field_type_id_str in
                 let field_type = btf_resolve_type btf_handle field_type_id in
                 (field_name, field_type) :: acc
               with
@@ -85,6 +86,7 @@ let parse_btf_file btf_path target_types =
               | 4 -> "struct"
               | 5 -> "union"
               | 6 -> "enum"
+              | 19 -> "enum64"
               | _ -> "unknown"
             in
 
@@ -96,8 +98,9 @@ let parse_btf_file btf_path target_types =
                   let member_array = btf_type_get_members btf_handle i in
                   let member_list = Array.to_list member_array in
                   (* Resolve each member's type and handle anonymous unions *)
-                  let resolved_members = List.fold_left (fun acc (field_name, field_type_id) ->
+                  let resolved_members = List.fold_left (fun acc (field_name, field_type_id_str) ->
                     try
+                      let field_type_id = int_of_string field_type_id_str in
                       let field_type = btf_resolve_type btf_handle field_type_id in
                       if field_name = "" && field_type = "union" then
                         (* Anonymous union: extract its members and flatten them *)
@@ -120,14 +123,14 @@ let parse_btf_file btf_path target_types =
                   Some (List.rev resolved_members)
                 with
                 | _ -> None
-              ) else if kind_int = 6 then (
-                (* Enum: extract enum values *)
+              ) else if kind_int = 6 || kind_int = 19 then (
+                (* Enum (kind 6) or Enum64 (kind 19): extract enum values *)
                 try
                   let member_array = btf_type_get_members btf_handle i in
                   let member_list = Array.to_list member_array in
-                  (* For enums, second element is the value, not type_id *)
+                  (* For enums, second element is the value (now as string) *)
                   let enum_values = List.map (fun (enum_name, enum_value) ->
-                    (enum_name, string_of_int enum_value)
+                    (enum_name, enum_value)
                   ) member_list in
                   Some enum_values
                 with

src/btf_parser.ml

@@ -597,7 +597,7 @@ let convert_btf_type_to_ks_definition btf_type =
         | None -> ""
       in
       sprintf "struct %s {\n%s\n}" btf_type.Btf_binary_parser.name fields
-  | "enum" ->
+  | "enum" | "enum64" ->
       let values = match btf_type.Btf_binary_parser.members with
         | Some members ->
             List.map (fun (name, value) -> 
@@ -676,6 +676,15 @@ let generate_program_header ~extract_kfuncs prog_type btf_path =
   let enum_section = if hardcoded_enums <> "" then hardcoded_enums ^ "\n\n" else "" in
   sprintf "%s%s\n\n%s%s\n" header type_definitions enum_section kfunc_declarations
 
+(* Get struct_ops-specific enum names to extract from BTF *)
+let get_struct_ops_enum_names struct_ops_name =
+  match struct_ops_name with
+  | "sched_ext_ops" -> [
+      "scx_public_consts";
+      "scx_dsq_id_flags";
+    ]
+  | _ -> []
+
 (* Generate struct_ops-specific header content using BTF *)
 let generate_struct_ops_header struct_ops_name btf_path =
   let header = sprintf {|// AUTO-GENERATED %s DEFINITIONS - DO NOT EDIT
@@ -700,7 +709,30 @@ let generate_struct_ops_header struct_ops_name btf_path =
         sprintf "// Warning: BTF extraction failed for %s" struct_ops_name
   in
 
-  sprintf "%s%s\n" header struct_definitions
+  (* Extract related enums for specific struct_ops types *)
+  let enum_definitions = 
+    let enum_names = get_struct_ops_enum_names struct_ops_name in
+    if enum_names <> [] then
+      try
+        let btf_enums = Btf_binary_parser.parse_btf_file btf_path enum_names in
+        let filtered_enums = List.filter (fun btf_type ->
+          (btf_type.Btf_binary_parser.kind = "enum" || btf_type.Btf_binary_parser.kind = "enum64") && 
+          List.mem btf_type.Btf_binary_parser.name enum_names
+        ) btf_enums in
+        if filtered_enums <> [] then
+          "\n// Related kernel enums\n" ^ 
+          (List.map convert_btf_type_to_ks_definition filtered_enums |> String.concat "\n\n")
+        else
+          sprintf "\n// Warning: No BTF enums found for %s" struct_ops_name
+      with
+      | exn ->
+          printf "Warning: Failed to extract BTF enums for %s: %s\n" struct_ops_name (Printexc.to_string exn);
+          sprintf "\n// Warning: BTF enum extraction failed for %s" struct_ops_name
+    else
+      ""
+  in
+  
+  sprintf "%s%s%s\n" header struct_definitions enum_definitions
 
 (* Generate tracepoint-specific header content using BTF *)
 let generate_tracepoint_header category_event btf_path =

src/btf_stubs.c

@@ -17,6 +17,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <inttypes.h>
 #include <bpf/libbpf.h>
 #include <bpf/btf.h>
 #include <caml/mlvalues.h>
@@ -149,6 +150,7 @@ value btf_type_by_id_stub(value btf_handle, value type_id) {
         case BTF_KIND_STRUCT:
         case BTF_KIND_UNION:
         case BTF_KIND_ENUM:
+        case BTF_KIND_ENUM64:
             size = t->size;
             break;
         case BTF_KIND_PTR:
@@ -206,8 +208,8 @@ value btf_type_get_members_stub(value btf_handle, value type_id) {
     }
 
     int kind = btf_kind(t);
-    if (kind != BTF_KIND_STRUCT && kind != BTF_KIND_UNION && kind != BTF_KIND_ENUM) {
-        /* Return empty array for non-struct/union/enum types */
+    if (kind != BTF_KIND_STRUCT && kind != BTF_KIND_UNION && kind != BTF_KIND_ENUM && kind != BTF_KIND_ENUM64) {
+        /* Return empty array for non-struct/union/enum/enum64 types */
         CAMLreturn(caml_alloc_tuple(0));
     }
 
@@ -227,8 +229,30 @@ value btf_type_get_members_stub(value btf_handle, value type_id) {
 
             member_tuple = caml_alloc_tuple(2);
             Store_field(member_tuple, 0, caml_copy_string(enum_name));
-            /* For enums, store the value instead of type_id */
-            Store_field(member_tuple, 1, Val_int(enums[i].val));
+            /* For enums, convert value to string for consistency with enum64 */
+            char value_str[16];
+            snprintf(value_str, sizeof(value_str), "%d", enums[i].val);
+            Store_field(member_tuple, 1, caml_copy_string(value_str));
+            
+            Store_field(result, i, member_tuple);
+        }
+    } else if (kind == BTF_KIND_ENUM64) {
+        /* Handle enum64 types - extract enum values */
+        const struct btf_enum64 *enums = btf_enum64(t);
+        for (int i = 0; i < vlen; i++) {
+            const char *enum_name = btf__name_by_offset(btf, enums[i].name_off);
+            if (!enum_name) enum_name = "";
+            
+            member_tuple = caml_alloc_tuple(2);
+            Store_field(member_tuple, 0, caml_copy_string(enum_name));
+            
+            /* For enum64, combine hi32 and lo32 to get the full 64-bit value */
+            uint64_t full_value = ((uint64_t)enums[i].val_hi32 << 32) | enums[i].val_lo32;
+            
+            /* Convert to string to preserve full precision */
+            char value_str[32];
+            snprintf(value_str, sizeof(value_str), "%" PRIu64, full_value);
+            Store_field(member_tuple, 1, caml_copy_string(value_str));
 
             Store_field(result, i, member_tuple);
         }
@@ -241,7 +265,10 @@ value btf_type_get_members_stub(value btf_handle, value type_id) {
 
             member_tuple = caml_alloc_tuple(2);
             Store_field(member_tuple, 0, caml_copy_string(member_name));
-            Store_field(member_tuple, 1, Val_int(members[i].type));
+            /* For struct/union, convert type_id to string for consistency */
+            char type_id_str[16];
+            snprintf(type_id_str, sizeof(type_id_str), "%u", members[i].type);
+            Store_field(member_tuple, 1, caml_copy_string(type_id_str));
 
             Store_field(result, i, member_tuple);
         }
@@ -305,13 +332,15 @@ static char* resolve_type_to_string(struct btf *btf, int type_id) {
         }
         case BTF_KIND_STRUCT:
         case BTF_KIND_UNION:
-        case BTF_KIND_ENUM: {
+        case BTF_KIND_ENUM:
+        case BTF_KIND_ENUM64: {
             const char *name = btf__name_by_offset(btf, t->name_off);
             if (name && strlen(name) > 0) {
                 return strdup(name);
             }
             return strdup(kind == BTF_KIND_STRUCT ? "struct" : 
-                         kind == BTF_KIND_UNION ? "union" : "enum");
+                         kind == BTF_KIND_UNION ? "union" : 
+                         kind == BTF_KIND_ENUM ? "enum" : "enum64");
         }
         default:
             return strdup("unknown");
@@ -477,22 +506,18 @@ value btf_resolve_type_stub(value btf_handle, value type_id) {
         }
         case BTF_KIND_STRUCT:
         case BTF_KIND_UNION:
-        case BTF_KIND_ENUM: {
-            const char *name = btf__name_by_offset(btf, t->name_off);
-            if (name && strlen(name) > 0) {
-                CAMLreturn(caml_copy_string(name));
-            }
-            /* For anonymous structs/unions */
-            CAMLreturn(caml_copy_string(kind == BTF_KIND_STRUCT ? "struct" : 
-                                      kind == BTF_KIND_UNION ? "union" : "enum"));
-        }
+        case BTF_KIND_ENUM:
         case BTF_KIND_ENUM64: {
             const char *name = btf__name_by_offset(btf, t->name_off);
             if (name && strlen(name) > 0) {
                 CAMLreturn(caml_copy_string(name));
             }
-            CAMLreturn(caml_copy_string("enum64"));
+            /* For anonymous structs/unions/enums */
+            CAMLreturn(caml_copy_string(kind == BTF_KIND_STRUCT ? "struct" : 
+                                      kind == BTF_KIND_UNION ? "union" : 
+                                      kind == BTF_KIND_ENUM ? "enum" : "enum64"));
         }
+
         case BTF_KIND_FWD: {
             const char *name = btf__name_by_offset(btf, t->name_off);
             if (name && strlen(name) > 0) {