Add support for program target in AST and related code generation

This update introduces a new optional field, `prog_target`, to the program definition structure, allowing for better handling of kprobe and tracepoint programs. The field is initialized to `None` in various program creation functions and is utilized in the IR generation and code generation processes to ensure proper target information is passed through. Additionally, tests have been updated to reflect this new field.

Author: congwang-mk

src/ast.ml

@@ -276,6 +276,7 @@ type program_def = {
   prog_functions: function_def list;
   prog_maps: map_declaration list; (* Maps local to this program *)
   prog_structs: struct_def list; (* Structs local to this program *)
+  prog_target: string option; (* Target for kprobe/tracepoint programs *)
   prog_pos: position;
 }
 
@@ -402,6 +403,7 @@ let make_program name prog_type functions pos = {
   prog_functions = functions;
   prog_maps = [];
   prog_structs = [];
+  prog_target = None;
   prog_pos = pos;
 }
 
@@ -411,6 +413,7 @@ let make_program_with_maps name prog_type functions maps pos = {
   prog_functions = functions;
   prog_maps = maps;
   prog_structs = [];
+  prog_target = None;
   prog_pos = pos;
 }
 
@@ -420,6 +423,7 @@ let make_program_with_all name prog_type functions maps structs pos = {
   prog_functions = functions;
   prog_maps = maps;
   prog_structs = structs;
+  prog_target = None;
   prog_pos = pos;
 }
 

src/ebpf_c_codegen.ml

@@ -3352,24 +3352,16 @@ let generate_c_function ctx ir_func =
          | Some target -> 
              (* If we have the target, convert KernelScript format to raw tracepoint format *)
              if String.contains target '/' then
-               let event_name = String.map (function '/' -> '_' | c -> c) target in
-               sprintf "SEC(\"raw_tracepoint/%s\")" event_name
+               (* For raw tracepoints, extract just the event name (part after slash) *)
+               let parts = String.split_on_char '/' target in
+               (match parts with
+                | [_category; event] -> sprintf "SEC(\"raw_tracepoint/%s\")" event
+                | _ -> sprintf "SEC(\"raw_tracepoint/%s\")" target)
              else
                sprintf "SEC(\"raw_tracepoint/%s\")" target
          | None ->
-             (* Fallback: try to extract from function name *)
-             let func_name = ir_func.func_name in
-             if String.contains func_name '_' then
-               (* Function name like "sched_sched_switch_handler" -> extract "sched_switch" *)
-               let parts = String.split_on_char '_' func_name in
-               (match parts with
-                | category :: event :: "handler" :: _ -> 
-                    sprintf "SEC(\"raw_tracepoint/%s_%s\")" category event
-                | category :: event :: _ when List.length parts >= 2 ->
-                    sprintf "SEC(\"raw_tracepoint/%s_%s\")" category event
-                | _ -> "SEC(\"raw_tracepoint\")")
-             else
-               "SEC(\"raw_tracepoint\")")
+             (* This should not happen now that we properly pass targets through *)
+             failwith "Tracepoint function is missing target information")
     | _ ->
         (* For non-struct_ops, non-kprobe, and non-tracepoint functions, only generate SEC if it's a main function *)
         if ir_func.is_main then

src/ir_generator.ml

@@ -2731,12 +2731,8 @@ let lower_single_program ctx prog_def _global_ir_maps _kernel_shared_functions =
     (* For attributed functions (single function programs), the function IS the entry function *)
     (* But struct_ops functions should NOT be marked as main functions *)
     let is_attributed_entry = (List.length prog_def.prog_functions = 1 && index = 0 && prog_def.prog_type <> Ast.StructOps) in
-    (* Extract target from attributed function if available *)
-    let func_target = 
-      (* This is a hack to find the original attributed function and extract the target *)
-      (* For now, we'll pass None and fix this in the eBPF codegen *)
-      None 
-    in
+    (* Extract target from program definition *)
+    let func_target = prog_def.prog_target in
     let temp_func = lower_function ctx prog_def.prog_name ~program_type:(Some prog_def.prog_type) ~func_target func in
     if is_attributed_entry then
       (* Mark the attributed function as entry by updating the is_main field *)
@@ -2830,9 +2826,10 @@ let lower_multi_program ast symbol_table source_name =
                     prog_functions = [attr_func.attr_function];
                     prog_maps = [];
                     prog_structs = [];
+                    prog_target = None;
                     prog_pos = attr_func.attr_pos;
                   })
-         | AttributeWithArg (attr_name, _target_func) :: _ ->
+         | AttributeWithArg (attr_name, target_func) :: _ ->
              (* Handle attributes with arguments like @kprobe("sys_read") *)
              (match attr_name with
               | "kprobe" -> 
@@ -2842,6 +2839,7 @@ let lower_multi_program ast symbol_table source_name =
                     prog_functions = [attr_func.attr_function];
                     prog_maps = [];
                     prog_structs = [];
+                    prog_target = Some target_func;
                     prog_pos = attr_func.attr_pos;
                   }
               | "tracepoint" ->
@@ -2851,6 +2849,7 @@ let lower_multi_program ast symbol_table source_name =
                     prog_functions = [attr_func.attr_function];
                     prog_maps = [];
                     prog_structs = [];
+                    prog_target = Some target_func;
                     prog_pos = attr_func.attr_pos;
                   }
               | _ -> None)
@@ -2871,6 +2870,7 @@ let lower_multi_program ast symbol_table source_name =
             prog_functions = [func];
             prog_maps = [];
             prog_structs = [];
+            prog_target = None; (* struct_ops don't have targets *)
             prog_pos = func.func_pos;
           }
       | Ast.ImplStaticField (_, _) -> None  (* Static fields are not programs *)

src/main.ml

@@ -814,7 +814,8 @@ let compile_source input_file output_dir _verbose generate_makefile btf_vmlinux_
 
     (* Create a program structure for safety analysis *)
     let safety_program = {
-      Ast.prog_name = base_name;
+              Ast.prog_name = base_name;
+        prog_target = None;
       prog_type = Xdp; (* Default - not used by safety checker *)
       prog_functions = all_functions;
       prog_maps = all_maps;

src/multi_program_analyzer.ml

@@ -122,6 +122,7 @@ let extract_programs (ast: declaration list) : program_def list =
                     prog_functions = [attr_func.attr_function];
                     prog_maps = [];
                     prog_structs = [];
+                    prog_target = None;
                     prog_pos = attr_func.attr_pos;
                   })
          | _ -> None)

src/type_checker.ml

@@ -2567,6 +2567,7 @@ let typed_program_to_program tprog original_prog =
     prog_functions = List.map typed_function_to_function tprog.tprog_functions;
     prog_maps = original_prog.prog_maps;  (* Preserve original map declarations *)
     prog_structs = original_prog.prog_structs;  (* Preserve original struct declarations *)
+    prog_target = original_prog.prog_target;  (* Preserve original target *)
     prog_pos = tprog.tprog_pos }
 
 (** Convert typed AST back to annotated AST declarations *)

tests/test_function_generation.ml

@@ -51,6 +51,7 @@ let test_function_parameters _ =
   (* Create program containing this function *)
   let prog_def = {
     prog_name = "test_prog";
+    prog_target = None;
     prog_type = Xdp;
     prog_maps = [];
     prog_structs = [];
@@ -127,6 +128,7 @@ let test_program_function_calls _ =
   (* Create program with both functions *)
   let prog_def = {
     prog_name = "test_prog";
+    prog_target = None;
     prog_type = Xdp;
     prog_maps = [];
     prog_structs = [];
@@ -195,6 +197,7 @@ let test_multiple_parameters _ =
 
   let prog_def = {
     prog_name = "test_prog";
+    prog_target = None;
     prog_type = Xdp;
     prog_maps = [];
     prog_structs = [];

tests/test_symbol_table.ml

@@ -88,7 +88,8 @@ let create_test_function name params return_type =
 
 let create_test_program name functions =
   {
-    prog_name = name;
+          prog_name = name;
+      prog_target = None;
     prog_type = Xdp;
     prog_functions = functions;
     prog_maps = [];

tests/test_tracepoint.ml

@@ -201,9 +201,75 @@ fn sched_switch_handler(ctx: *trace_event_raw_sched_switch) -> i32 {
   check bool "Function should be marked as main" true main_func.is_main;
   check string "Function name should match" "sched_switch_handler" main_func.func_name
 
+(* NEW: Target Propagation Tests *)
+let test_tracepoint_target_propagation _ =
+  let source = "@tracepoint(\"sched/sched_switch\")
+fn sched_switch_handler(ctx: *trace_event_raw_sched_switch) -> i32 {
+    return 0
+}" in
+  let ast = parse_string source in
+  let typed_ast = type_check_ast ast in
+  let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in
+  let ir_multi_prog = generate_ir typed_ast symbol_table "test_tracepoint" in
+  let program = List.hd ir_multi_prog.programs in
+  let main_func = program.entry_function in
+  
+  (* Test that the target is properly propagated through IR generation *)
+  check (option string) "Function should have correct target" (Some "sched/sched_switch") main_func.func_target
+
+let test_multiple_tracepoint_targets _ =
+  (* Test various tracepoint targets to ensure they all work correctly *)
+  let test_cases = [
+    ("sched/sched_switch", "SEC(\"raw_tracepoint/sched_switch\")");
+    ("net/netif_rx", "SEC(\"raw_tracepoint/netif_rx\")");
+    ("syscalls/sys_enter_read", "SEC(\"raw_tracepoint/sys_enter_read\")");
+    ("syscalls/sys_exit_write", "SEC(\"raw_tracepoint/sys_exit_write\")");
+    ("irq/irq_handler_entry", "SEC(\"raw_tracepoint/irq_handler_entry\")");
+  ] in
+  
+  List.iter (fun (target, expected_sec) ->
+    let source = Printf.sprintf "@tracepoint(\"%s\")
+fn handler(ctx: *trace_event_raw_context) -> i32 {
+    return 0
+}" target in
+    let ast = parse_string source in
+    let typed_ast = type_check_ast ast in
+    let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in
+    let ir_multi_prog = generate_ir typed_ast symbol_table "test" in
+    let c_code = generate_c_multi_program ir_multi_prog in
+    
+    check bool (Printf.sprintf "Should generate %s for target %s" expected_sec target) true
+      (Str.search_forward (Str.regexp_string expected_sec) c_code 0 >= 0)
+  ) test_cases
+
+let test_sched_switch_bug_regression _ =
+  (* Regression test: Ensure we don't generate the buggy SEC("raw_tracepoint/sched_sched") *)
+  let source = "@tracepoint(\"sched/sched_switch\")
+fn sched_switch_handler(ctx: *trace_event_raw_sched_switch) -> i32 {
+    return 0
+}" in
+  let ast = parse_string source in
+  let typed_ast = type_check_ast ast in
+  let symbol_table = Kernelscript.Symbol_table.build_symbol_table typed_ast in
+  let ir_multi_prog = generate_ir typed_ast symbol_table "test_regression" in
+  let c_code = generate_c_multi_program ir_multi_prog in
+  
+  (* Ensure correct SEC() is generated *)
+  check bool "Should generate correct SEC(raw_tracepoint/sched_switch)" true
+    (Str.search_forward (Str.regexp_string "SEC(\"raw_tracepoint/sched_switch\")") c_code 0 >= 0);
+  
+  (* Ensure buggy SEC() is NOT generated *)
+  check bool "Should NOT generate buggy SEC(raw_tracepoint/sched_sched)" true
+    (try 
+       let _ = Str.search_forward (Str.regexp_string "SEC(\"raw_tracepoint/sched_sched\")") c_code 0 in
+       false  (* Found the buggy pattern - test should fail *)
+     with Not_found -> 
+       true   (* Didn't find the buggy pattern - test should pass *)
+    )
+
 (* 4. Code Generation Tests *)
 let test_raw_tracepoint_section_name_generation _ =
-  (* Test minimal raw tracepoint section name conversion logic *)
+  (* Test correct raw tracepoint section name generation *)
   let source = "@tracepoint(\"sched/sched_switch\")
 fn sched_switch_handler(ctx: *trace_event_raw_sched_switch) -> i32 {
     return 0
@@ -214,9 +280,9 @@ fn sched_switch_handler(ctx: *trace_event_raw_sched_switch) -> i32 {
   let ir_multi_prog = generate_ir typed_ast symbol_table "test_raw_tracepoint" in
   let c_code = generate_c_multi_program ir_multi_prog in
 
-  (* Check that forward slash is converted to underscore in section name *)
-  check bool "Should contain raw_tracepoint section with underscore" true
-    (String.contains c_code (String.get "SEC(\"raw_tracepoint/sched_sched_switch\")" 0))
+  (* Check that the correct SEC() is generated with just the event name *)
+  check bool "Should contain correct raw_tracepoint/sched_switch section" true
+    (Str.search_forward (Str.regexp_string "SEC(\"raw_tracepoint/sched_switch\")") c_code 0 >= 0)
 
 let test_tracepoint_ebpf_codegen _ =
   let source = "@tracepoint(\"sched/sched_switch\")
@@ -230,8 +296,8 @@ fn sched_switch_handler(ctx: *trace_event_raw_sched_switch) -> i32 {
   let c_code = generate_c_multi_program ir_multi_prog in
 
   (* Check for tracepoint-specific C code elements *)
-  check bool "Should contain SEC(\"tracepoint\")" true 
-    (String.contains c_code (String.get "SEC(\"tracepoint\")" 0));
+  check bool "Should contain correct raw_tracepoint SEC" true 
+    (Str.search_forward (Str.regexp_string "SEC(\"raw_tracepoint/sched_switch\")") c_code 0 >= 0);
   check bool "Should contain function definition" true
     (String.contains c_code (String.get "sched_switch_handler" 0));
   check bool "Should contain struct parameter" true
@@ -342,8 +408,8 @@ fn sys_enter_open_handler(ctx: *trace_event_raw_sys_enter) -> i32 {
   let c_code = generate_c_multi_program ir_multi_prog in
 
   (* Comprehensive end-to-end validation *)
-  check bool "Contains tracepoint section" true
-    (String.contains c_code (String.get "SEC(\"tracepoint\")" 0));
+  check bool "Contains correct raw_tracepoint section" true
+    (Str.search_forward (Str.regexp_string "SEC(\"raw_tracepoint/sys_enter_open\")") c_code 0 >= 0);
   check bool "Contains function name" true
     (String.contains c_code (String.get "sys_enter_open_handler" 0));
   check bool "Contains context struct" true
@@ -381,6 +447,9 @@ let type_checking_tests = [
 let ir_generation_tests = [
   "tracepoint IR generation", `Quick, test_tracepoint_ir_generation;
   "tracepoint function signature validation", `Quick, test_tracepoint_function_signature_validation;
+  "tracepoint target propagation", `Quick, test_tracepoint_target_propagation;
+  "multiple tracepoint targets", `Quick, test_multiple_tracepoint_targets;
+  "sched_switch bug regression", `Quick, test_sched_switch_bug_regression;
 ]
 
 let code_generation_tests = [

tests/test_utils.ml

@@ -185,6 +185,7 @@ module Helpers = struct
   let make_test_program name prog_type main_func =
     {
       prog_name = name;
+      prog_target = None;
       prog_type = prog_type;
       prog_functions = [main_func];
       prog_maps = [];