delay cmj optimizations

jscomp/core/js_cmj_format.ml

@@ -33,8 +33,6 @@ type cmj_value = {
       (** Either constant or closed functor *)
 }
 
-type effect = string option
-
 let single_na = Single Lam_arity.na
 
 type keyed_cmj_value = {
@@ -47,7 +45,7 @@ type keyed_cmj_values = keyed_cmj_value array
 
 type t = {
   values : keyed_cmj_values;
-  pure : bool;
+  effect : string option;
   package_spec : Js_packages_info.t;
   case : Js_packages_info.file_case;
   delayed_program : J.deps_program;
@@ -63,7 +61,7 @@ let make ~(values : cmj_value String.Map.t) ~effect ~package_spec ~case
             arity = v.arity;
             persistent_closed_lambda = v.persistent_closed_lambda;
           });
-    pure = effect = None;
+    effect;
     package_spec;
     case;
     delayed_program;

jscomp/core/js_cmj_format.mli

@@ -54,8 +54,6 @@ type cmj_value = {
       (* Either constant or closed functor *)
 }
 
-type effect = string option
-
 type keyed_cmj_value = {
   name : string;
   arity : arity;
@@ -64,15 +62,15 @@ type keyed_cmj_value = {
 
 type t = {
   values : keyed_cmj_value array;
-  pure : bool;
+  effect : string option;
   package_spec : Js_packages_info.t;
   case : Js_packages_info.file_case;
   delayed_program : J.deps_program;
 }
 
 val make :
   values:cmj_value String.Map.t ->
-  effect:effect ->
+  effect:string option ->
   package_spec:Js_packages_info.t ->
   case:Js_packages_info.file_case ->
   delayed_program:J.deps_program ->

jscomp/core/js_implementation.ml

@@ -167,16 +167,20 @@ let after_parsing_impl ppf fname (ast : Parsetree.structure) =
       let lambda =
         Translmod.transl_implementation modulename typedtree_coercion
       in
-      let js_program =
-        print_if_pipe ppf Clflags.dump_rawlambda Printlambda.lambda lambda.code
-        |> Lam_compile_main.compile outputprefix
+      let cmj =
+        print_if ppf Clflags.dump_rawlambda Printlambda.lambda lambda.code;
+        Lam_compile_main.compile_coercion ~output_prefix:outputprefix
+          lambda.code
       in
       if not !Js_config.cmj_only then
         (* XXX(anmonteiro): important that we get package_info after
            processing, as `[@@@config {flags = [| ... |]}]` could have added to
            package specs. *)
-        let package_info = Js_packages_state.get_packages_info () in
-        Lam_compile_main.lambda_as_module ~package_info js_program outputprefix
+        let cmj =
+          let package_spec = Js_packages_state.get_packages_info () in
+          { cmj with package_spec }
+        in
+        Lam_compile_main.lambda_as_module ~output_prefix:outputprefix cmj
 (* process_with_gentype (Artifact_extension.append_extension outputprefix Cmt) *)
 
 let implementation ~parser ppf fname =
@@ -200,5 +204,4 @@ let implementation_cmj _ppf fname =
   let output_prefix =
     output_prefix ~f:Filename.chop_all_extensions_maybe fname
   in
-  Lam_compile_main.lambda_as_module ~package_info:cmj.package_spec
-    cmj.delayed_program output_prefix
+  Lam_compile_main.lambda_as_module cmj ~output_prefix

jscomp/core/lam_compile.ml

@@ -190,7 +190,8 @@ type initialization = J.block
 let rec compile_external_field (* Like [List.empty]*)
     (lamba_cxt : Lam_compile_context.t) (id : Ident.t) name : Js_output.t =
   match Lam_compile_env.query_external_id_info id name with
-  | { persistent_closed_lambda = Some lam; _ } when Lam_util.not_function lam ->
+  | { persistent_closed_lambda = Some lam; _ }
+    when not (Lam_util.is_function lam) ->
       compile_lambda lamba_cxt lam
   | _ ->
       Js_output.output_of_expression lamba_cxt.continuation

jscomp/core/lam_compile_env.ml

@@ -130,9 +130,9 @@ let is_pure_module (oid : Lam_module_ident.t) =
           match Js_cmj_format.load_unit (Lam_module_ident.name oid) with
           | cmj_load_info ->
               oid +> Ml cmj_load_info;
-              cmj_load_info.cmj_table.pure
+              cmj_load_info.cmj_table.effect = None
           | exception _ -> false)
-      | Some (Ml { cmj_table; _ }) -> cmj_table.pure
+      | Some (Ml { cmj_table; _ }) -> cmj_table.effect = None
       | Some External -> false)
 
 let populate_required_modules extras

jscomp/core/lam_compile_main.cppo.ml

@@ -109,12 +109,9 @@ let _d  = fun  s lam ->
 
 let _j = Js_pass_debug.dump
 
-(* Actually simplify_lets is kind of global optimization since it requires you to know whether
-    it's used or not
-*)
-let compile
-    (output_prefix : string)
-    (lam : Lambda.lambda)   =
+(* Actually simplify_lets is kind of global optimization since it requires you
+   to know whether it's used or not *)
+let compile_coercion ~output_prefix (lam : Lambda.lambda) =
   let export_idents = Translmod.get_export_identifiers() in
   let export_ident_sets = Ident.Set.of_list export_idents in
   (* To make toplevel happy - reentrant for js-demo *)
@@ -130,7 +127,7 @@ let compile
     Lam_compile_env.reset () ;
   in
   let lam = Tmc.rewrite lam in
-  let lam, may_required_modules = Lam_convert.convert export_ident_sets lam in
+  let lam, maybe_required_modules = Lam_convert.convert export_ident_sets lam in
 
 
   let lam = _d "initial"  lam in
@@ -146,7 +143,7 @@ let compile
       lam
       |> _d "flatten1"
       |>  Lam_pass_exits.simplify_exits
-      |> _d "simplyf_exits"
+      |> _d "simplify_exits"
       |> (fun lam -> Lam_pass_collect.collect_info meta lam;
 #ifndef BS_RELEASE_BUILD
       let () =
@@ -230,7 +227,9 @@ let () =
 
 in
 #endif
-let maybe_pure = no_side_effects groups in
+
+  let maybe_impure = no_side_effects groups in
+
 #ifndef BS_RELEASE_BUILD
 let () =
   Log.warn ~loc:(Loc.of_pos __POS__)
@@ -249,33 +248,26 @@ let () =
     (Pp.textf "[TIME:]Post-compile: %f"  (Sys.time () *. 1000.))
 in
 #endif
-let meta_exports = meta.exports in
-let export_set = Ident.Set.of_list meta_exports in
-let js : J.program =
-  {
-    exports = meta_exports ;
-    export_set;
-    block = body}
-in
-js
-|> _j "initial"
-|> Js_pass_flatten.program
-|> _j "flatten"
-|> Js_pass_tailcall_inline.tailcall_inline
-|> _j "inline_and_shake"
-|> Js_pass_flatten_and_mark_dead.program
-|> _j "flatten_and_mark_dead"
-|> Js_pass_scope.program
-|> Js_shake.shake_program
-|> _j "shake"
-|> ( fun (program:  J.program) ->
+    let program : J.program =
+      let meta_exports = meta.exports in
+      let export_set =
+        Ident.Set.of_list meta_exports
+      in
+      {
+        exports = meta_exports ;
+        export_set;
+        block = body}
+    in
+
+    Warnings.check_fatal();
+
     let external_module_ids : Lam_module_ident.t list =
       if !Js_config.all_module_aliases then []
       else
         let hard_deps =
           Js_fold_basic.calculate_hard_dependencies program.block
         in
-        Lam_compile_env.populate_required_modules may_required_modules hard_deps;
+        Lam_compile_env.populate_required_modules maybe_required_modules hard_deps;
         let module_ids =
           let arr =
             Lam_module_ident.Hash_set.to_list hard_deps
@@ -289,17 +281,19 @@ js
         in
         module_ids
     in
-    Warnings.check_fatal();
     let effect =
       Lam_stats_export.get_dependent_module_effect
-        maybe_pure external_module_ids in
-    let delayed_program = {
-      J.program = program ;
-      side_effect = effect ;
-      preamble = !Js_config.preamble;
-      modules = external_module_ids
-    }
+        maybe_impure external_module_ids in
+
+    let delayed_program: J.deps_program =
+      {
+        J.program;
+        side_effect = effect;
+        preamble = !Js_config.preamble;
+        modules = external_module_ids
+      }
     in
+
     let case =
       Js_packages_info.module_case
         ~output_prefix
@@ -310,33 +304,86 @@ js
         ~case
         ~delayed_program
         meta
-        effect
+        ~effect
         coerced_input.export_map
     in
     (if not !Clflags.dont_write_files then
        Js_cmj_format.to_file (Artifact_extension.append_extension output_prefix Cmj) cmj);
-    delayed_program
-  )
+    cmj
 ;;
 
-let write_to_file ~package_info ~output_info ~output_prefix lambda_output file  =
-  let oc = open_out_bin file in
-  Fun.protect
-    ~finally:(fun () -> close_out oc)
-    (fun () ->
-      Js_dump_program.dump_deps_program
-        ~package_info
-        ~output_info
-        ~output_prefix
-        lambda_output
-        oc)
+let optimize_program (cmj: Js_cmj_format.t) =
+  let deps_program =
+    let deps_program = cmj.delayed_program in
+    deps_program.program
+    |> _j "initial"
+    |> Js_pass_flatten.program
+    |> _j "flatten"
+    |> Js_pass_tailcall_inline.tailcall_inline
+    |> _j "inline_and_shake"
+    |> Js_pass_flatten_and_mark_dead.program
+    |> _j "flatten_and_mark_dead"
+    |> Js_pass_scope.program
+    |> Js_shake.shake_program
+    |> _j "shake"
+    |> (fun (program:  J.program) ->
+        {deps_program with program })
+  in
+  let program = deps_program.program in
+  let external_module_ids : Lam_module_ident.t list =
+    if !Js_config.all_module_aliases then []
+    else
+      let hard_deps =
+        Js_fold_basic.calculate_hard_dependencies program.block
+      in
+      Lam_compile_env.populate_required_modules
+        (deps_program.modules |> Array.of_list |> Lam_module_ident.Hash_set.of_array)
+        hard_deps;
+      let module_ids =
+        let arr =
+          Lam_module_ident.Hash_set.to_list hard_deps
+          |> Array.of_list
+        in
+        Array.sort
+          ~cmp:(fun id1 id2 ->
+            String.compare (Lam_module_ident.name id1) (Lam_module_ident.name id2))
+          arr;
+        Array.to_list arr
+      in
+      module_ids
+  in
+  let effect =
+    Lam_stats_export.get_dependent_module_effect
+      cmj.effect external_module_ids
+  in
+  let deps_program = {
+    deps_program with
+    side_effect = effect ;
+    modules = external_module_ids
+  }
+  in
+  deps_program
 
 let lambda_as_module =
   let (//) = Path.(//) in
-  fun ~package_info (lambda_output : J.deps_program) (output_prefix : string) ->
+  let write_to_file ~package_info ~output_info ~output_prefix lambda_output file  =
+    let oc = open_out_bin file in
+    Fun.protect
+      ~finally:(fun () -> close_out oc)
+      (fun () ->
+        Js_dump_program.dump_deps_program
+          ~package_info
+          ~output_info
+          ~output_prefix
+          lambda_output
+          oc)
+  in
+  fun ~output_prefix (cmj : Js_cmj_format.t) ->
     let make_basename suffix =
       (Filename.basename output_prefix) ^ (Js_suffix.to_string suffix)
     in
+    let package_info = cmj.package_spec in
+    let lambda_output : J.deps_program = optimize_program cmj in
     match (!Js_config.js_stdout, !Clflags.output_name) with
     | (true, None) ->
       Js_dump_program.dump_deps_program

jscomp/core/lam_compile_main.mli

@@ -28,9 +28,9 @@ open Import
 
 (** Compile and register the hook of function to compile  a lambda to JS IR *)
 
-val compile : string -> Lambda.lambda -> J.deps_program
+val compile_coercion : output_prefix:string -> Lambda.lambda -> Js_cmj_format.t
 (** For toplevel, [filename] is [""] which is the same as
     {!Env.get_unit_name ()} *)
 
-val lambda_as_module :
-  package_info:Js_packages_info.t -> J.deps_program -> string -> unit
+val optimize_program : Js_cmj_format.t -> J.deps_program
+val lambda_as_module : output_prefix:string -> Js_cmj_format.t -> unit

jscomp/core/lam_stats_export.ml

@@ -117,14 +117,15 @@ let values_of_export (meta : Lam_stats.t) (export_map : Lam.t Ident.Map.t) :
 *)
 let get_dependent_module_effect (maybe_pure : string option)
     (external_ids : Lam_module_ident.t list) =
-  if maybe_pure = None then
-    let non_pure_module =
-      List.find_opt
-        ~f:(fun id -> not (Lam_compile_env.is_pure_module id))
-        external_ids
-    in
-    Option.map (fun x -> Lam_module_ident.name x) non_pure_module
-  else maybe_pure
+  match maybe_pure with
+  | None ->
+      let non_pure_module =
+        List.find_opt
+          ~f:(fun id -> not (Lam_compile_env.is_pure_module id))
+          external_ids
+      in
+      Option.map (fun x -> Lam_module_ident.name x) non_pure_module
+  | Some _ -> maybe_pure
 
 (* Note that
    [lambda_exports] is
@@ -137,7 +138,7 @@ let get_dependent_module_effect (maybe_pure : string option)
    ]}
    TODO: check that we don't do this in browser environment
 *)
-let export_to_cmj ~case meta effect export_map =
+let export_to_cmj ~case meta ~effect export_map =
   let values = values_of_export meta export_map in
 
   Js_cmj_format.make ~values ~effect

jscomp/core/lam_stats_export.mli

@@ -30,7 +30,7 @@ val get_dependent_module_effect :
 val export_to_cmj :
   case:Js_packages_info.file_case ->
   Lam_stats.t ->
-  Js_cmj_format.effect ->
+  effect:string option ->
   Lam.t Ident.Map.t ->
   (* FIXME: this is a leaky abstraction *)
   delayed_program:J.deps_program ->

jscomp/core/lam_util.ml

@@ -198,8 +198,6 @@ let generate_label ?(name = "") () =
 let is_function (lam : Lam.t) =
   match lam with Lfunction _ -> true | _ -> false
 
-let not_function (lam : Lam.t) =
-  match lam with Lfunction _ -> false | _ -> true
 (*
 let is_var (lam : Lam.t) id =
   match lam with