Fix function calls in variable initializers return directly to the variable's register

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

Author: congwang-mk

src/ebpf_c_codegen.ml

@@ -2773,8 +2773,35 @@ let generate_c_basic_block ctx ir_block =
     increase_indent ctx
   );
 
-  (* Emit instructions *)
-  List.iter (generate_c_instruction ctx) ir_block.instructions
+  (* Emit instructions with special handling for function call + variable declaration sequences *)
+  let rec emit_instructions_optimized = function
+    | [] -> ()
+    | call_instr :: decl_instr :: rest 
+      when (match call_instr.instr_desc, decl_instr.instr_desc with
+        | IRCall (call_target, args, Some result_val), IRDeclareVariable (dest_val, typ, None) ->
+            (* Check if the call result register matches the declaration register *)
+            (match result_val.value_desc, dest_val.value_desc with
+             | IRRegister call_reg, IRRegister decl_reg when call_reg = decl_reg ->
+                 (* Generate combined declaration with function call initialization *)
+                 let var_name = get_meaningful_var_name ctx decl_reg typ in
+                 let args_str = String.concat ", " (List.map (generate_c_value ctx) args) in
+                 let call_str = match call_target with
+                   | DirectCall name -> sprintf "%s(%s)" name args_str
+                   | FunctionPointerCall func_ptr -> sprintf "(*%s)(%s)" (generate_c_value ctx func_ptr) args_str
+                 in
+                 let decl_str = generate_ebpf_c_declaration typ var_name in
+                 emit_line ctx (sprintf "%s = %s;" decl_str call_str);
+                 true
+             | _ -> false)
+        | _ -> false) ->
+        (* Skip the next instruction since we handled it in the combined declaration *)
+        emit_instructions_optimized rest
+    | instr :: rest ->
+        (* Regular instruction processing *)
+        generate_c_instruction ctx instr;
+        emit_instructions_optimized rest
+  in
+  emit_instructions_optimized ir_block.instructions
 
 (** Collect mapping from registers to variable names *)
 let collect_register_variable_mapping ir_func =

src/ir_generator.ml

@@ -1017,20 +1017,67 @@ and resolve_declaration_type_and_init ctx reg typ_opt expr_opt =
   | Some ast_type, Some expr ->
       (* Use explicitly declared type, but process initialization expression *)
       let target_type = resolve_type_alias ctx reg ast_type in
-      let value = lower_expression ctx expr in
-      (target_type, Some value)
+      (* For function calls, manually handle them to use the target register *)
+      (match expr.Ast.expr_desc with
+       | Ast.Call (callee_expr, args) ->
+           (* Handle function call that should return to the target register *)
+           let arg_vals = List.map (lower_expression ctx) args in
+           let result_val = make_ir_value (IRRegister reg) target_type expr.Ast.expr_pos in
+           let call_target = match callee_expr.Ast.expr_desc with
+             | Ast.Identifier name ->
+                 if Hashtbl.mem ctx.variables name || Hashtbl.mem ctx.function_parameters name then
+                   let callee_val = lower_expression ctx callee_expr in
+                   FunctionPointerCall callee_val
+                 else
+                   DirectCall name
+             | _ ->
+                 let callee_val = lower_expression ctx callee_expr in
+                 FunctionPointerCall callee_val
+           in
+           let instr = make_ir_instruction (IRCall (call_target, arg_vals, Some result_val)) expr.Ast.expr_pos in
+           emit_instruction ctx instr;
+           (target_type, None)
+       | _ ->
+           (* Non-function call - use normal processing *)
+           let value = lower_expression ctx expr in
+           (target_type, Some value))
   | None, Some expr ->
       (* No declared type - use type checker annotation if available, otherwise infer from expression *)
-      let value = lower_expression ctx expr in
-      let inferred_type = match expr.expr_type with
-        | Some ast_type -> 
-            (* Prioritize type checker annotation as single source of truth *)
-            ast_type_to_ir_type_with_context ctx.symbol_table ast_type
-        | None -> 
-            (* Fallback to IR type inference only when type checker didn't provide annotation *)
-            value.val_type
-      in
-      (inferred_type, Some value)
+      (match expr.Ast.expr_desc with
+       | Ast.Call (callee_expr, args) ->
+           (* Handle function call in type inference *)
+           let inferred_type = match expr.Ast.expr_type with
+             | Some ast_type -> ast_type_to_ir_type_with_context ctx.symbol_table ast_type
+             | None -> IRU32 (* Default fallback *)
+           in
+           let arg_vals = List.map (lower_expression ctx) args in
+           let result_val = make_ir_value (IRRegister reg) inferred_type expr.Ast.expr_pos in
+           let call_target = match callee_expr.Ast.expr_desc with
+             | Ast.Identifier name ->
+                 if Hashtbl.mem ctx.variables name || Hashtbl.mem ctx.function_parameters name then
+                   let callee_val = lower_expression ctx callee_expr in
+                   FunctionPointerCall callee_val
+                 else
+                   DirectCall name
+             | _ ->
+                 let callee_val = lower_expression ctx callee_expr in
+                 FunctionPointerCall callee_val
+           in
+           let instr = make_ir_instruction (IRCall (call_target, arg_vals, Some result_val)) expr.Ast.expr_pos in
+           emit_instruction ctx instr;
+           (inferred_type, None)
+       | _ ->
+           (* Non-function call - use normal processing *)
+           let value = lower_expression ctx expr in
+           let inferred_type = match expr.Ast.expr_type with
+             | Some ast_type -> 
+                 (* Prioritize type checker annotation as single source of truth *)
+                 ast_type_to_ir_type_with_context ctx.symbol_table ast_type
+             | None -> 
+                 (* Fallback to IR type inference only when type checker didn't provide annotation *)
+                 value.val_type
+           in
+           (inferred_type, Some value))
   | Some ast_type, None ->
       (* Declared type, no initialization *)
       let target_type = resolve_type_alias ctx reg ast_type in

tests/test_ebpf_c_codegen.ml

@@ -925,6 +925,31 @@ let test_map_field_access_pointer_fix () =
   check bool "dot notation used for regular struct field access" true (contains_substr regular_result "my_struct.size");
   ()
 
+(** Test variable declaration with function call initialization *)
+let test_variable_function_call_declaration () =
+  let ctx = create_c_context () in
+  ctx.indent_level <- 1; (* Set valid indent level *)
+  
+  (* Create a function call that returns to a register *)
+  let result_reg = 0 in
+  let result_val = make_ir_value (IRRegister result_reg) IRU32 test_pos in
+  let call_instr = make_ir_instruction (IRCall (DirectCall "helper_function", [make_ir_value (IRLiteral (IntLit (5, None))) IRU32 test_pos], Some result_val)) test_pos in
+  
+  (* Create a variable declaration for the same register with no initialization *)
+  let decl_instr = make_ir_instruction (IRDeclareVariable (result_val, IRU32, None)) test_pos in
+  
+  (* Test the optimization that combines these into a single declaration *)
+  let ir_block = make_ir_basic_block "test" [call_instr; decl_instr] 0 in
+  generate_c_basic_block ctx ir_block;
+  
+  let output = String.concat "\n" ctx.output_lines in
+  
+  (* Should generate: __u32 val_0 = helper_function(5); *)
+  check bool "combined declaration with function call" true (contains_substr output "val_0 = helper_function(5)");
+  
+  (* Should NOT generate separate variable declaration without initialization *)
+  check bool "no uninitialized declaration" false (contains_substr output "__u32 val_0;")
+
 (** Test suite definition *)
 let suite =
   [
@@ -962,6 +987,7 @@ let suite =
     ("String size collection from userspace structs", `Quick, test_string_size_collection_from_userspace_structs);
     ("Declaration ordering fix", `Quick, test_declaration_ordering_fix);
     ("BPF printk string literal fix", `Quick, test_bpf_printk_string_literal_fix);
+    ("Variable function call declaration", `Quick, test_variable_function_call_declaration);
   ]
 
 (** Run all tests *)

tests/test_ir.ml

@@ -208,13 +208,84 @@ let test_userspace_binding_generation () =
   | None ->
       fail "No C bindings found"
 
+let test_variable_function_call_initialization () =
+  (* Test for the bug where function calls in variable initializers 
+     return to wrong registers, causing uninitialized variable usage *)
+  let input = {|
+@xdp fn test_handler(ctx: *xdp_md) -> xdp_action {
+    return 2  // XDP_PASS
+}
+
+fn main() -> i32 {
+    var prog = load(test_handler)  // Should assign to same register as 'prog'
+    var result = attach(prog, "eth0", 0)  // Should use 'prog' register correctly
+    return result
+}
+|} in
+
+  try
+    let ast = Kernelscript.Parse.parse_string input in
+    let symbol_table = Kernelscript.Symbol_table.build_symbol_table ast in
+    let (typed_ast, _) = Kernelscript.Type_checker.type_check_and_annotate_ast ~symbol_table:(Some symbol_table) ast in
+    let ir_multi_prog = generate_ir typed_ast symbol_table "test_var_func_init" in
+    
+    (* Extract the main function from userspace program *)
+    let userspace_program = match ir_multi_prog.userspace_program with
+      | Some prog -> prog
+      | None -> failwith "No userspace program found"
+    in
+    let main_func = List.find (fun func -> func.func_name = "main") userspace_program.userspace_functions in
+    
+    (* Collect all instructions from all basic blocks *)
+    let all_instructions = List.flatten (List.map (fun block -> block.instructions) main_func.basic_blocks) in
+    
+    (* Find variable declarations and function calls *)
+    let declarations = List.filter_map (fun instr ->
+      match instr.instr_desc with
+      | IRDeclareVariable (dest_val, _, _) -> Some dest_val
+      | _ -> None
+    ) all_instructions in
+    
+    let function_calls = List.filter_map (fun instr ->
+      match instr.instr_desc with
+      | IRCall (_, _, Some result_val) -> Some result_val
+      | _ -> None
+    ) all_instructions in
+    
+    (* Verify we have the expected number of declarations and calls *)
+    check int "Should have variable declarations" 2 (List.length declarations);
+    check int "Should have function calls" 2 (List.length function_calls);
+    
+    (* The key test: verify that function call returns go to the same registers as variable declarations *)
+    let get_register_from_value val_desc = match val_desc with
+      | IRRegister reg -> Some reg
+      | _ -> None
+    in
+    
+    let declaration_registers = List.filter_map (fun val_desc -> get_register_from_value val_desc.value_desc) declarations in
+    let call_result_registers = List.filter_map (fun val_desc -> get_register_from_value val_desc.value_desc) function_calls in
+    
+    (* Verify that function call results use the same registers as variable declarations *)
+    (* This catches the bug where function calls returned to different registers *)
+    check bool "Function call results should use declaration registers" true 
+      (List.for_all (fun reg -> List.mem reg declaration_registers) call_result_registers);
+    
+    (* Verify register consistency - each variable should map to exactly one register *)
+    let sorted_decl_regs = List.sort compare declaration_registers in
+    let sorted_call_regs = List.sort compare call_result_registers in
+    check (list int) "Declaration and call registers should match" sorted_decl_regs sorted_call_regs
+    
+  with
+  | e -> failwith (Printf.sprintf "Variable function call initialization test failed: %s" (Printexc.to_string e))
+
 let ir_tests = [
   "program_lowering", `Quick, test_program_lowering;
   "context_access_lowering", `Quick, test_context_access_lowering;
   "map_operation_lowering", `Quick, test_map_operation_lowering;
   "bounds_check_insertion", `Quick, test_bounds_check_insertion;
   "stack_usage_tracking", `Quick, test_stack_usage_tracking;
   "userspace_binding_generation", `Quick, test_userspace_binding_generation;
+  "variable_function_call_initialization", `Quick, test_variable_function_call_initialization;
 ]
 
 let () =