Refactor return type extraction in block arms for improved type checking consistency.

Author: congwang-mk

examples/type_checking.ks

@@ -72,7 +72,7 @@ fn extract_header(ctx: *xdp_md) -> *PacketHeader {
 @helper
 fn classify_protocol(proto: u8) -> ProtocolType {
   // Type checker validates enum constant access
-  match (proto) {
+  return match (proto) {
     6: TCP,
     17: UDP,  
     1: ICMP,
@@ -99,7 +99,7 @@ fn make_decision(header: PacketHeader) -> FilterDecision {
   // Type checker validates function call signatures
   var proto_type = classify_protocol(header.protocol)
 
-  match (proto_type) {
+  return match (proto_type) {
     TCP: {
       // Type checker validates field access on struct types
       if (header.length > 1500) {
@@ -128,7 +128,7 @@ fn make_decision(header: PacketHeader) -> FilterDecision {
   var decision = make_decision(*packet_header)
 
   // Type checker validates match expressions and enum types
-  match (decision) {
+  return match (decision) {
     Allow: XDP_PASS,
     Block: XDP_DROP,
     Log: {

src/ir_generator.ml

@@ -1323,32 +1323,55 @@ and lower_statement ctx stmt =
                               let ret_val = lower_expression ctx expr in
                               IRReturnValue ret_val)
                      | Block stmts ->
-                         (* For block arms, look for the return statement *)
-                         let rec find_return_action = function
-                           | [] -> failwith "Block arm must have a return statement"
-                           | stmt :: rest ->
-                               (match stmt.stmt_desc with
-                                | Ast.Return (Some return_expr) ->
-                                    (match return_expr.expr_desc with
-                                     | Ast.Call (callee_expr, args) ->
-                                         (* Check if this is a simple function call that could be a tail call *)
-                                         (match callee_expr.expr_desc with
-                                          | Ast.Identifier name ->
-                                              let arg_vals = List.map (lower_expression ctx) args in
-                                              IRReturnCall (name, arg_vals)
-                                          | _ ->
-                                              (* Function pointer call - treat as regular return *)
-                                              let ret_val = lower_expression ctx return_expr in
-                                              IRReturnValue ret_val)
-                                     | Ast.TailCall (name, args) ->
+                         (* For block arms, extract return action from the last statement *)
+                         let rec extract_return_action_from_stmt stmt =
+                           match stmt.stmt_desc with
+                           | Ast.Return (Some return_expr) ->
+                               (match return_expr.expr_desc with
+                                | Ast.Call (callee_expr, args) ->
+                                    (* Check if this is a simple function call that could be a tail call *)
+                                    (match callee_expr.expr_desc with
+                                     | Ast.Identifier name ->
                                          let arg_vals = List.map (lower_expression ctx) args in
-                                         IRReturnTailCall (name, arg_vals, 0)
+                                         IRReturnCall (name, arg_vals)
                                      | _ ->
+                                         (* Function pointer call - treat as regular return *)
                                          let ret_val = lower_expression ctx return_expr in
                                          IRReturnValue ret_val)
-                                | _ -> find_return_action rest)
+                                | Ast.TailCall (name, args) ->
+                                    let arg_vals = List.map (lower_expression ctx) args in
+                                    IRReturnTailCall (name, arg_vals, 0)
+                                | _ ->
+                                    let ret_val = lower_expression ctx return_expr in
+                                    IRReturnValue ret_val)
+                           | Ast.ExprStmt expr ->
+                               (* Handle implicit return from expression statement *)
+                               (match expr.expr_desc with
+                                | Ast.Call (callee_expr, args) ->
+                                    (match callee_expr.expr_desc with
+                                     | Ast.Identifier name ->
+                                         let arg_vals = List.map (lower_expression ctx) args in
+                                         IRReturnCall (name, arg_vals)
+                                     | _ ->
+                                         let ret_val = lower_expression ctx expr in
+                                         IRReturnValue ret_val)
+                                | Ast.TailCall (name, args) ->
+                                    let arg_vals = List.map (lower_expression ctx) args in
+                                    IRReturnTailCall (name, arg_vals, 0)
+                                | _ ->
+                                    let ret_val = lower_expression ctx expr in
+                                    IRReturnValue ret_val)
+                           | Ast.If (_, then_stmts, Some _) ->
+                               (* For if-else statements, we'll use the then branch action (both should be compatible) *)
+                               extract_return_action_from_block then_stmts
+                           | _ ->
+                               failwith "Block arm must end with a return statement, expression, or if-else statement"
+                         and extract_return_action_from_block stmts =
+                           match List.rev stmts with
+                           | last_stmt :: _ -> extract_return_action_from_stmt last_stmt
+                           | [] -> failwith "Empty block in match arm"
                          in
-                         find_return_action stmts
+                         extract_return_action_from_block stmts
                    in
 
                    { match_pattern = ir_pattern; return_action = return_action; arm_pos = arm.arm_pos }

src/type_checker.ml

@@ -618,6 +618,31 @@ let is_truthy_type bpf_type =
   | Enum _ -> true                                        (* enums: based on numeric value *)
   | _ -> false                                            (* other types not allowed in boolean context *)
 
+(** Helper function to extract return type from a block of statements *)
+let rec extract_block_return_type stmts arm_pos =
+  let extract_type_from_stmt stmt =
+    match stmt.tstmt_desc with
+    | TReturn (Some return_expr) -> return_expr.texpr_type
+    | TExprStmt expr -> expr.texpr_type
+    | TIf (_, then_stmts, Some else_stmts) ->
+        (* For if-else statements, both branches must return compatible types *)
+        let then_type = extract_block_return_type then_stmts arm_pos in
+        let else_type = extract_block_return_type else_stmts arm_pos in
+        (match unify_types then_type else_type with
+         | Some unified_type -> unified_type
+         | None -> type_error ("If-else branches have incompatible types: " ^ 
+                               string_of_bpf_type then_type ^ " vs " ^ 
+                               string_of_bpf_type else_type) arm_pos)
+    | TIf (_, _, None) ->
+        (* If without else - this doesn't work as a return value *)
+        type_error "If statement without else cannot be used as return value in match arm" arm_pos
+    | _ -> 
+        type_error "Block arms must end with a return statement, expression, or if-else statement" arm_pos
+  in
+  match List.rev stmts with
+  | last_stmt :: _ -> extract_type_from_stmt last_stmt
+  | [] -> type_error "Empty block in match arm" arm_pos
+
 (** Type check a user function call *)
 let rec type_check_user_function_call ctx name typed_args arg_types pos =
   try
@@ -1197,19 +1222,13 @@ and type_check_expression ctx expr =
             let first_type = match first_arm.tarm_body with
               | TSingleExpr expr -> expr.texpr_type
               | TBlock stmts -> 
-                  (* For block arms, we need to find the return type from the last statement *)
-                  match List.rev stmts with
-                  | { tstmt_desc = TReturn (Some return_expr); _ } :: _ -> return_expr.texpr_type
-                  | _ -> type_error "Block arms must end with a return statement" first_arm.tarm_pos
+                  extract_block_return_type stmts first_arm.tarm_pos
             in
             List.iter (fun arm ->
               let arm_type = match arm.tarm_body with
                 | TSingleExpr expr -> expr.texpr_type
                 | TBlock stmts -> 
-                    (* For block arms, we need to find the return type from the last statement *)
-                    match List.rev stmts with
-                    | { tstmt_desc = TReturn (Some return_expr); _ } :: _ -> return_expr.texpr_type
-                    | _ -> type_error "Block arms must end with a return statement" arm.tarm_pos
+                    extract_block_return_type stmts arm.tarm_pos
               in
               match unify_types first_type arm_type with
               | Some _ -> () (* Compatible *)

tests/test_match.ml

@@ -461,6 +461,89 @@ let test_nested_match_structures () =
 
   check bool "nested match should generate nested conditional structures" true has_conditional_structure
 
+(** Test match arms with implicit returns from block expressions - bug fix test *)
+let test_match_block_implicit_returns () =
+  let input = {|
+    enum Decision {
+      Accept = 0,
+      Reject = 1,
+      Review = 2
+    }
+    
+    fn process_value(value: u32) -> Decision {
+      return match (value) {
+        1: Accept,
+        2: {
+          if (value > 10) {
+            Reject
+          } else {
+            Review  
+          }
+        },
+        3: {
+          Review
+        },
+        default: Reject
+      }
+    }
+  |} in
+  
+  let ast = Parse.parse_string input in
+  let symbol_table = Test_utils.Helpers.create_test_symbol_table ast in
+  
+  (* The main test: Type checking should succeed (this would fail before the bug fix) *)
+  (try 
+    let (_typed_ast, _typed_functions) = Type_checker.type_check_and_annotate_ast ~symbol_table:(Some symbol_table) ast in
+    check bool "type checking should succeed for match arms with implicit returns" true true
+  with
+  | Failure msg when msg = "Block arms must end with a return statement" ->
+      fail "Bug regression: type checker still requires explicit returns in match arm blocks"
+  | Failure msg -> 
+      failwith ("Type checking failed with different error: " ^ msg)
+  | exn ->
+      failwith ("Unexpected type checking error: " ^ (Printexc.to_string exn)));
+  
+  (* Verify the structure was parsed correctly *)
+  let func = match List.find (function 
+    | GlobalFunction f when f.func_name = "process_value" -> true 
+    | _ -> false) ast with
+    | GlobalFunction f -> f
+    | _ -> failwith "Expected process_value function"
+  in
+  
+  let return_stmt = List.hd func.func_body in
+  let match_expr = match return_stmt.stmt_desc with
+    | Return (Some expr) -> (match expr.expr_desc with
+        | Match (_, arms) -> arms
+        | _ -> failwith "Expected match expression")
+    | _ -> failwith "Expected return statement with match"
+  in
+  
+  (* Verify we have the expected structure *)
+  check int "should have 4 match arms" 4 (List.length match_expr);
+  
+  (* Verify the second arm has a block with if-else (implicit return) *)
+  let second_arm = List.nth match_expr 1 in
+  (match second_arm.arm_body with
+   | Block stmts -> 
+       check bool "second arm should have statements" true (List.length stmts > 0);
+       (* Verify it's an if statement (implicit return, no explicit return needed) *)
+       (match (List.hd stmts).stmt_desc with
+        | If (_, _, Some _) -> () (* if-else statement - good *)
+        | _ -> failwith "Expected if-else statement in second arm")
+   | _ -> failwith "Expected block in second arm");
+   
+  (* Verify the third arm has a block with expression (implicit return) *)
+  let third_arm = List.nth match_expr 2 in
+  (match third_arm.arm_body with
+   | Block stmts -> 
+       check bool "third arm should have statements" true (List.length stmts > 0);
+       (* Verify it's an expression statement (implicit return) *)
+       (match (List.hd stmts).stmt_desc with
+        | ExprStmt _ -> () (* expression statement - good *)
+        | _ -> failwith "Expected expression statement in third arm")
+   | _ -> failwith "Expected block in third arm")
+
 let suite = [
   "test_basic_match_parsing", `Quick, test_basic_match_parsing;
   "test_match_with_enums", `Quick, test_match_with_enums;
@@ -471,6 +554,7 @@ let suite = [
   "test_match_conditional_control_flow", `Quick, test_match_conditional_control_flow;
   "test_match_no_premature_execution", `Quick, test_match_no_premature_execution;
   "test_nested_match_structures", `Quick, test_nested_match_structures;
+  "test_match_block_implicit_returns", `Quick, test_match_block_implicit_returns;
 ]
 
 let () = run "Match Construct Tests" [