Do more aggressive lambda lifting

compiler/bin-js_of_ocaml/build_fs.ml

@@ -80,6 +80,7 @@ function jsoo_create_file_extern(name,content){
           ~standalone:true
           ~wrap_with_fun:`Iife
           ~link:`Needed
+          ~lambda_lift_all:false
           ~formatter:pfs_fmt
           ~source_map:false
           code

compiler/bin-js_of_ocaml/cmd_arg.ml

@@ -78,6 +78,7 @@ type t =
   ; fs_external : bool
   ; keep_unit_names : bool
   ; effects : Config.effects_backend
+  ; lambda_lift_all : bool
   }
 
 let set_param =
@@ -306,6 +307,14 @@ let options =
           None
       & info [ "effects" ] ~docv:"KIND" ~doc)
   in
+  let lambda_lift_all =
+    let doc =
+      "Lambda-lift all functions in the compilation result. This can improve the \
+       performance of some programs on some engines (such as V8). Ignored when effects \
+       are enabled."
+    in
+    Arg.(value & flag & info [ "lambda-lift-all" ] ~doc)
+  in
   let build_t
       common
       set_param
@@ -335,6 +344,7 @@ let options =
       js_files
       keep_unit_names
       effects
+      lambda_lift_all
       shape_files =
     let inline_source_content = not sourcemap_don't_inline_content in
     let chop_extension s = try Filename.chop_extension s with Invalid_argument _ -> s in
@@ -406,6 +416,7 @@ let options =
       ; source_map
       ; keep_unit_names
       ; effects
+      ; lambda_lift_all
       ; shape_files
       }
   in
@@ -440,6 +451,7 @@ let options =
       $ js_files
       $ keep_unit_names
       $ effects
+      $ lambda_lift_all
       $ shape_files)
   in
   Term.ret t
@@ -662,6 +674,7 @@ let options_runtime_only =
       ; keep_unit_names = false
       ; effects
       ; shape_files = []
+      ; lambda_lift_all = false
       }
   in
   let t =

compiler/bin-js_of_ocaml/cmd_arg.mli

@@ -51,6 +51,7 @@ type t =
   ; fs_external : bool
   ; keep_unit_names : bool
   ; effects : Config.effects_backend
+  ; lambda_lift_all : bool
   }
 
 val options : t Cmdliner.Term.t

compiler/bin-js_of_ocaml/compile.ml

@@ -156,6 +156,7 @@ let run
     ; keep_unit_names
     ; include_runtime
     ; effects
+    ; lambda_lift_all
     ; shape_files
     } =
   let source_map_base =
@@ -279,6 +280,7 @@ let run
             ~wrap_with_fun
             ~source_map:(source_map_enabled source_map)
             ~formatter
+            ~lambda_lift_all
             code
       | `File, formatter ->
           let fs_instr1, fs_instr2 =
@@ -300,6 +302,7 @@ let run
               ~shapes
               ?profile
               ~link
+              ~lambda_lift_all
               ~wrap_with_fun
               ~source_map:(source_map_enabled source_map)
               ~formatter

compiler/lib/driver.ml

@@ -142,10 +142,21 @@ let collects_shapes ~shapes (p : Code.program) =
     map)
   else StringMap.empty
 
+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_and_exact_calls
     ~keep_flow_data
     ~deadcode_sentinal
     ~shapes
+    ~lambda_lift_all
     (profile : Profile.t)
     p =
   let fast =
@@ -165,6 +176,14 @@ let effects_and_exact_calls
       Deadcode.f pure_fun p
     else Deadcode.f pure_fun p
   in
+  let p =
+    match lambda_lift_all, Config.target (), Config.effects () with
+    | true, `JavaScript, `Disabled ->
+        let to_lift = all_functions p in
+        let p, _ = Lambda_lifting_simple.f ~to_lift p in
+        p
+    | _ -> p
+  in
   match Config.effects () with
   | `Cps | `Double_translation ->
       if debug () then Format.eprintf "Effects...@.";
@@ -696,7 +715,7 @@ let link_and_pack ?(standalone = true) ?(wrap_with_fun = `Iife) ?(link = `No) p
   |> pack ~wrap_with_fun ~standalone
   |> check_js
 
-let optimize ~shapes ~profile ~keep_flow_data p =
+let optimize ~shapes ~profile ~keep_flow_data ~lambda_lift_all p =
   let deadcode_sentinal =
     (* If deadcode is disabled, this field is just fresh variable *)
     Code.Var.fresh_n "dummy"
@@ -709,7 +728,12 @@ let optimize ~shapes ~profile ~keep_flow_data p =
        | O2 -> o2
        | O3 -> o3)
     +> specialize_js_once_after
-    +> effects_and_exact_calls ~keep_flow_data ~deadcode_sentinal ~shapes profile
+    +> effects_and_exact_calls
+         ~keep_flow_data
+         ~deadcode_sentinal
+         ~shapes
+         ~lambda_lift_all
+         profile
     +> map_fst5
          (match Config.target (), Config.effects () with
          | `JavaScript, `Disabled -> Generate_closure.f
@@ -729,15 +753,26 @@ let optimize ~shapes ~profile ~keep_flow_data p =
 
 let optimize_for_wasm ~shapes ~profile p =
   let optimized_code, global_flow_data =
-    optimize ~shapes ~profile ~keep_flow_data:true p
+    optimize ~shapes ~profile ~keep_flow_data:true ~lambda_lift_all:false p
   in
   ( optimized_code
   , match global_flow_data with
     | Some data -> data
     | None -> Global_flow.f ~fast:false optimized_code.program )
 
-let full ~standalone ~wrap_with_fun ~shapes ~profile ~link ~source_map ~formatter p =
-  let optimized_code, _ = optimize ~shapes ~profile ~keep_flow_data:false p in
+let full
+    ~standalone
+    ~wrap_with_fun
+    ~shapes
+    ~profile
+    ~link
+    ~source_map
+    ~formatter
+    ~lambda_lift_all
+    p =
+  let optimized_code, _ =
+    optimize ~shapes ~profile ~keep_flow_data:false ~lambda_lift_all p
+  in
   let exported_runtime = not standalone in
   let emit formatter =
     generate ~exported_runtime ~wrap_with_fun ~warn_on_unhandled_effect:standalone
@@ -757,9 +792,26 @@ let full ~standalone ~wrap_with_fun ~shapes ~profile ~link ~source_map ~formatte
     shapes_v;
   emit formatter optimized_code, shapes_v
 
-let full_no_source_map ~formatter ~shapes ~standalone ~wrap_with_fun ~profile ~link p =
+let full_no_source_map
+    ~formatter
+    ~shapes
+    ~standalone
+    ~wrap_with_fun
+    ~profile
+    ~link
+    ~lambda_lift_all
+    p =
   let (_ : Source_map.info * _) =
-    full ~shapes ~standalone ~wrap_with_fun ~profile ~link ~source_map:false ~formatter p
+    full
+      ~shapes
+      ~standalone
+      ~wrap_with_fun
+      ~profile
+      ~link
+      ~source_map:false
+      ~formatter
+      ~lambda_lift_all
+      p
   in
   ()
 
@@ -771,17 +823,36 @@ let f
     ~link
     ~source_map
     ~formatter
+    ~lambda_lift_all
     p =
-  full ~standalone ~wrap_with_fun ~shapes ~profile ~link ~source_map ~formatter p
+  full
+    ~standalone
+    ~wrap_with_fun
+    ~shapes
+    ~profile
+    ~link
+    ~source_map
+    ~formatter
+    ~lambda_lift_all
+    p
 
 let f'
     ?(standalone = true)
     ?(wrap_with_fun = `Iife)
     ?(profile = Profile.O1)
+    ?(lambda_lift_all = false)
     ~link
     formatter
     p =
-  full_no_source_map ~formatter ~shapes:false ~standalone ~wrap_with_fun ~profile ~link p
+  full_no_source_map
+    ~formatter
+    ~shapes:false
+    ~standalone
+    ~wrap_with_fun
+    ~profile
+    ~link
+    ~lambda_lift_all
+    p
 
 let from_string ~prims ~debug s formatter =
   let p = Parse_bytecode.from_string ~prims ~debug s in
@@ -792,4 +863,5 @@ let from_string ~prims ~debug s formatter =
     ~wrap_with_fun:`Anonymous
     ~profile:O1
     ~link:`No
+    ~lambda_lift_all:false
     p

compiler/lib/driver.mli

@@ -43,13 +43,15 @@ val f :
   -> link:[ `All | `All_from of string list | `Needed | `No ]
   -> source_map:bool
   -> formatter:Pretty_print.t
+  -> lambda_lift_all:bool
   -> Code.program
   -> Source_map.info * Shape.t StringMap.t
 
 val f' :
      ?standalone:bool
   -> ?wrap_with_fun:[ `Iife | `Anonymous | `Named of string ]
   -> ?profile:Profile.t
+  -> ?lambda_lift_all:bool
   -> link:[ `All | `All_from of string list | `Needed | `No ]
   -> Pretty_print.t
   -> Code.program