Started implementing early returns (todo: typecheck, assemble)

parser.mly

@@ -130,13 +130,8 @@ dimlist:
   { Dim (name, loc $startpos(name) $endpos(name)) :: head }
 
 fun_body:
-  | llist(var_dec) local_fun_dec* statement* loption(return_statement)
-  { VarDecs (List.rev $1) :: LocalFuns $2 :: ($3 @ $4) }
-
-return_statement:
-  (* return statement: use location of return value *)
-  | RETURN value=expr SEMICOL
-  { [Return (value, loc $startpos(value) $endpos(value))] }
+  | llist(var_dec) local_fun_dec* statement*
+  { VarDecs (List.rev $1) :: LocalFuns $2 :: $3 }
 
 (* function: use location of function name *)
 local_fun_dec:
@@ -195,6 +190,10 @@ statement:
   { let loc = loc $startpos(cnt) $endpos(cnt) in
     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)) }
+
 block:
   | LBRACE stats=statement* RBRACE  { stats }
   | stat=statement                  { [stat] }

phases/assemble.ml

@@ -89,8 +89,9 @@ let assemble program =
       let store = Store (typeof dec, get_scope node dec, indexof dec) in
       trav value @ [InlineComment (store, node2str node)]
 
-    | Return (value, _) ->
-      trav value @ [InlineComment (Ret (typeof value), node2str node)]
+    | Return (expr, _) ->
+      let ty = match expr with Some value -> typeof value | _ -> None in
+      trav value @ [InlineComment (Ret ty, node2str node)]
 
     | If (cond, body, _) ->
       let endlabel = genlabel "end" in

stringify.ml

@@ -93,8 +93,10 @@ and node2str node =
     name ^ "[" ^ concat ", " dims ^ "] = " ^ str value ^ ";"
   | Expr expr ->
     str expr ^ ";"
-  | Return (value, _) ->
-    "return " ^ (str value) ^ ";"
+  | Return (None, _) ->
+    "return;"
+  | Return (Some value, _) ->
+    "return " ^ str value ^ ";"
   | If (cond, body, _) ->
     "if (" ^ str cond ^ ") " ^ str body
   | IfElse (cond, true_body, false_body, _) ->

types.mli

@@ -87,8 +87,8 @@ and node =
   | For of string * node * node * node * node * ann
     (** counter, start, stop, step, body
         [for (int <counter> = <start>, <stop>[, <step>]) <body>] *)
-  | Return of node * ann
-    (** Return value [return <value>;] *)
+  | Return of node option * ann
+    (** Return value [return[ <value>];] *)
   | Expr of node
     (** Expression statement [<expr>;] *)
   | Block of node list

util.ml

@@ -146,9 +146,9 @@ let rec traverse u ( *: ) trav node =
     let dims, res_dims = trav_all dims in
     let value, res_value = trav value in
     (VarLet (dec, Some dims, value, ann), res_dims *: res_value)
-  | Return (value, ann) ->
+  | Return (Some value, ann) ->
     let value, res_value = trav value in
-    (Return (value, ann), res_value)
+    (Return (Some value, ann), res_value)
   | If (cond, body, ann) ->
     let cond, res_cond = trav cond in
     let body, res_body = trav body in