1 files changed, 129 insertions, 4 deletions
diff --git a/bytecomp/simplif.ml b/bytecomp/simplif.ml
index 03d503171f..8011953fb8 100644
--- a/bytecomp/simplif.ml
+++ b/bytecomp/simplif.ml
@@ -711,12 +711,137 @@ module Hooks = Misc.MakeHooks(struct
     type t = lambda
   end)
 
+(* Simplify local let-bound functions: if all occurrences are
+   fully-applied function calls in the same "tail scope", replace the
+   function by a staticcatch handler (on that scope).
+
+   This handles as a special case functions used exactly once (in any
+   scope) for a full application.
+*)
+
+type slot =
+  {
+    nargs: int;
+    mutable scope: lambda option;
+  }
+
+module LamTbl = Hashtbl.Make(struct
+    type t = lambda
+    let equal = (==)
+    let hash = Hashtbl.hash
+  end)
+
+let simplify_local_functions lam =
+  let slots = Hashtbl.create 16 in
+  let static_id = Hashtbl.create 16 in (* function id -> static id *)
+  let static = LamTbl.create 16 in (* scope -> static function on that scope *)
+  (* We keep track of the current "tail scope", identified
+     by the outermost lambda for which the the current lambda
+     is in tail position. *)
+  let current_scope = ref lam in
+  let check_static lf =
+    if lf.attr.local = Always_local then
+      Location.prerr_warning lf.loc
+        (Warnings.Inlining_impossible
+           "This function cannot be compiled into a static continuation")
+  in
+  let enabled = function
+    | {local = Always_local; _}
+    | {local = Default_local; inline = (Never_inline | Default_inline); _}
+      -> true
+    | {local = Default_local; inline = (Always_inline | Unroll _); _}
+    | {local = Never_local; _}
+      -> false
+  in
+  let rec tail = function
+    | Llet (_str, _kind, id, Lfunction lf, cont) when enabled lf.attr ->
+        let r = {nargs=List.length lf.params; scope=None} in
+        Hashtbl.add slots id r;
+        tail cont;
+        begin match Hashtbl.find_opt slots id with
+        | Some {scope = Some scope; _} ->
+            let st = next_raise_count () in
+            let sc =
+              (* Do not move higher than current lambda *)
+              if scope == !current_scope then cont
+              else scope
+            in
+            Hashtbl.add static_id id st;
+            LamTbl.add static sc (st, lf);
+            (* The body of the function will become an handler
+               in that "scope". *)
+            with_scope ~scope lf.body
+        | _ ->
+            check_static lf;
+            (* note: if scope = None, the function is unused *)
+            non_tail lf.body
+        end
+    | Lapply {ap_func = Lvar id; ap_args; _} ->
+        begin match Hashtbl.find_opt slots id with
+        | Some {nargs; _} when nargs <> List.length ap_args ->
+            (* Wrong arity *)
+            Hashtbl.remove slots id
+        | Some {scope = Some scope; _} when scope != !current_scope ->
+            (* Different "tail scope" *)
+            Hashtbl.remove slots id
+        | Some ({scope = None; _} as slot) ->
+            (* First use of the function: remember the current tail scope *)
+            slot.scope <- Some !current_scope
+        | _ ->
+            ()
+        end;
+        List.iter non_tail ap_args
+    | Lvar id ->
+        Hashtbl.remove slots id
+    | Lfunction lf as lam ->
+        check_static lf;
+        Lambda.shallow_iter ~tail ~non_tail lam
+    | lam ->
+        Lambda.shallow_iter ~tail ~non_tail lam
+  and non_tail lam =
+    with_scope ~scope:lam lam
+  and with_scope ~scope lam =
+    let old_scope = !current_scope in
+    current_scope := scope;
+    tail lam;
+    current_scope := old_scope
+  in
+  tail lam;
+  let rec rewrite lam0 =
+    let lam =
+      match lam0 with
+      | Llet (_, _, id, _, cont) when Hashtbl.mem static_id id ->
+          rewrite cont
+      | Lapply {ap_func = Lvar id; ap_args; _} when Hashtbl.mem static_id id ->
+          Lstaticraise (Hashtbl.find static_id id, List.map rewrite ap_args)
+      | lam ->
+          Lambda.shallow_map rewrite lam
+    in
+    List.fold_right
+      (fun (st, lf) lam ->
+         Lstaticcatch (lam, (st, lf.params), rewrite lf.body)
+      )
+      (LamTbl.find_all static lam0)
+      lam
+  in
+  if LamTbl.length static = 0 then
+    lam
+  else
+    rewrite lam
+
 (* The entry point:
    simplification + emission of tailcall annotations, if needed. *)
 
 let simplify_lambda sourcefile lam =
-  let res = simplify_lets (simplify_exits lam) in
-  let res = Hooks.apply_hooks { Misc.sourcefile } res in
+  let lam =
+    lam
+    |> (if !Clflags.native_code || not !Clflags.debug
+        then simplify_local_functions else Fun.id
+       )
+    |> simplify_exits
+    |> simplify_lets
+    |> Hooks.apply_hooks { Misc.sourcefile }
+  in
   if !Clflags.annotations || Warnings.is_active Warnings.Expect_tailcall
-    then emit_tail_infos true res;
-  res
+    then emit_tail_infos true lam;
+  lam