Do more aggressive lambda lifting

CHANGES.md

@@ -2,6 +2,7 @@
 
 ## Features/Changes
 * Compiler/wasm: omit code pointer from closures when not used (#2059, #2093)
+* Compiler: add optional full lambda lifting for the Javascript compiler (#1886)
 * Compiler/wasm: unbox numbers within functions (#2069)
 
 ## Bug fixes
@@ -85,7 +86,7 @@
 ## Bug fixes
 * Compiler: fix stack overflow issues with double translation (#1869)
 * Compiler: minifier fix (#1867)
-* Compiler: fix shortvar with --enable es6 (AssignTarget was not properly handled)  
+* Compiler: fix shortvar with --enable es6 (AssignTarget was not properly handled)
 * Compiler: fix assert failure with double translation (#1870)
 * Compiler: fix path rewriting of Wasm source maps (#1882)
 * Compiler: fix global dead code in presence of dead tailcall (#2010)

compiler/lib/config.ml

@@ -106,6 +106,8 @@ module Flag = struct
   let es6 = o ~name:"es6" ~default:false
 
   let load_shapes_auto = o ~name:"load-shapes-auto" ~default:false
+
+  let lambda_lift_all = o ~name:"lambda-lift-all" ~default:false
 end
 
 module Param = struct

compiler/lib/config.mli

@@ -76,6 +76,8 @@ module Flag : sig
 
   val load_shapes_auto : unit -> bool
 
+  val lambda_lift_all : unit -> bool
+
   val enable : string -> unit
 
   val disable : string -> unit

compiler/lib/driver.ml

@@ -142,6 +142,16 @@ 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
@@ -165,6 +175,14 @@ let effects_and_exact_calls
       Deadcode.f pure_fun p
     else Deadcode.f pure_fun p
   in
+  let p =
+    match Config.(Flag.lambda_lift_all (), target (), 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...@.";

compiler/tests-compiler/direct_calls_lift_all.ml

@@ -0,0 +1,215 @@
+(* Js_of_ocaml compiler
+ * http://www.ocsigen.org/js_of_ocaml/
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, with linking exception;
+ * either version 2.1 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
+ *)
+
+open Util
+
+let%expect_test "direct calls with --effects=none --disable lambda-lift-all" =
+  let code =
+    compile_and_parse
+      ~lambda_lift_all:true
+      {|
+         let l = ref []
+
+         (* Arity of the argument of a function / direct call *)
+         let test1 () =
+           let f g x =
+             l := (fun () -> ()) :: !l; (* pervent inlining *)
+             try g x with e -> raise e in
+           ignore (f (fun x -> x + 1) 7);
+           ignore (f (fun x -> x *. 2.) 4.)
+
+         (* Arity of the argument of a function / CPS call *)
+         let test2 () =
+           let f g x =
+             l := (fun () -> ()) :: !l; (* pervent inlining *)
+             try g x with e -> raise e in
+           ignore (f (fun x -> x + 1) 7);
+           ignore (f (fun x -> x ^ "a") "a")
+
+         (* Arity of functions in a functor / direct call *)
+         let test3 x =
+       let module F(_ : sig end) = struct
+         let r = ref 0
+         let () = for _ = 0 to 2 do incr r done (* pervent inlining *)
+         let f x = x + 1
+       end in
+           let module M1 = F (struct end) in
+           let module M2 = F (struct end) in
+           (M1.f 1, M2.f 2)
+
+         (* Arity of functions in a functor / CPS call *)
+         let test4 x =
+           let module F(_ : sig end) =
+             struct
+               let r = ref 0
+               let () = for _ = 0 to 2 do incr r done (* pervent inlining *)
+               let f x = Printf.printf "%d" x
+             end in
+           let module M1 = F (struct end) in
+           let module M2 = F (struct end) in
+           M1.f 1; M2.f 2
+|}
+  in
+  print_fun_decl code (Some "test1");
+  print_fun_decl code (Some "test2");
+  print_fun_decl code (Some "test3");
+  print_fun_decl code (Some "test4");
+  [%expect
+    {|
+    function test1(param){var f = f$2(); f(_f_(), 7); f(_g_(), 4.); return 0;}
+    //end
+    function test2(param){var f = f$1(); f(_c_(), 7); f(_d_(), cst_a); return 0;}
+    //end
+    function test3(x){
+     var F = F$0(), M1 = F([0]), M2 = F([0]), _g_ = M2[2].call(null, 2);
+     return [0, M1[2].call(null, 1), _g_];
+    }
+    //end
+    function test4(x){
+     var F$0 = F(), M1 = F$0([0]), M2 = F$0([0]);
+     M1[2].call(null, 1);
+     return M2[2].call(null, 2);
+    }
+    //end
+    |}]
+
+let%expect_test "direct calls with --effects=cps" =
+  let code =
+    compile_and_parse
+      ~lambda_lift_all:true
+      ~effects:`Cps
+      {|
+         let l = ref []
+
+         (* Arity of the argument of a function / direct call *)
+         let test1 () =
+           let f g x =
+             l := (fun () -> ()) :: !l; (* pervent inlining *)
+             try g x with e -> raise e in
+           ignore (f (fun x -> x + 1) 7);
+           ignore (f (fun x -> x *. 2.) 4.)
+
+         (* Arity of the argument of a function / CPS call *)
+         let test2 () =
+           let f g x =
+             l := (fun () -> ()) :: !l; (* pervent inlining *)
+             try g x with e -> raise e in
+           ignore (f (fun x -> x + 1) 7);
+           ignore (f (fun x -> x ^ "a") "a")
+
+         (* Arity of functions in a functor / direct call *)
+         let test3 x =
+       let module F(_ : sig end) = struct
+         let r = ref 0
+         let () = for _ = 0 to 2 do incr r done (* pervent inlining *)
+         let f x = x + 1
+       end in
+           let module M1 = F (struct end) in
+           let module M2 = F (struct end) in
+           (M1.f 1, M2.f 2)
+
+         (* Arity of functions in a functor / CPS call *)
+         let test4 x =
+           let module F(_ : sig end) =
+             struct
+               let r = ref 0
+               let () = for _ = 0 to 2 do incr r done (* pervent inlining *)
+               let f x = Printf.printf "%d" x
+             end in
+           let module M1 = F (struct end) in
+           let module M2 = F (struct end) in
+           M1.f 1; M2.f 2
+|}
+  in
+  print_fun_decl code (Some "test1");
+  print_fun_decl code (Some "test2");
+  print_fun_decl code (Some "test3");
+  print_fun_decl code (Some "test4");
+  [%expect
+    {|
+    function test1(param, cont){
+     function f(g, x){
+      l[1] = [0, function(param, cont){return cont(0);}, l[1]];
+      try{g(); return;}
+      catch(e$0){
+       var e = caml_wrap_exception(e$0);
+       throw caml_maybe_attach_backtrace(e, 0);
+      }
+     }
+     f(function(x){});
+     f(function(x){});
+     return cont(0);
+    }
+    //end
+    function test2(param, cont){
+     function f(g, x, cont){
+      l[1] = [0, function(param, cont){return cont(0);}, l[1]];
+      runtime.caml_push_trap
+       (function(e){
+         var raise = caml_pop_trap(), e$0 = caml_maybe_attach_backtrace(e, 0);
+         return raise(e$0);
+        });
+      return caml_exact_trampoline_cps_call
+              (g, x, function(_b_){caml_pop_trap(); return cont();});
+     }
+     return caml_exact_trampoline_cps_call$0
+             (f,
+              function(x, cont){return cont();},
+              7,
+              function(_b_){
+               return caml_exact_trampoline_cps_call$0
+                       (f,
+                        function(x, cont){
+                         return caml_trampoline_cps_call3
+                                 (Stdlib[28], x, cst_a$0, cont);
+                        },
+                        cst_a,
+                        function(_b_){return cont(0);});
+              });
+    }
+    //end
+    function test3(x, cont){
+     function F(symbol){
+      var r = [0, 0], for$ = 0;
+      for(;;){r[1]++; var _b_ = for$ + 1 | 0; if(2 === for$) break; for$ = _b_;}
+      function f(x){return x + 1 | 0;}
+      return [0, , f];
+     }
+     var M1 = F(), M2 = F(), _b_ = M2[2].call(null, 2);
+     return cont([0, M1[2].call(null, 1), _b_]);
+    }
+    //end
+    function test4(x, cont){
+     function F(symbol){
+      var r = [0, 0], for$ = 0;
+      for(;;){r[1]++; var _b_ = for$ + 1 | 0; if(2 === for$) break; for$ = _b_;}
+      function f(x, cont){
+       return caml_trampoline_cps_call3(Stdlib_Printf[2], _a_, x, cont);
+      }
+      return [0, , f];
+     }
+     var M1 = F(), M2 = F();
+     return caml_exact_trampoline_cps_call
+             (M1[2],
+              1,
+              function(_a_){
+               return caml_exact_trampoline_cps_call(M2[2], 2, cont);
+              });
+    }
+    //end
+    |}]

compiler/tests-compiler/dune.inc

@@ -89,6 +89,21 @@
  (preprocess
   (pps ppx_expect)))
 
+(library
+ ;; compiler/tests-compiler/direct_calls_lift_all.ml
+ (name direct_calls_lift_all_15)
+ (enabled_if true)
+ (modules direct_calls_lift_all)
+ (libraries js_of_ocaml_compiler unix str jsoo_compiler_expect_tests_helper)
+ (inline_tests
+  (enabled_if true)
+  (deps
+   (file %{project_root}/compiler/bin-js_of_ocaml/js_of_ocaml.exe)
+   (file %{project_root}/compiler/bin-jsoo_minify/jsoo_minify.exe)))
+ (flags (:standard -open Jsoo_compiler_expect_tests_helper))
+ (preprocess
+  (pps ppx_expect)))
+
 (library
  ;; compiler/tests-compiler/effects.ml
  (name effects_15)

compiler/tests-compiler/util/util.ml

@@ -266,6 +266,7 @@ let extract_sourcemap file =
 let compile_to_javascript
     ?(flags = [])
     ?(use_js_string = false)
+    ?(lambda_lift_all = false)
     ?(effects = `Disabled)
     ?(werror = true)
     ~pretty
@@ -283,6 +284,9 @@ let compile_to_javascript
       ; (if use_js_string
          then [ "--enable=use-js-string" ]
          else [ "--disable=use-js-string" ])
+      ; (if lambda_lift_all
+         then [ "--enable=lambda-lift-all" ]
+         else [ "--disable=lambda-lift-all" ])
       ; flags
       ; (if werror then [ "--Werror" ] else [])
       ]
@@ -324,17 +328,26 @@ let compile_bc_to_javascript
     ?flags
     ?effects
     ?use_js_string
+    ?lambda_lift_all
     ?(pretty = true)
     ?(sourcemap = true)
     ?werror
     file =
   Filetype.path_of_bc_file file
-  |> compile_to_javascript ?flags ?effects ?use_js_string ?werror ~pretty ~sourcemap
+  |> compile_to_javascript
+       ?flags
+       ?effects
+       ?use_js_string
+       ?lambda_lift_all
+       ?werror
+       ~pretty
+       ~sourcemap
 
 let compile_cmo_to_javascript
     ?(flags = [])
     ?effects
     ?use_js_string
+    ?lambda_lift_all
     ?(pretty = true)
     ?(sourcemap = true)
     ?werror
@@ -343,6 +356,7 @@ let compile_cmo_to_javascript
   |> compile_to_javascript
        ?effects
        ?use_js_string
+       ?lambda_lift_all
        ?werror
        ~flags:([ "--disable"; "header" ] @ flags)
        ~pretty
@@ -578,6 +592,7 @@ let compile_and_parse_whole_program
     ?flags
     ?effects
     ?use_js_string
+    ?lambda_lift_all
     ?unix
     ?werror
     s =
@@ -591,11 +606,20 @@ let compile_and_parse_whole_program
            ?flags
            ?effects
            ?use_js_string
+           ?lambda_lift_all
            ?werror
            ~sourcemap:debug
       |> parse_js)
 
-let compile_and_parse ?(debug = true) ?pretty ?flags ?effects ?use_js_string ?werror s =
+let compile_and_parse
+    ?(debug = true)
+    ?pretty
+    ?flags
+    ?effects
+    ?use_js_string
+    ?lambda_lift_all
+    ?werror
+    s =
   with_temp_dir ~f:(fun () ->
       s
       |> Filetype.ocaml_text_of_string
@@ -606,6 +630,7 @@ let compile_and_parse ?(debug = true) ?pretty ?flags ?effects ?use_js_string ?we
            ?flags
            ?effects
            ?use_js_string
+           ?lambda_lift_all
            ?werror
            ~sourcemap:debug
       |> parse_js)