feat(compiler)!: Early return (#1464)

compiler/src/codegen/compcore.re

@@ -2815,6 +2815,14 @@ let do_backpatches = (wasm_mod, env, backpatches) => {
   );
 };
 
+let current_function = ref(None);
+let get_current_function = () => {
+  switch (current_function^) {
+  | Some(func) => func
+  | None => failwith("No current function set")
+  };
+};
+let set_current_function = func => current_function := Some(func);
 let loop_stack = ref([]: list((string, string)));
 
 let rec compile_store = (wasm_mod, env, binds) => {
@@ -3151,6 +3159,24 @@ and compile_instr = (wasm_mod, env, instr) =>
       Expression.Const.make(wasm_mod, const_void()),
     );
   | MError(err, args) => call_error_handler(wasm_mod, env, err, args)
+  | MReturn(value) =>
+    let current_function = get_current_function();
+    let value =
+      Option.fold(
+        ~none=Expression.Const.make(wasm_mod, const_void()),
+        ~some=compile_instr(wasm_mod, env),
+        value,
+      );
+    Expression.Return.make(
+      wasm_mod,
+      cleanup_locals(
+        wasm_mod,
+        env,
+        value,
+        current_function.args,
+        current_function.return_type,
+      ),
+    );
   | MArityOp(_) => failwith("NYI: (compile_instr): MArityOp")
   | MTagOp(_) => failwith("NYI: (compile_instr): MTagOp")
   };
@@ -3325,9 +3351,10 @@ let compile_function =
       ~preamble=?,
       wasm_mod,
       env,
-      {id, args, return_type, stack_size, body: body_instrs, func_loc},
+      {id, args, return_type, stack_size, body: body_instrs, func_loc} as func,
     ) => {
   sources := [];
+  set_current_function(func);
   let arity = List.length(args);
   let func_name =
     switch (name) {

compiler/src/codegen/mashtree.re

@@ -426,6 +426,7 @@ and instr_desc =
   | MFor(option(block), option(block), block)
   | MContinue
   | MBreak
+  | MReturn(option(instr))
   | MSwitch(immediate, list((int32, block)), block, Types.allocation_type) /* value, branches, default, return type */
   | MPrim0(prim0)
   | MPrim1(prim1, immediate)

compiler/src/codegen/transl_anf.re

@@ -233,6 +233,7 @@ module RegisterAllocation = {
         )
       | MContinue => MContinue
       | MBreak => MBreak
+      | MReturn(v) => MReturn(Option.map(apply_allocations(ty, allocs), v))
       | MSwitch(v, bs, d, ty) =>
         MSwitch(
           apply_allocation_to_imm(v),
@@ -322,6 +323,7 @@ let run_register_allocation = (instrs: list(Mashtree.instr)) => {
     | MPrim0(_)
     | MContinue
     | MBreak => []
+    | MReturn(v) => Option.fold(~none=[], ~some=live_locals, v)
     | MSwitch(v, bs, d, ty) =>
       imm_live_local(v)
       @ List.concat(List.map(((_, b)) => block_live_locals(b), bs))
@@ -787,6 +789,7 @@ let rec compile_comp = (~id=?, env, c) => {
       )
     | CContinue => MContinue
     | CBreak => MBreak
+    | CReturn(e) => MReturn(Option.map(compile_comp(env), e))
     | CPrim0(p0) => MPrim0(p0)
     | CPrim1(Box, arg)
     | CPrim1(BoxBind, arg) => MAllocate(MBox(compile_imm(env, arg)))

compiler/src/formatting/debug.re

@@ -76,6 +76,7 @@ let debug_expression = (expr: Parsetree.expression) => {
     print_loc("PExpFor", expr.pexp_loc)
   | PExpContinue => print_loc("PExpContinue", expr.pexp_loc)
   | PExpBreak => print_loc("PExpBreak", expr.pexp_loc)
+  | PExpReturn(expression) => print_loc("PExpReturn", expr.pexp_loc)
   | PExpConstraint(expression, parsed_type) =>
     print_loc("PExpConstraint", expr.pexp_loc)
   | PExpLambda(patterns, expression) =>

compiler/src/formatting/format.re

@@ -3473,6 +3473,23 @@ and print_expression_inner =
       ]);
     | PExpContinue => Doc.text("continue")
     | PExpBreak => Doc.text("break")
+    | PExpReturn(expr) =>
+      Doc.concat([
+        Doc.text("return"),
+        switch (expr) {
+        | Some(expr) =>
+          Doc.concat([
+            Doc.space,
+            print_expression(
+              ~expression_parent=GenericExpression,
+              ~original_source,
+              ~comments,
+              expr,
+            ),
+          ])
+        | None => Doc.nil
+        },
+      ])
     | PExpConstraint(expression, parsed_type) =>
       let comments_in_expression =
         Comment_utils.get_comments_inside_location(

compiler/src/middle_end/analyze_free_vars.re

@@ -60,6 +60,8 @@ module FreeVarsArg: Anf_iterator.IterArgument = {
         Ident.Set.union(cond, Ident.Set.union(inc, body));
       | CContinue
       | CBreak => Ident.Set.empty
+      | CReturn(expr) =>
+        Option.fold(~none=Ident.Set.empty, ~some=comp_free_vars, expr)
       | CSwitch(arg, branches, _) =>
         List.fold_left(
           (acc, (_, b)) => Ident.Set.union(anf_free_vars(b), acc),

compiler/src/middle_end/analyze_purity.re

@@ -122,7 +122,8 @@ module PurityArg: Anf_iterator.IterArgument = {
         && Option.fold(~none=true, ~some=anf_expression_purity_internal, inc)
         && anf_expression_purity_internal(body)
       | CContinue
-      | CBreak => false
+      | CBreak
+      | CReturn(_) => false
       | CSwitch(_, branches, _) =>
         let branches_purities =
           List.map(

compiler/src/middle_end/analyze_tail_calls.re

@@ -54,6 +54,9 @@ let rec analyze_comp_expression =
       push_tail_call(analyses);
     };
     false;
+  // An explicit return is definitionally in tail position
+  | CReturn(Some(e)) => analyze_comp_expression(true, e)
+  | CReturn(None)
   | CBoxAssign(_)
   | CAssign(_)
   | CLocalAssign(_)

compiler/src/middle_end/anf_helper.re

@@ -168,6 +168,14 @@ module Comp = {
       ~env?,
       CBreak,
     );
+  let return = (~loc=?, ~attributes=?, ~env=?, ret) =>
+    mk(
+      ~loc?,
+      ~attributes?,
+      ~allocation_type=Unmanaged(WasmI32),
+      ~env?,
+      CReturn(ret),
+    );
   let switch_ =
       (
         ~loc=?,

compiler/src/middle_end/anf_helper.rei

@@ -263,6 +263,14 @@ module Comp: {
   let break:
     (~loc: loc=?, ~attributes: attributes=?, ~env: env=?, unit) =>
     comp_expression;
+  let return:
+    (
+      ~loc: loc=?,
+      ~attributes: attributes=?,
+      ~env: env=?,
+      option(comp_expression)
+    ) =>
+    comp_expression;
   let switch_:
     (
       ~loc: loc=?,

compiler/src/middle_end/anf_iterator.re

@@ -91,6 +91,7 @@ module MakeIter = (Iter: IterArgument) => {
       iter_anf_expression(body);
     | CContinue
     | CBreak => ()
+    | CReturn(e) => Option.iter(iter_comp_expression, e)
     | CSwitch(c, branches, _) =>
       iter_imm_expression(c);
       List.iter(((_, body)) => iter_anf_expression(body), branches);

compiler/src/middle_end/anf_mapper.re

@@ -36,7 +36,8 @@ and comp_marker =
       comp_expression,
     )
   | Switch(list((int, anf_expression)) => comp_expression)
-  | Lambda(anf_expression => comp_expression);
+  | Lambda(anf_expression => comp_expression)
+  | Return(option(comp_expression) => comp_expression);
 
 let inputs = ref([]);
 let outputs = ref([]);
@@ -214,6 +215,11 @@ module MakeMap = (Iter: MapArgument) => {
       push_input(OptNode(Option.map(cond => AnfNode(cond), cond)));
     | CContinue => leave_with(CContinue)
     | CBreak => leave_with(CBreak)
+    | CReturn(expr) =>
+      push_input(
+        CompMarker(Return(expr => {...c, comp_desc: CReturn(expr)})),
+      );
+      push_input(OptNode(Option.map(expr => CompNode(expr), expr)));
     | CSwitch(cond, branches, partial) =>
       let cond = process_imm_expression(cond);
       push_input(
@@ -402,6 +408,19 @@ module MakeMap = (Iter: MapArgument) => {
           outputs := [CompNode(node), ...rest];
         | _ => failwith("Impossible: invalid output stack")
         }
+      | CompMarker(Return(f)) =>
+        switch (outputs^) {
+        | [OptNode(expr), ...rest] =>
+          let expr =
+            switch (expr) {
+            | Some(CompNode(expr)) => Some(expr)
+            | None => None
+            | _ => failwith("Impossible: invalid output stack")
+            };
+          let node = Iter.leave_comp_expression(f(expr));
+          outputs := [CompNode(node), ...rest];
+        | _ => failwith("Impossible: invalid output stack")
+        }
       };
       process_worklist();
     };

compiler/src/middle_end/anftree.re

@@ -341,6 +341,7 @@ and comp_expression_desc =
   | CFor(option(anf_expression), option(anf_expression), anf_expression)
   | CContinue
   | CBreak
+  | CReturn(option(comp_expression))
   | CSwitch(imm_expression, list((int, anf_expression)), partial)
   | CApp(
       (imm_expression, (list(allocation_type), allocation_type)),

compiler/src/middle_end/anftree.rei

@@ -321,6 +321,7 @@ and comp_expression_desc =
   | CFor(option(anf_expression), option(anf_expression), anf_expression)
   | CContinue
   | CBreak
+  | CReturn(option(comp_expression))
   | CSwitch(imm_expression, list((int, anf_expression)), partial)
   | CApp(
       (imm_expression, (list(allocation_type), allocation_type)),

compiler/src/middle_end/linearize.re

@@ -553,6 +553,18 @@ let rec transl_imm =
       Imm.const(Const_void),
       [BSeq(Comp.break(~loc, ~env, ()))],
     )
+  | TExpReturn(value) =>
+    let (value_comp, value_setup) =
+      switch (value) {
+      | Some(value) =>
+        let (value_comp, value_setup) = transl_comp_expression(value);
+        (Some(value_comp), value_setup);
+      | None => (None, [])
+      };
+    (
+      Imm.const(Const_void),
+      value_setup @ [BSeq(Comp.return(~loc, ~env, value_comp))],
+    );
   | TExpApp(
       {exp_desc: TExpIdent(_, _, {val_kind: TValPrim("@throw")})},
       _,

compiler/src/parsing/ast_helper.re

@@ -269,6 +269,8 @@ module Exp = {
     mk(~loc?, ~attributes?, PExpContinue);
   let break = (~loc=?, ~attributes=?, ()) =>
     mk(~loc?, ~attributes?, PExpBreak);
+  let return = (~loc=?, ~attributes=?, a) =>
+    mk(~loc?, ~attributes?, PExpReturn(a));
   let constraint_ = (~loc=?, ~attributes=?, a, b) =>
     mk(~loc?, ~attributes?, PExpConstraint(a, b));
   let box_assign = (~loc=?, ~attributes=?, a, b) =>

compiler/src/parsing/ast_helper.rei

@@ -198,6 +198,8 @@ module Exp: {
     expression;
   let continue: (~loc: loc=?, ~attributes: attributes=?, unit) => expression;
   let break: (~loc: loc=?, ~attributes: attributes=?, unit) => expression;
+  let return:
+    (~loc: loc=?, ~attributes: attributes=?, option(expression)) => expression;
   let constraint_:
     (~loc: loc=?, ~attributes: attributes=?, expression, parsed_type) =>
     expression;

compiler/src/parsing/ast_iterator.re

@@ -108,6 +108,7 @@ module E = {
       sub.expr(sub, b);
     | PExpContinue
     | PExpBreak => ()
+    | PExpReturn(e) => Option.iter(sub.expr(sub), e)
     | PExpConstraint(e, t) =>
       sub.expr(sub, e);
       sub.typ(sub, t);

compiler/src/parsing/ast_mapper.re

@@ -123,6 +123,8 @@ module E = {
       )
     | PExpContinue => continue(~loc, ~attributes, ())
     | PExpBreak => break(~loc, ~attributes, ())
+    | PExpReturn(e) =>
+      return(~loc, ~attributes, Option.map(sub.expr(sub), e))
     | PExpLambda(pl, e) =>
       lambda(
         ~loc,

compiler/src/parsing/lexer.re

@@ -220,6 +220,7 @@ let rec token = lexbuf => {
   | "for" => positioned(FOR)
   | "continue" => positioned(CONTINUE)
   | "break" => positioned(BREAK)
+  | "return" => positioned(RETURN)
   | "if" => positioned(IF)
   | "when" => positioned(WHEN)
   | "else" => positioned(ELSE)

compiler/src/parsing/parser.messages

@@ -5268,6 +5268,19 @@ program: LIDENT EQUAL WHEN
 
 Expected an expression to assign.
 
+program: RETURN WHEN
+##
+## Ends in an error in state: 13.
+##
+## stmt_expr -> RETURN . [ THICKARROW STAR SLASH SEMI RPAREN RCARET RBRACK RBRACE PIPE LCARET INFIX_90 INFIX_80 INFIX_70 INFIX_60 INFIX_50 INFIX_40 INFIX_30 INFIX_120 INFIX_110 INFIX_100 EOL EOF ELSE DASH COMMA COLON ]
+## stmt_expr -> RETURN . expr [ THICKARROW STAR SLASH SEMI RPAREN RCARET RBRACK RBRACE PIPE LCARET INFIX_90 INFIX_80 INFIX_70 INFIX_60 INFIX_50 INFIX_40 INFIX_30 INFIX_120 INFIX_110 INFIX_100 EOL EOF ELSE DASH COMMA COLON ]
+##
+## The known suffix of the stack is as follows:
+## RETURN
+##
+
+Expected an expression to return.
+
 program: WASMI64 INFIX_70 WHEN
 ##
 ## Ends in an error in state: 567.

compiler/src/parsing/parser.mly

@@ -28,7 +28,7 @@ module Grain_parsing = struct end
 
 %token TRUE FALSE VOID
 
-%token LET MUT REC IF WHEN ELSE MATCH WHILE FOR CONTINUE BREAK
+%token LET MUT REC IF WHEN ELSE MATCH WHILE FOR CONTINUE BREAK RETURN
 %token AT
 
 %token <string> INFIX_10 INFIX_30 INFIX_40 INFIX_50 INFIX_60 INFIX_70
@@ -63,7 +63,7 @@ module Grain_parsing = struct end
 %left INFIX_110 DASH
 %left INFIX_120 STAR SLASH
 
-%right SEMI EOL COMMA DOT COLON LPAREN
+%right SEMI EOL COMMA DOT COLON LPAREN RETURN
 
 %nonassoc _if
 %nonassoc ELSE
@@ -531,6 +531,7 @@ stmt_expr:
   | THROW expr { Exp.apply ~loc:(to_loc $loc) (mkid_expr $loc($1) [mkstr $loc($1) "throw"]) [$2] }
   | ASSERT expr { Exp.apply ~loc:(to_loc $loc) (mkid_expr $loc($1) [mkstr $loc($1) "assert"]) [$2] }
   | FAIL expr { Exp.apply ~loc:(to_loc $loc) (mkid_expr $loc($1) [mkstr $loc($1) "fail"]) [$2] }
+  | RETURN ioption(expr) { Exp.return ~loc:(to_loc $loc) $2 }
   | CONTINUE { Exp.continue ~loc:(to_loc $loc) () }
   | BREAK { Exp.break ~loc:(to_loc $loc) () }
 

compiler/src/parsing/parsetree.re

@@ -455,6 +455,7 @@ and expression_desc =
     )
   | PExpContinue
   | PExpBreak
+  | PExpReturn(option(expression))
   | PExpConstraint(expression, parsed_type)
   | PExpLambda(list(pattern), expression)
   | PExpApp(expression, list(expression))