Merge pull request #14 from SiyuanSun0736/fix_in_ringbuf_demo

Auto-deref map access in assignments and Add Forward Declarations

Author: congwang-mk

src/ebpf_c_codegen.ml

@@ -2559,7 +2559,10 @@ and generate_assignment ctx dest_val expr is_const =
              | IRValue src_val ->
                  (* Simple value assignment *)
                  let dest_str = generate_c_value ctx dest_val in
-                 let src_str = generate_c_value ctx src_val in
+                 (* Auto-dereference map access to get the value, not the pointer *)
+                 let src_str = (match src_val.value_desc with
+                   | IRMapAccess (_, _, _) -> generate_c_value ~auto_deref_map_access:true ctx src_val
+                   | _ -> generate_c_value ctx src_val) in
                  emit_line ctx (sprintf "%s%s = %s;" assignment_prefix dest_str src_str)
              | _ ->
                  (* Other expressions *)

src/userspace_codegen.ml

@@ -2817,6 +2817,19 @@ let generate_pinned_globals_support _project_name global_variables =
 
 (** Generate ring buffer event handler functions *)
 let generate_ringbuf_handlers_from_registry (registry : Ir.ir_ring_buffer_registry) ~dispatch_used =
+  (* Generate forward declarations for callback functions *)
+  let forward_declarations = List.map (fun rb_decl ->
+    let ringbuf_name = rb_decl.rb_name in
+    let value_type = c_type_from_ir_type rb_decl.rb_value_type in
+    let handler_name = match List.assoc_opt ringbuf_name registry.event_handler_registrations with
+      | Some handler -> handler
+      | None -> 
+          (* Try callback function naming convention: {ringbuf_name}_callback *)
+          ringbuf_name ^ "_callback"
+    in
+    sprintf "int %s(%s *event);" handler_name value_type
+  ) registry.ring_buffer_declarations |> String.concat "\n" in
+  
   let event_handlers = List.map (fun rb_decl ->
     let ringbuf_name = rb_decl.rb_name in
     let value_type = c_type_from_ir_type rb_decl.rb_value_type in
@@ -2842,7 +2855,11 @@ static int %s_event_handler(void *ctx, void *data, size_t data_sz) {
   else "" in
 
   (* Only generate event handlers if dispatch is actually used *)
-  let final_event_handlers = if dispatch_used then event_handlers else "" in
+  let final_event_handlers = if dispatch_used then 
+    if List.length registry.ring_buffer_declarations > 0 then
+      sprintf "\n// Forward declarations for ring buffer callbacks\n%s\n%s" forward_declarations event_handlers
+    else ""
+  else "" in
 
   final_event_handlers ^ combined_rb_declaration
 

tests/test_ebpf_c_codegen.ml

@@ -1146,6 +1146,37 @@ let test_global_map_redefinition_fix () =
 
   ()
 
+(** Test map access auto-dereference in variable assignment (covers lines 2559-2566) *)
+let test_map_access_auto_deref_in_assignment () =
+  let ctx = create_c_context () in
+
+  (* Build: count = my_map[user_key]
+     IRMapAccess carries the raw lookup-pointer as its underlying value. *)
+  let key_val = make_ir_value (IRVariable "user_key") IRU32 test_pos in
+  let underlying_desc = IRVariable "map_ptr_0" in
+  let underlying_type = IRPointer (IRU64, make_bounds_info ()) in
+  let src_val = make_ir_value
+    (IRMapAccess ("my_map", key_val, (underlying_desc, underlying_type)))
+    IRU64 test_pos in
+
+  let dest_val = make_ir_value (IRVariable "count") IRU64 test_pos in
+  let assign_instr = make_ir_instruction
+    (IRAssign (dest_val, make_ir_expr (IRValue src_val) IRU64 test_pos))
+    test_pos in
+
+  generate_c_instruction ctx assign_instr;
+
+  let output = String.concat "\n" ctx.output_lines in
+
+  (* auto_deref_map_access:true emits a guarded dereference:
+     count = ({ __u64 __val = {0}; if (map_ptr_0) { __val = *(map_ptr_0); } __val; }); *)
+  check bool "dest variable present in output" true (contains_substr output "count =");
+  check bool "__val used for safe dereference" true (contains_substr output "__val");
+  check bool "null-guard if-check emitted" true (contains_substr output "if (map_ptr_0)");
+  check bool "pointer dereference emitted" true (contains_substr output "*(map_ptr_0)");
+  (* Without the fix this branch would fall through to the raw-pointer path *)
+  check bool "no raw pointer assignment" false (contains_substr output "count = map_ptr_0")
+
 (** Test suite definition *)
 let suite =
   [
@@ -1191,6 +1222,8 @@ let suite =
     ("eBPF function generation bug fix", `Quick, test_ebpf_function_generation_bug_fix);
     (* Test to prevent global variable map redefinition regression *)
     ("Global map redefinition fix", `Quick, test_global_map_redefinition_fix);
+    (* Coverage for IRMapAccess auto-dereference path in generate_assignment *)
+    ("Map access auto-deref in assignment", `Quick, test_map_access_auto_deref_in_assignment);
   ]
 
 (** Run all tests *)

tests/test_ringbuf.ml

@@ -1095,6 +1095,107 @@ fn main() -> i32 {
   check bool "Should generate user handler functions" true 
     (contains_substr c_code "handle_network" && contains_substr c_code "handle_security")
 
+(** ── Unit tests for generate_ringbuf_handlers_from_registry (lines 2812-2826) ── *)
+
+(** Helper: build a minimal ir_ring_buffer_registry for unit testing *)
+let make_registry ?(event_handler_registrations = []) rb_decls =
+  {
+    Ir.ring_buffer_declarations = rb_decls;
+    Ir.event_handler_registrations = event_handler_registrations;
+    Ir.usage_summary = {
+      Ir.used_in_ebpf = [];
+      Ir.used_in_userspace = [];
+      Ir.needs_event_processing = [];
+    };
+  }
+
+let make_rb_decl name value_type =
+  {
+    Ir.rb_name = name;
+    Ir.rb_value_type = value_type;
+    Ir.rb_size = 4096;
+    Ir.rb_is_global = true;
+    Ir.rb_declaration_pos = { Ast.line = 1; column = 1; filename = "test.ks" };
+  }
+
+(** Test: empty registry produces empty output regardless of dispatch_used *)
+let test_handlers_empty_registry () =
+  let registry = make_registry [] in
+  let out_true  = Userspace_codegen.generate_ringbuf_handlers_from_registry
+                    registry ~dispatch_used:true in
+  let out_false = Userspace_codegen.generate_ringbuf_handlers_from_registry
+                    registry ~dispatch_used:false in
+  check string "empty registry + dispatch_used:true  → empty" "" out_true;
+  check string "empty registry + dispatch_used:false → empty" "" out_false
+
+(** Test: None branch – no event_handler_registrations entry → fallback to {rb_name}_callback *)
+let test_handlers_none_branch_fallback () =
+  (* No registration entry for "events" → handler name must fall back to "events_callback" *)
+  let decl     = make_rb_decl "events" Ir.IRU32 in
+  let registry = make_registry [decl] in            (* event_handler_registrations = [] *)
+  let out = Userspace_codegen.generate_ringbuf_handlers_from_registry
+              registry ~dispatch_used:true in
+  check bool "None branch: _callback fallback in event handler wrapper"
+    true (contains_substr out "events_callback");
+  check bool "None branch: event handler wrapper function generated"
+    true (contains_substr out "events_event_handler");
+  check bool "None branch: wrapper calls fallback handler"
+    true (contains_substr out "return events_callback(event)")
+
+(** Test: Some branch – registered handler name is used instead of fallback *)
+let test_handlers_some_branch_registered_name () =
+  let decl     = make_rb_decl "events" (Ir.IRStruct ("Event", [("id", Ir.IRU32)])) in
+  let registry = make_registry ~event_handler_registrations:[("events", "handle_event")] [decl] in
+  let out = Userspace_codegen.generate_ringbuf_handlers_from_registry
+              registry ~dispatch_used:true in
+  check bool "Some branch: registered handler name used in wrapper"
+    true (contains_substr out "handle_event");
+  check bool "Some branch: wrapper calls registered handler"
+    true (contains_substr out "return handle_event(event)");
+  check bool "Some branch: fallback name NOT used"
+    false (contains_substr out "events_callback")
+
+(** Test: dispatch_used:false → no event handler wrappers emitted *)
+let test_handlers_dispatch_false_no_wrappers () =
+  let decl     = make_rb_decl "events" Ir.IRU32 in
+  let registry = make_registry ~event_handler_registrations:[("events", "handle_event")] [decl] in
+  let out = Userspace_codegen.generate_ringbuf_handlers_from_registry
+              registry ~dispatch_used:false in
+  check bool "dispatch_used:false → no event_handler wrapper"
+    false (contains_substr out "events_event_handler");
+  check bool "dispatch_used:false → no combined_rb declaration"
+    false (contains_substr out "combined_rb")
+
+(** Test: dispatch_used:true → combined_rb declaration emitted *)
+let test_handlers_dispatch_true_combined_rb () =
+  let decl     = make_rb_decl "events" Ir.IRU32 in
+  let registry = make_registry ~event_handler_registrations:[("events", "handle_event")] [decl] in
+  let out = Userspace_codegen.generate_ringbuf_handlers_from_registry
+              registry ~dispatch_used:true in
+  check bool "dispatch_used:true → combined_rb NULL declaration emitted"
+    true (contains_substr out "combined_rb = NULL")
+
+(** Test: multiple ring buffers – every buffer gets its own event handler wrapper *)
+let test_handlers_multiple_ringbufs () =
+  let decl1 = make_rb_decl "net_events"  (Ir.IRStruct ("NetEvent",  [("src_ip", Ir.IRU32)])) in
+  let decl2 = make_rb_decl "sec_events"  (Ir.IRStruct ("SecEvent",  [("severity", Ir.IRU32)])) in
+  let registry = make_registry
+    ~event_handler_registrations:[
+      ("net_events", "handle_net");
+      ("sec_events", "handle_sec");
+    ]
+    [decl1; decl2] in
+  let out = Userspace_codegen.generate_ringbuf_handlers_from_registry
+              registry ~dispatch_used:true in
+  check bool "multiple: net_events_event_handler generated"
+    true (contains_substr out "net_events_event_handler");
+  check bool "multiple: sec_events_event_handler generated"
+    true (contains_substr out "sec_events_event_handler");
+  check bool "multiple: handle_net referenced"
+    true (contains_substr out "handle_net");
+  check bool "multiple: handle_sec referenced"
+    true (contains_substr out "handle_sec")
+
 (** Run all tests *)
 let () =
   run "Ring Buffer Tests" [
@@ -1151,4 +1252,12 @@ let () =
       test_case "complex operations" `Quick test_ringbuf_complex_operations;
       test_case "on_event and dispatch integration" `Quick test_on_event_dispatch_integration;
     ];
+    "handler registry unit tests", [
+      test_case "empty registry" `Quick test_handlers_empty_registry;
+      test_case "None branch fallback to _callback" `Quick test_handlers_none_branch_fallback;
+      test_case "Some branch registered handler name" `Quick test_handlers_some_branch_registered_name;
+      test_case "dispatch_used false suppresses wrappers" `Quick test_handlers_dispatch_false_no_wrappers;
+      test_case "dispatch_used true emits combined_rb" `Quick test_handlers_dispatch_true_combined_rb;
+      test_case "multiple ring buffers each get wrapper" `Quick test_handlers_multiple_ringbufs;
+    ];
   ]