From 49cde7ca2b81393296e263931c6ea3becdacd055 Mon Sep 17 00:00:00 2001
From: Olivier Nicole <olivier@chnik.fr>
Date: Fri, 14 Mar 2025 17:10:38 +0000
Subject: [PATCH] Do more aggressive lambda lifting

With @vouillon we realized that the `Lambda_lifting_simple` pass that is
performed by double translation makes some programs significantly
faster. We measured roughly a 1.45 speedup on a large (proprietary)
Bonsai benchmark. Presumably V8 is much faster with more toplevel
functions and less nested closures.
---
 compiler/lib/driver.ml | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/compiler/lib/driver.ml b/compiler/lib/driver.ml
index 484489069e..f27580dd33 100644
--- a/compiler/lib/driver.ml
+++ b/compiler/lib/driver.ml
@@ -673,6 +673,33 @@ let link_and_pack ?(standalone = true) ?(wrap_with_fun = `Iife) ?(link = `No) p
   |> pack ~wrap_with_fun ~standalone
   |> check_js
 
+let all_functions p =
+  let open Code in
+  fold_closures
+    p
+    (fun name _ _ _ acc ->
+      match name with
+      | Some name -> Var.Set.add name acc
+      | None -> acc)
+    Var.Set.empty
+
+let effects_or_lambda_lift ~deadcode_sentinal p =
+  (* If effects are disabled, we lambda-lift aggressively. While not necessary, it results
+     in a substantial gain in performance for Javascript. *)
+  match Config.(target (), effects ()) with
+  | `JavaScript, `Disabled ->
+      let to_lift = all_functions p in
+      let p, _ = Lambda_lifting_simple.f ~to_lift p in
+      ( p
+      , (Code.Var.Set.empty : Effects.trampolined_calls)
+      , (Code.Var.Set.empty : Effects.in_cps) )
+  | _, (`Cps | `Double_translation) -> effects ~deadcode_sentinal p
+  | `Wasm, (`Disabled | `Jspi) ->
+      ( p
+      , (Code.Var.Set.empty : Effects.trampolined_calls)
+      , (Code.Var.Set.empty : Effects.in_cps) )
+  | `JavaScript, `Jspi -> assert false
+
 let optimize ~profile p =
   let deadcode_sentinal =
     (* If deadcode is disabled, this field is just fresh variable *)
@@ -687,7 +714,7 @@ let optimize ~profile p =
        | O3 -> o3)
     +> specialize_js_once_after
     +> exact_calls ~deadcode_sentinal profile
-    +> effects ~deadcode_sentinal
+    +> effects_or_lambda_lift ~deadcode_sentinal
     +> map_fst
          (match Config.target (), Config.effects () with
          | `JavaScript, `Disabled -> Generate_closure.f