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

Author: ospencer

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=?,