Implement rest of early returns + add some tests

parser.mly

@@ -191,8 +191,10 @@ statement:
     For (cnt, start, stop, step, Block body, loc) }
 
   (* return statement: use location of return value *)
-  | RETURN value=expr? SEMICOL
-  { Return (value, loc $startpos(value) $endpos(value)) }
+  | RETURN value=expr SEMICOL
+  { Return (Some value, loc $startpos(value) $endpos(value)) }
+  | ret=RETURN SEMICOL
+  { Return (None, loc $startpos(ret) $endpos(ret)) }
 
 block:
   | LBRACE stats=statement* RBRACE  { stats }

phases/assemble.ml

@@ -57,6 +57,16 @@ let assemble program =
        ImportFun (name, ret_type, List.map typeof params)]
 
     | FunDef (export, ret_type, name, params, body, _) ->
+      let last_if_return body =
+        match List.fold_left (fun _ x -> x) DummyNode (block_body body) with
+        | Return _ as ret -> Some ret
+        | _               -> None
+      in
+      let add_explicit_return body =
+        match last_if_return body with
+        | Some _ -> body
+        | None   -> Block (block_body body @ [Return (None, [])])
+      in
       let label = labelof node in
       begin
         if export
@@ -68,8 +78,7 @@ let assemble program =
        Label label] @
       (match indexof node with 0 -> [] | n -> [RtnEnter n]) @
       (trav_all params) @
-      (trav body) @
-      (match ret_type with Void -> [Ret Void] | _ -> []) @
+      (trav (add_explicit_return body)) @
       [EmptyLine] @
       (traverse_localfuns body)
 
@@ -89,9 +98,11 @@ let assemble program =
       let store = Store (typeof dec, get_scope node dec, indexof dec) in
       trav value @ [InlineComment (store, node2str node)]
 
-    | Return (expr, _) ->
-      let ty = match expr with Some value -> typeof value | _ -> None in
-      trav value @ [InlineComment (Ret ty, node2str node)]
+    | Return (None, _) ->
+      [InlineComment (Ret Void, node2str node)]
+
+    | Return (Some value, _) ->
+      trav value @ [InlineComment (Ret (typeof value), node2str node)]
 
     | If (cond, body, _) ->
       let endlabel = genlabel "end" in

phases/context.ml

@@ -123,6 +123,8 @@ let analyse do_rename program =
     | Program (decls, ann) ->
       Program (Block decls |> trav |> trav_funs scope |> block_body, ann)
 
+    (* Do not traverse into external function declarations, since their
+     * parameters must not be added to the namespace *)
     | FunDec _ | FunDef _ -> node
 
     | GlobalDef (export, ctype, name, init, ann) ->
@@ -223,10 +225,6 @@ let analyse do_rename program =
     | Allocate (dec, dims, ann) ->
       Allocate (Hashtbl.find vars (nameof dec), List.map trav dims, ann)
 
-    (* Do not traverse into external function declarations, since their
-     * parameters must not be added to the namespace *)
-    | FunDec _ -> node
-
     | _ -> traverse_unit trav node
 
   and trav_funs scope = function

phases/typecheck.ml

@@ -172,31 +172,29 @@ let rec typecheck node =
   | Dim (name, ann) ->
     (Dim (name, Type Int :: ann), [])
 
-  (* Void functions may have no return statement, other functions must have a
-   * return statement of valid type *)
+  (* Void functions may have a return statement without a value, other functions
+   * must have a return statement of a valid type *)
   | FunDef (export, ret_type, name, params, body, ann) ->
     let params, perr = err_map typecheck params in
     let body, berr = typecheck body in
-    let rec find_return = function
-      | []                         -> None
-      | [Return (value, _) as ret] -> Some (ret, typeof value)
-      | hd :: tl                   -> find_return tl
+    let rec check_types = function
+      | Return (Some _, _) as ret when ret_type = Void ->
+        ret, [NodeMsg (ret, "void function should not have a return value")]
+
+      | Return (Some value, _) as ret when typeof value <> ret_type ->
+        ret, [NodeMsg (ret, sprintf
+          "function \"%s\" has return type %s, got %s"
+          name (type2str ret_type) (type2str (typeof value)))]
+
+      | node -> traverse_list check_types node
     in
     let err =
-      match (ret_type, find_return (block_body body)) with
-      | (Void, Some (ret, _)) ->
-        [NodeMsg (ret, "void function should not have a return value")]
-
-      | ((Bool | Int | Float), None) ->
+      snd (check_types body) @
+      match last_if_return body with
+      | None when ret_type <> Void ->
         [NodeMsg (node, sprintf
-          "expected return value of type %s for function \"%s\""
+          "function \"%s\" must end with a return of type %s"
           (type2str ret_type) name)]
-
-      | ((Bool | Int | Float), Some (ret, t)) when t != ret_type ->
-        [NodeMsg (ret, sprintf
-          "function \"%s\" has return type %s, got %s"
-          name (type2str ret_type) (type2str t))]
-
       | _ -> []
     in
     (FunDef (export, ret_type, name, params, body, ann), perr @ berr @ err)

test/basic/check_error/only_early_return.cvc

@@ -0,0 +1,4 @@
+int foo(int p) {
+    if (p > 1)
+        return 1;
+}

test/basic/check_success/early_return.cvc

@@ -0,0 +1,6 @@
+void foo(int p) {
+    int i = 1;
+    if (p > 1)
+        return;
+    i = 2;
+}

test/basic/check_success/void_return.cvc

@@ -0,0 +1,3 @@
+void baz() {
+    return;
+}

test/basic/functional/early_return.cvc

@@ -0,0 +1,24 @@
+extern void printInt(int val);
+extern void printNewlines(int num);
+
+void foo(int p) {
+    int i = 1;
+    if (p < 2)
+        return;
+    i = -p;
+    printInt(i);
+    printNewlines(1);
+}
+
+void bar() {
+    return;
+    printInt(1);
+    printNewlines(1);
+}
+
+export int main() {
+    foo(0);
+    foo(2);  // prints -2
+    bar();
+    return 0;
+}

test/basic/functional/early_return.out

@@ -0,0 +1 @@
+-2

util.ml

@@ -464,9 +464,13 @@ let prerr_loc_msg loc msg =
 
 let block_body = function
   | Block nodes -> nodes
-  (*| node -> [node]*)
   | _ -> raise InvalidNode
 
+let last_if_return body =
+  match List.fold_left (fun _ x -> x) DummyNode (block_body body) with
+  | Return _ as ret -> Some ret
+  | _               -> None
+
 let basetypeof node = match typeof node with
   | ArrayDims (ctype, _)
   | Array ctype

util.mli

@@ -70,6 +70,9 @@ val flatten_blocks : node list -> node list
 (** Extract the list of child nodes from a [Block] node. *)
 val block_body : node -> node list
 
+(** Get the last node of a [Block] node if it is a [Return]. *)
+val last_if_return : node -> node option
+
 (** {2 AST node annotations} *)
 
 (** Add a single annotation to a node. *)