Add early return with provisional extension syntax

CHANGELOG.md

@@ -29,6 +29,7 @@
 
 #### :rocket: New Feature
 
+- Add early return from function bodies, provisionally encoded as `%return(value)`, including loops and async functions. https://github.com/rescript-lang/rescript/pull/8656
 - Support UTF-16 surrogate-pair escapes such as `"\uD83D\uDE00"` in ordinary string literals. https://github.com/rescript-lang/rescript/pull/8606
 - Support dynamic imports of external bindings annotated with `@scope`; the generated import follows the complete property path. These imports were previously rejected. https://github.com/rescript-lang/rescript/pull/8582
 - Add `@res.hoistedFunction` for emitting nested module functions as flat JavaScript exports. https://github.com/rescript-lang/rescript/pull/8402

analysis/reanalyze/src/side_effects.ml

@@ -28,7 +28,7 @@ let rec expr_no_side_effects (expr : Typedtree.expression) =
   | Texp_function _ -> true
   (* Loop control changes whether subsequent code in the enclosing loop runs,
      so it should not be treated as a removable pure expression. *)
-  | Texp_break | Texp_continue -> false
+  | Texp_return _ | Texp_break | Texp_continue -> false
   | Texp_apply {funct = {exp_desc = Texp_ident (path, _, _)}; args}
     when path |> path_is_whitelisted_for_side_effects ->
     args |> List.for_all (fun (_, eo) -> eo |> expr_opt_no_side_effects)

analysis/src/completion_front_end.ml

@@ -1253,6 +1253,9 @@ let completion_with_parser1 ~debug ~offset ~pos_cursor ~kind_file
         match expr.pexp_desc with
         | Pexp_object_literal fields ->
           List.iter (fun (_, e) -> iterator.expr iterator e) fields
+        | Pexp_extension
+            ({txt = "return"}, PStr [{pstr_desc = Pstr_eval (value, [])}]) ->
+          iterator.expr iterator value
         | Pexp_extension ({txt}, _) -> set_result (CextensionNode txt)
         | Pexp_constant _ -> set_result Cnone
         | Pexp_ident lid ->

compiler/core/lam_analysis.ml

@@ -122,7 +122,7 @@ let rec no_side_effects (lam : Lambda.t) : bool =
   | Lifthenelse (a, b, c) ->
     no_side_effects a && no_side_effects b && no_side_effects c
   | Lsequence (a, b) -> no_side_effects a && no_side_effects b
-  | Lbreak | Lcontinue -> false
+  | Lreturn _ | Lbreak | Lcontinue -> false
   | Lletrec (bindings, body) ->
     Ext_list.for_all_snd bindings no_side_effects && no_side_effects body
   | Lwhile _ ->
@@ -186,6 +186,7 @@ let rec size (lam : Lambda.t) =
     | Ltrywith _ -> really_big ()
     | Lifthenelse (l1, l2, l3) -> 1 + size l1 + size l2 + size l3
     | Lsequence (l1, l2) -> size l1 + size l2
+    | Lreturn value -> 1 + size value
     | Lbreak | Lcontinue -> 1
     | Lwhile _ -> really_big ()
     | Lfor _ -> really_big ()
@@ -249,9 +250,18 @@ let destruct_pattern (body : Lambda.t) params args =
     | Some _ | None -> false)
   | _ -> false
 
+(* Returns are relative to the nearest function, so ordinary beta reduction
+   must preserve that boundary. Nested functions retain their own boundary. *)
+let rec contains_return = function
+  | Lambda.Lreturn _ -> true
+  | Lfunction _ -> false
+  | lam -> Lambda_traverse.shallow_exists contains_return lam
+
 (* Async functions cannot be beta reduced *)
 let lfunction_can_be_inlined (lfunction : Lambda.lfunction) =
-  (not lfunction.attr.async) && lfunction.attr.directive = None
+  (not lfunction.attr.async)
+  && lfunction.attr.directive = None
+  && not (contains_return lfunction.body)
 
 (** Hints to inlining *)
 let ok_to_inline_fun_when_app (m : Lambda.lfunction) (args : Lambda.t list) =

compiler/core/lam_analysis.mli

@@ -29,6 +29,10 @@ val no_side_effects : Lambda.t -> bool
 
 val size : Lambda.t -> int
 
+val contains_return : Lambda.t -> bool
+(** Whether a return targets the surrounding function. Does not descend into
+    nested functions, which establish their own return scope. *)
+
 val lfunction_can_be_inlined : Lambda.lfunction -> bool
 
 val ok_to_inline_fun_when_app : Lambda.lfunction -> Lambda.t list -> bool

compiler/core/lam_arity_analysis.ml

@@ -123,7 +123,7 @@ let rec get_arity (meta : Lam_stats.t) (lam : Lambda.t) : Lam_arity.t =
   | Lifthenelse (_, l2, l3) -> all_lambdas meta [l2; l3]
   | Lsequence (_, l2) -> get_arity meta l2
   | Lstaticraise _ (* since it will not be in tail position *) -> Lam_arity.na
-  | Lbreak | Lcontinue -> Lam_arity.non_function_arity_info
+  | Lreturn _ | Lbreak | Lcontinue -> Lam_arity.non_function_arity_info
   | Lwhile _ | Lfor _ | Lfor_of _ | Lfor_await_of _ | Lassign _ ->
     Lam_arity.non_function_arity_info
 

compiler/core/lam_bounded_vars.ml

@@ -160,6 +160,7 @@ let rewrite (map : _ Hash_ident.t) (lam : Lambda.t) : Lambda.t =
       let l1 = aux l1 in
       let l2 = aux l2 in
       Lambda.seq l1 l2
+    | Lreturn value -> Lambda.return (aux value)
     | Lbreak -> Lambda.break
     | Lcontinue -> Lambda.continue
     | Lwhile (l1, l2) ->

compiler/core/lam_check.ml

@@ -95,6 +95,7 @@ let check ~file ~pass lam =
       check_staticfails e2 cxt
     | Lifthenelse (e1, e2, e3) -> check_list [e1; e2; e3] cxt
     | Lsequence (e1, e2) -> check_list [e1; e2] cxt
+    | Lreturn value -> check_staticfails value cxt
     | Lassign (_id, e) -> check_staticfails e cxt
   in
   let rec iter_list xs = Ext_list.iter xs iter
@@ -164,6 +165,7 @@ let check ~file ~pass lam =
       iter e1;
       def v;
       iter e2
+    | Lreturn e -> iter e
     | Lassign (id, e) ->
       use id;
       iter e

compiler/core/lam_closure.ml

@@ -143,6 +143,7 @@ let free_variables (export_idents : Set_ident.t) (params : stats Map_ident.t)
       local_add v;
       iter sink_pos e1;
       iter sink_pos e2
+    | Lreturn e -> iter top e
     | Lassign (id, e) ->
       used top id;
       iter top e

compiler/core/lam_compile.ml

@@ -42,6 +42,7 @@ let rec source_loc_of_lam (lam : Lambda.t) =
     | None -> source_loc_of_lam body)
   | Lletrec (_, body) | Lsequence (_, body) -> source_loc_of_lam body
   | Lifthenelse (_, then_, _) -> source_loc_of_lam then_
+  | Lreturn body -> source_loc_of_lam body
   | Lstaticcatch (body, _, _) | Ltrywith (body, _, _) -> source_loc_of_lam body
   | Lstringswitch (_, cases, default) -> (
     match cases with
@@ -1190,19 +1191,22 @@ let compile output_prefix =
   *)
   and compile_while (predicate : Lambda.t) (body : Lambda.t)
       (lambda_cxt : Lam_compile_context.t) =
+    let direct_condition = Lam_analysis.contains_return predicate in
+    let predicate_cxt =
+      if direct_condition then
+        (* The condition will sit inside a new while(true) below. Any break or
+           continue targeting an outer loop must cross that new loop, just as
+           it must cross a switch, so force an explicit outer-loop label. *)
+        Lam_compile_context.enter_switch lambda_cxt
+      else lambda_cxt
+    in
     match
       compile_lambda
-        {lambda_cxt with continuation = NeedValue Not_tail}
+        {predicate_cxt with continuation = NeedValue Not_tail}
         predicate
     with
     | {value = None} -> assert false
     | {block; value = Some e} ->
-      (* st = NeedValue -- this should be optimized and never happen *)
-      let e =
-        match block with
-        | [] -> e
-        | _ -> E.of_block block ~e
-      in
       let loop_cxt, loop_frame = Lam_compile_context.push_loop lambda_cxt in
       let body_block =
         Js_output.output_as_block
@@ -1212,7 +1216,18 @@ let compile output_prefix =
       in
       (* The label stays absent for ordinary loops and is filled in lazily if a
          nested switch emits break/continue for this loop. *)
-      let block = [S.while_ ?label:loop_frame.label e body_block] in
+      let condition, body_block =
+        if direct_condition then
+          (* A helper function would capture the condition's return. Evaluate
+             it at the top of each iteration in the original function. *)
+          (E.true_, block @ [S.if_ (E.not e) [S.break_ ()]] @ body_block)
+        else
+          ( (match block with
+            | [] -> e
+            | _ -> E.of_block block ~e),
+            body_block )
+      in
+      let block = [S.while_ ?label:loop_frame.label condition body_block] in
       Js_output.output_of_block_and_expression lambda_cxt.continuation block
         E.unit
   (* all non-tail
@@ -1899,6 +1914,20 @@ let compile output_prefix =
     | Lswitch (switch_arg, sw) -> compile_switch switch_arg sw lambda_cxt
     | Lstaticraise (i, largs) -> compile_staticraise i largs lambda_cxt
     | Lstaticcatch _ -> compile_staticcatch cur_lam lambda_cxt
+    | Lreturn value ->
+      (* Preserve handlers around operand evaluation. Do not infer a self-tail
+         call here from a continuation that may have been reset by a loop or
+         an expression context. The dummy value satisfies NeedValue callers
+         on this nonreturning path. *)
+      let output =
+        compile_lambda
+          {
+            lambda_cxt with
+            continuation = EffectCall (Maybe_tail_is_return Tail_in_try);
+          }
+          value
+      in
+      {output with value = Some E.undefined; output_finished = True}
     | Lbreak -> (
       match lambda_cxt.loop_stack with
       | [] -> assert false

compiler/core/lam_exit_count.ml

@@ -113,7 +113,7 @@ let count_helper (lam : Lambda.t) : collection option =
       | Lfor_await_of (_, l1, l2) ->
         count l1;
         count l2
-      | Lassign (_, l) -> count l
+      | Lreturn l | Lassign (_, l) -> count l
     and count_default sw =
       match sw.sw_failaction with
       | None -> ()

compiler/core/lam_hit.ml

@@ -36,6 +36,7 @@ let hit_variables (fv : Set_ident.t) (l : t) : bool =
   and hit (l : t) =
     match (l : t) with
     | Lvar id -> hit_var id
+    | Lreturn value -> hit value
     | Lassign (id, e) -> hit_var id || hit e
     | Lstaticcatch (e1, (_, _vars), e2) -> hit e1 || hit e2
     | Ltrywith (e1, _exn, e2) -> hit e1 || hit e2
@@ -73,6 +74,7 @@ let hit_variable (fv : Ident.t) (l : t) : bool =
   and hit (l : t) =
     match (l : t) with
     | Lvar id -> hit_var id
+    | Lreturn value -> hit value
     | Lassign (id, e) -> hit_var id || hit e
     | Lstaticcatch (e1, (_, _vars), e2) -> hit e1 || hit e2
     | Ltrywith (e1, _exn, e2) -> hit e1 || hit e2

compiler/core/lam_pass_collect.ml

@@ -139,6 +139,7 @@ let collect_info (meta : Lam_stats.t) (lam : Lambda.t) =
     | Lfor_await_of (_, l1, l2) ->
       collect l1;
       collect l2
+    | Lreturn l -> collect l
     | Lassign (_v, l) ->
       (* Lalias-bound variables are never assigned, so don't increase
          v's refcollect *)

compiler/core/lam_pass_count.ml

@@ -124,6 +124,7 @@ let collect_occurs lam : occ_tbl =
       (* count [l2] first,
          If v is unused, l1 will be removed, so don't count its variables *)
       if kind = Strict || used v then count bv l1
+    | Lreturn l -> count bv l
     | Lassign (_, l) ->
       (* Lalias-bound variables are never assigned, so don't increase
          this ident's refcount *)

compiler/core/lam_pass_exits.ml

@@ -33,6 +33,7 @@ and no_bounded_variables (l : Lambda.t) =
   match l with
   | Lvar _ -> true
   | Lconst _ -> true
+  | Lreturn e -> no_bounded_variables e
   | Lassign (_id, e) -> no_bounded_variables e
   | Lapply {ap_func; ap_args; _} ->
     no_bounded_variables ap_func && no_list ap_args

compiler/ext/config.ml

@@ -13,6 +13,6 @@ and ast0_impl_magic_number = "Caml1999M022"
 
 and ast0_intf_magic_number = "Caml1999N022"
 
-and cmt_magic_number = "Caml1999T037"
+and cmt_magic_number = "Caml1999T038"
 
 let load_path = ref ([] : string list)

compiler/frontend/bs_builtin_ppx.ml

@@ -129,6 +129,10 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
       ~attrs:(self.attributes self e.pexp_attributes)
       {e with pexp_desc = raw; pexp_attributes = []}
       (Ast_comb.to_regexp_type loc)
+  | Pexp_extension ({txt = "return"}, _) ->
+    (* Unlike unknown extensions, return contains an ordinary expression:
+       lower await, JSX, and other builtins in its operand before typing. *)
+    Ast_mapper.default_mapper.expr self e
   | Pexp_extension extension ->
     Ast_exp_extension.handle_extension e self extension
   | Pexp_constant (Pconst_integer (s, Some 'l')) ->

compiler/ml/lambda.ml

@@ -349,6 +349,10 @@ type t =
   | Ltrywith of t * Ident.t * t
   | Lifthenelse of t * t * t
   | Lsequence of t * t
+  | Lreturn of t
+      (** Evaluates its operand, then exits the nearest [Lfunction]. This is
+          effectful even for a pure operand. Moving it across a function
+          boundary (in particular beta reduction) is invalid. *)
   | Lbreak
   | Lcontinue
   | Lwhile of t * t
@@ -664,6 +668,7 @@ let letrec bindings body : t = Lletrec (bindings, body)
 let staticraise i args : t = Lstaticraise (i, args)
 let staticcatch body catch handler : t = Lstaticcatch (body, catch, handler)
 let try_ body id handler : t = Ltrywith (body, id, handler)
+let return value : t = Lreturn value
 let break : t = Lbreak
 let continue : t = Lcontinue
 let while_ cond body : t = Lwhile (cond, body)
@@ -728,6 +733,10 @@ let rec eq_approx (l1 : t) (l2 : t) =
     match l2 with
     | Lsequence (a0, b0) -> eq_approx a a0 && eq_approx b b0
     | _ -> false)
+  | Lreturn v -> (
+    match l2 with
+    | Lreturn v2 -> eq_approx v v2
+    | _ -> false)
   | Lbreak -> l2 = Lbreak
   | Lcontinue -> l2 = Lcontinue
   | Lwhile (p, b) -> (

compiler/ml/lambda.mli

@@ -363,6 +363,10 @@ type t = private
   | Ltrywith of t * Ident.t * t
   | Lifthenelse of t * t * t
   | Lsequence of t * t
+  | Lreturn of t
+      (** Evaluates its operand, then exits the nearest [Lfunction]. This is
+          effectful even for a pure operand. Moving it across a function
+          boundary (in particular beta reduction) is invalid. *)
   | Lbreak
   | Lcontinue
   | Lwhile of t * t
@@ -506,6 +510,7 @@ val if_ : t -> t -> t -> t
 
 val seq : t -> t -> t
 
+val return : t -> t
 val break : t
 
 val continue : t

compiler/ml/lambda_scc.ml

@@ -37,6 +37,7 @@ let exists_var (p : Ident.t -> bool) (l : Lambda.t) : bool =
   and hit (l : Lambda.t) =
     match l with
     | Lvar id -> p id
+    | Lreturn value -> hit value
     | Lassign (id, e) -> p id || hit e
     | Lstaticcatch (e1, _, e2)
     | Ltrywith (e1, _, e2)

compiler/ml/lambda_traverse.ml

@@ -27,6 +27,9 @@ open Lambda
 let shallow_map_sharing (f : t -> t) (lam : t) : t =
   match lam with
   | Lvar _ | Lglobal_module _ | Lconst _ | Lbreak | Lcontinue -> lam
+  | Lreturn value ->
+    let value' = f value in
+    if value == value' then lam else return value'
   | Lapply ap ->
     let fn = f ap.ap_func in
     let args = Ext_list.map_sharing ap.ap_args f in
@@ -142,6 +145,7 @@ let make_key e =
     | Lifthenelse (cond, ifso, ifnot) ->
       if_ (tr_rec env cond) (tr_rec env ifso) (tr_rec env ifnot)
     | Lsequence (e1, e2) -> seq (tr_rec env e1) (tr_rec env e2)
+    | Lreturn value -> return (tr_rec env value)
     | Lbreak -> break
     | Lcontinue -> continue
     | Lassign (x, e) -> assign x (tr_rec env e)
@@ -168,6 +172,7 @@ let shallow_exists (f : t -> bool) (lam : t) : bool =
   match lam with
   | Lvar _ | Lglobal_module _ | Lconst _ | Lbreak | Lcontinue -> false
   | Lapply {ap_func; ap_args} -> f ap_func || Ext_list.exists ap_args f
+  | Lreturn value -> f value
   | Lfunction {body} -> f body
   | Llet (_, _, arg, body) -> f arg || f body
   | Lletrec (decl, body) -> f body || Ext_list.exists_snd decl f
@@ -217,7 +222,7 @@ let free_ids get l =
     | Lassign (id, _e) -> fv := Set_ident.add !fv id
     | Lvar _ | Lglobal_module _ | Lconst _ | Lapply _ | Lprim _ | Lswitch _
     | Lstringswitch _ | Lstaticraise _ | Lifthenelse _ | Lsequence _ | Lbreak
-    | Lcontinue | Lwhile _ ->
+    | Lreturn _ | Lcontinue | Lwhile _ ->
       ()
   in
   free l;

compiler/ml/printlambda.ml

@@ -355,6 +355,7 @@ let rec lam ppf = function
   | Lifthenelse (lcond, lif, lelse) ->
     fprintf ppf "@[<2>(if@ %a@ %a@ %a)@]" lam lcond lam lif lam lelse
   | Lsequence (l1, l2) -> fprintf ppf "@[<2>(seq@ %a@ %a)@]" lam l1 sequence l2
+  | Lreturn value -> fprintf ppf "(return %a)" lam value
   | Lbreak -> fprintf ppf "break"
   | Lcontinue -> fprintf ppf "continue"
   | Lwhile (lcond, lbody) ->

compiler/ml/printtyped.ml

@@ -342,6 +342,9 @@ and expression i ppf x =
     line i ppf "Texp_sequence\n";
     expression i ppf e1;
     expression i ppf e2
+  | Texp_return exp ->
+    line i ppf "Texp_return\n";
+    expression (i + 1) ppf exp
   | Texp_break -> line i ppf "Texp_break\n"
   | Texp_continue -> line i ppf "Texp_continue\n"
   | Texp_while (e1, e2) ->

compiler/ml/rec_check.ml

@@ -207,9 +207,9 @@ let rec classify_expression : Typedtree.expression -> sd =
     Static
   | Texp_apply {funct = {exp_desc = Texp_ident (_, _, vd)}} when is_ref vd ->
     Static
-  | Texp_apply _ | Texp_match _ | Texp_ifthenelse _ | Texp_object_get _
-  | Texp_object_set _ | Texp_field _ | Texp_assert _ | Texp_try _
-  | Texp_tagged_template _ | Texp_template _ ->
+  | Texp_return _ | Texp_apply _ | Texp_match _ | Texp_ifthenelse _
+  | Texp_object_get _ | Texp_object_set _ | Texp_field _ | Texp_assert _
+  | Texp_try _ | Texp_tagged_template _ | Texp_template _ ->
     Dynamic
 
 let rec expression : Env.env -> Typedtree.expression -> Use.t =
@@ -252,6 +252,7 @@ let rec expression : Env.env -> Typedtree.expression -> Use.t =
         (* The body is evaluated, but not used, and not available
            for inclusion in another value *)
         (discard (expression env body)))
+  | Texp_return exp -> Use.inspect (expression env exp)
   | Texp_constant _ -> Use.empty
   | Texp_break | Texp_continue -> Use.empty
   | Texp_apply