Propagate arity between compilation units

CHANGES.md

@@ -39,6 +39,7 @@
 * Compiler: deadcode elimination of cyclic values (#1978)
 * Compiler: directly write Wasm binary modules (#2000, #2003)
 * Compiler: rewrote inlining pass (#1935)
+* Compiler: Propagate arity between compilation units (#1594)
 
 ## Bug fixes
 * Compiler: fix stack overflow issues with double translation (#1869)

compiler/bin-js_of_ocaml/build_fs.ml

@@ -75,7 +75,7 @@ function jsoo_create_file_extern(name,content){
   let code = Code.prepend Code.empty instr in
   Filename.gen_file output_file (fun chan ->
       let pfs_fmt = Pretty_print.to_out_channel chan in
-      let (_ : Source_map.info) =
+      let (_ : Source_map.info * Shape.t StringMap.t) =
         Driver.f
           ~standalone:true
           ~wrap_with_fun:`Iife

compiler/bin-js_of_ocaml/cmd_arg.ml

@@ -64,6 +64,8 @@ type t =
   ; static_env : (string * string) list
   ; wrap_with_fun : [ `Iife | `Named of string | `Anonymous ]
   ; target_env : Target_env.t
+  ; shape_files : string list
+  ; write_shape : bool
   ; (* toplevel *)
     dynlink : bool
   ; linkall : bool
@@ -114,6 +116,14 @@ let options =
     let doc = "Set output file name to [$(docv)]." in
     Arg.(value & opt (some string) None & info [ "o" ] ~docv:"FILE" ~doc)
   in
+  let shape_files =
+    let doc = "load shape file [$(docv)]." in
+    Arg.(value & opt_all string [] & info [ "load-shape" ] ~docv:"FILE" ~doc)
+  in
+  let write_shape =
+    let doc = "Emit shape files" in
+    Arg.(value & flag & info [ "write-shape" ] ~doc)
+  in
   let input_file =
     let doc =
       "Compile the bytecode program [$(docv)]. "
@@ -309,7 +319,9 @@ let options =
       input_file
       js_files
       keep_unit_names
-      effects =
+      effects
+      shape_files
+      write_shape =
     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
     let runtime_files = js_files in
@@ -380,6 +392,8 @@ let options =
       ; source_map
       ; keep_unit_names
       ; effects
+      ; shape_files
+      ; write_shape
       }
   in
   let t =
@@ -412,7 +426,9 @@ let options =
       $ input_file
       $ js_files
       $ keep_unit_names
-      $ effects)
+      $ effects
+      $ shape_files
+      $ write_shape)
   in
   Term.ret t
 
@@ -633,6 +649,8 @@ let options_runtime_only =
       ; source_map
       ; keep_unit_names = false
       ; effects
+      ; shape_files = []
+      ; write_shape = false
       }
   in
   let t =

compiler/bin-js_of_ocaml/cmd_arg.mli

@@ -37,6 +37,8 @@ type t =
       | `Anonymous
       ]
   ; target_env : Target_env.t
+  ; shape_files : string list
+  ; write_shape : bool
   ; (* toplevel *)
     dynlink : bool
   ; linkall : bool

compiler/bin-js_of_ocaml/compile.ml

@@ -42,6 +42,7 @@ let source_map_enabled : Source_map.Encoding_spec.t option -> bool = function
   | Some _ -> true
 
 let output_gen
+    ~write_shape
     ~standalone
     ~custom_header
     ~build_info
@@ -53,7 +54,15 @@ let output_gen
     Driver.configure fmt;
     if standalone then header ~custom_header fmt;
     if Config.Flag.header () then jsoo_header fmt build_info;
-    let sm = f ~standalone ~source_map (k, fmt) in
+    let sm, shapes = f ~standalone ~shapes:write_shape ~source_map (k, fmt) in
+    (if write_shape
+     then
+       match output_file with
+       | `Stdout -> ()
+       | `Name name ->
+           Shape.Store.save'
+             (Filename.remove_extension name ^ Shape.Store.ext)
+             (StringMap.bindings shapes));
     match source_map, sm with
     | None, _ | _, None -> ()
     | Some { output_file = output; source_map; keep_empty }, Some sm ->
@@ -71,7 +80,6 @@ let output_gen
         Pretty_print.newline fmt;
         Pretty_print.string fmt (Printf.sprintf "//# sourceMappingURL=%s\n" urlData)
   in
-
   match output_file with
   | `Stdout -> f stdout `Stdout
   | `Name name -> Filename.gen_file name (fun chan -> f chan `File)
@@ -132,6 +140,11 @@ let sourcemap_of_infos ~base l =
 
 let sourcemap_of_info ~base info = sourcemap_of_infos ~base [ info ]
 
+let map_fst f (x, y) = f x, y
+
+let merge_shape a b =
+  StringMap.union (fun _name s1 s2 -> if Shape.equal s1 s2 then Some s1 else None) a b
+
 let run
     { Cmd_arg.common
     ; profile
@@ -156,6 +169,8 @@ let run
     ; keep_unit_names
     ; include_runtime
     ; effects
+    ; shape_files
+    ; write_shape
     } =
   let source_map_base =
     Option.map ~f:(fun spec -> spec.Source_map.Encoding_spec.source_map) source_map
@@ -172,6 +187,7 @@ let run
   | `Name _, _ -> ());
   List.iter params ~f:(fun (s, v) -> Config.Param.set s v);
   List.iter static_env ~f:(fun (s, v) -> Eval.set_static_env s v);
+  List.iter shape_files ~f:(fun fn -> Shape.Store.load' fn);
   let t = Timer.make () in
   let include_dirs =
     List.filter_map (include_dirs @ [ "+stdlib/" ]) ~f:(fun d -> Findlib.find [] d)
@@ -251,6 +267,7 @@ let run
       (one : Parse_bytecode.one)
       ~check_sourcemap
       ~standalone
+      ~shapes
       ~(source_map : Source_map.Encoding_spec.t option)
       ~link
       output_file =
@@ -269,6 +286,7 @@ let run
           let code = Code.prepend one.code instr in
           Driver.f
             ~standalone
+            ~shapes
             ?profile
             ~link
             ~wrap_with_fun
@@ -292,6 +310,7 @@ let run
           let res =
             Driver.f
               ~standalone
+              ~shapes
               ?profile
               ~link
               ~wrap_with_fun
@@ -313,14 +332,22 @@ let run
   let output_partial
       (cmo : Cmo_format.compilation_unit)
       ~standalone
+      ~shapes
       ~source_map
       code
       ((_, fmt) as output_file) =
     assert (not standalone);
     let uinfo = Unit_info.of_cmo cmo in
     Pretty_print.string fmt "\n";
     Pretty_print.string fmt (Unit_info.to_string uinfo);
-    output code ~check_sourcemap:true ~source_map ~standalone ~link:`No output_file
+    output
+      code
+      ~check_sourcemap:true
+      ~source_map
+      ~standalone
+      ~shapes
+      ~link:`No
+      output_file
   in
   let output_partial_runtime ~standalone ~source_map ((_, fmt) as output_file) =
     assert (not standalone);
@@ -371,22 +398,24 @@ let run
         { code; cmis = StringSet.empty; debug = Parse_bytecode.Debug.default_summary }
       in
       output_gen
+        ~write_shape
         ~standalone:true
         ~custom_header
         ~build_info:(Build_info.create `Runtime)
         ~source_map
         (fst output_file)
-        (fun ~standalone ~source_map ((_, fmt) as output_file) ->
+        (fun ~standalone ~shapes ~source_map ((_, fmt) as output_file) ->
           Pretty_print.string fmt "\n";
           Pretty_print.string fmt (Unit_info.to_string uinfo);
           output
             code
             ~check_sourcemap:false
             ~source_map
             ~standalone
+            ~shapes
             ~link:`All
             output_file
-          |> sourcemap_of_info ~base:source_map_base)
+          |> map_fst (sourcemap_of_info ~base:source_map_base))
   | (`Stdin | `File _) as bytecode ->
       let kind, ic, close_ic, include_dirs =
         match bytecode with
@@ -419,20 +448,22 @@ let run
           in
           if times () then Format.eprintf "  parsing: %a@." Timer.print t1;
           output_gen
+            ~write_shape
             ~standalone:true
             ~custom_header
             ~build_info:(Build_info.create `Exe)
             ~source_map
             (fst output_file)
-            (fun ~standalone ~source_map output_file ->
+            (fun ~standalone ~shapes ~source_map output_file ->
               output
                 code
                 ~check_sourcemap:true
                 ~standalone
+                ~shapes
                 ~source_map
                 ~link:(if linkall then `All else `Needed)
                 output_file
-              |> sourcemap_of_info ~base:source_map_base)
+              |> map_fst (sourcemap_of_info ~base:source_map_base))
       | `Cmo cmo ->
           let output_file =
             match output_file, keep_unit_names with
@@ -457,20 +488,26 @@ let run
           in
           if times () then Format.eprintf "  parsing: %a@." Timer.print t1;
           output_gen
+            ~write_shape
             ~standalone:false
             ~custom_header
             ~build_info:(Build_info.create `Cmo)
             ~source_map
             output_file
-            (fun ~standalone ~source_map output ->
+            (fun ~standalone ~shapes ~source_map output ->
               match include_runtime with
               | true ->
-                  let sm1 = output_partial_runtime ~standalone ~source_map output in
-                  let sm2 = output_partial cmo code ~standalone ~source_map output in
-                  sourcemap_of_infos ~base:source_map_base [ sm1; sm2 ]
+                  let sm1, sh1 =
+                    output_partial_runtime ~standalone ~shapes ~source_map output
+                  in
+                  let sm2, sh2 =
+                    output_partial cmo code ~standalone ~shapes ~source_map output
+                  in
+                  ( sourcemap_of_infos ~base:source_map_base [ sm1; sm2 ]
+                  , merge_shape sh1 sh2 )
               | false ->
-                  output_partial cmo code ~standalone ~source_map output
-                  |> sourcemap_of_info ~base:source_map_base)
+                  output_partial cmo code ~standalone ~shapes ~source_map output
+                  |> map_fst (sourcemap_of_info ~base:source_map_base))
       | `Cma cma when keep_unit_names ->
           (if include_runtime
            then
@@ -486,14 +523,15 @@ let run
                    failwith "use [-o dirname/] or remove [--keep-unit-names]"
              in
              output_gen
+               ~write_shape
                ~standalone:false
                ~custom_header
                ~build_info:(Build_info.create `Runtime)
                ~source_map
                (`Name output_file)
-               (fun ~standalone ~source_map output ->
-                 output_partial_runtime ~standalone ~source_map output
-                 |> sourcemap_of_info ~base:source_map_base));
+               (fun ~standalone ~shapes ~source_map output ->
+                 output_partial_runtime ~standalone ~shapes ~source_map output
+                 |> map_fst (sourcemap_of_info ~base:source_map_base)));
           List.iter cma.lib_units ~f:(fun cmo ->
               let output_file =
                 match output_file with
@@ -522,23 +560,26 @@ let run
                   t1
                   (Ocaml_compiler.Cmo_format.name cmo);
               output_gen
+                ~write_shape
                 ~standalone:false
                 ~custom_header
                 ~build_info:(Build_info.create `Cma)
                 ~source_map
                 (`Name output_file)
-                (fun ~standalone ~source_map output ->
-                  output_partial ~standalone ~source_map cmo code output
-                  |> sourcemap_of_info ~base:source_map_base))
+                (fun ~standalone ~shapes ~source_map output ->
+                  output_partial ~standalone ~shapes ~source_map cmo code output
+                  |> map_fst (sourcemap_of_info ~base:source_map_base)))
       | `Cma cma ->
-          let f ~standalone ~source_map output =
-            let source_map_runtime =
+          let f ~standalone ~shapes ~source_map output =
+            (* Always compute shapes because it can be used by other units of the cma *)
+            let shapes = shapes || true in
+            let runtime =
               if not include_runtime
               then None
-              else Some (output_partial_runtime ~standalone ~source_map output)
+              else Some (output_partial_runtime ~standalone ~shapes ~source_map output)
             in
 
-            let source_map_units =
+            let units =
               List.map cma.lib_units ~f:(fun cmo ->
                   let t1 = Timer.make () in
                   let code =
@@ -556,16 +597,22 @@ let run
                       Timer.print
                       t1
                       (Ocaml_compiler.Cmo_format.name cmo);
-                  output_partial ~standalone ~source_map cmo code output)
+                  output_partial ~standalone ~shapes ~source_map cmo code output)
+            in
+            let sm_and_shapes =
+              match runtime with
+              | None -> units
+              | Some x -> x :: units
             in
-            let sm =
-              match source_map_runtime with
-              | None -> source_map_units
-              | Some x -> x :: source_map_units
+            let shapes =
+              List.fold_left sm_and_shapes ~init:StringMap.empty ~f:(fun acc (_, s) ->
+                  merge_shape s acc)
             in
-            sourcemap_of_infos ~base:source_map_base sm
+            ( sourcemap_of_infos ~base:source_map_base (List.map sm_and_shapes ~f:fst)
+            , shapes )
           in
           output_gen
+            ~write_shape
             ~standalone:false
             ~custom_header
             ~build_info:(Build_info.create `Cma)

compiler/bin-wasm_of_ocaml/compile.ml

@@ -245,7 +245,7 @@ let generate_prelude ~out_file =
   let code, uinfo = Parse_bytecode.predefined_exceptions () in
   let profile = Profile.O1 in
   let Driver.{ program; variable_uses; in_cps; deadcode_sentinal; _ } =
-    Driver.optimize ~profile code
+    Driver.optimize ~profile ~shapes:false code
   in
   let context = Generate.start () in
   let _ =
@@ -397,7 +397,7 @@ let run
     let code = one.code in
     let standalone = Option.is_none unit_name in
     let Driver.{ program; variable_uses; in_cps; deadcode_sentinal; _ } =
-      Driver.optimize ~profile code
+      Driver.optimize ~profile ~shapes:false code
     in
     let context = Generate.start () in
     let toplevel_name, generated_js =