Propagate arity between compilation units

CHANGES.md

@@ -4,6 +4,7 @@
 * Lib: fix the type of some DOM properties and methods (#1747)
 * Test: use dune test stanzas (#1631)
 * Merged Wasm_of_ocaml (#1724)
+* Compiler: Propagate arity between compilation units (#1594)
 
 # 5.9.1 (02-12-2024) - Lille
 

compiler/bin-js_of_ocaml/build_fs.ml

@@ -74,7 +74,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

@@ -53,6 +53,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
@@ -102,6 +104,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)]. "
@@ -279,6 +289,8 @@ let options =
       output_file
       input_file
       js_files
+      shape_files
+      write_shape
       keep_unit_names =
     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
@@ -341,6 +353,8 @@ let options =
       ; bytecode
       ; source_map
       ; keep_unit_names
+      ; shape_files
+      ; write_shape
       }
   in
   let t =
@@ -371,6 +385,8 @@ let options =
       $ output_file
       $ input_file
       $ js_files
+      $ shape_files
+      $ write_shape
       $ keep_unit_names)
   in
   Term.ret t
@@ -567,6 +583,8 @@ let options_runtime_only =
       ; bytecode = `None
       ; source_map
       ; keep_unit_names = false
+      ; 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

@@ -46,13 +46,28 @@ let source_map_enabled = function
   | No_sourcemap -> false
   | Inline | File _ -> true
 
-let output_gen ~standalone ~custom_header ~build_info ~source_map output_file f =
+let output_gen
+    ~write_shape
+    ~standalone
+    ~custom_header
+    ~build_info
+    ~source_map
+    output_file
+    f =
   let f chan k =
     let fmt = Pretty_print.to_out_channel chan in
     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 ~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
     | No_sourcemap, _ | _, None -> ()
     | ((Inline | File _) as output), Some sm ->
@@ -70,7 +85,6 @@ let output_gen ~standalone ~custom_header ~build_info ~source_map output_file f
         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)
@@ -130,6 +144,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
@@ -153,6 +172,8 @@ let run
     ; export_file
     ; keep_unit_names
     ; include_runtime
+    ; shape_files
+    ; write_shape
     } =
   let source_map_base = Option.map ~f:snd source_map in
   let source_map =
@@ -172,6 +193,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)
@@ -366,6 +388,7 @@ let run
         }
       in
       output_gen
+        ~write_shape
         ~standalone:true
         ~custom_header
         ~build_info:(Build_info.create `Runtime)
@@ -381,7 +404,7 @@ let run
             ~standalone
             ~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
@@ -414,6 +437,7 @@ 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)
@@ -427,7 +451,7 @@ let run
                 ~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
@@ -452,6 +476,7 @@ 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)
@@ -460,12 +485,13 @@ let run
             (fun ~standalone ~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 ~source_map output in
+                  let sm2, sh2 = output_partial cmo code ~standalone ~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)
+                  |> map_fst (sourcemap_of_info ~base:source_map_base))
       | `Cma cma when keep_unit_names ->
           (if include_runtime
            then
@@ -481,14 +507,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));
+                 |> map_fst (sourcemap_of_info ~base:source_map_base)));
           List.iter cma.lib_units ~f:(fun cmo ->
               let output_file =
                 match output_file with
@@ -517,23 +544,24 @@ 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))
+                  |> map_fst (sourcemap_of_info ~base:source_map_base)))
       | `Cma cma ->
           let f ~standalone ~source_map output =
-            let source_map_runtime =
+            let runtime =
               if not include_runtime
               then None
               else Some (output_partial_runtime ~standalone ~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 =
@@ -553,14 +581,20 @@ let run
                       (Ocaml_compiler.Cmo_format.name cmo);
                   output_partial ~standalone ~source_map cmo code output)
             in
-            let sm =
-              match source_map_runtime with
-              | None -> source_map_units
-              | Some x -> x :: source_map_units
+            let sm_and_shapes =
+              match runtime with
+              | None -> units
+              | Some x -> x :: units
+            in
+            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/lib/code.ml

@@ -556,13 +556,17 @@ module Print = struct
         if exact
         then Format.fprintf f "%a!(%a)" Var.print g var_list args
         else Format.fprintf f "%a(%a)" Var.print g var_list args
-    | Block (t, a, _, mut) ->
+    | Block (t, a, k, mut) ->
         Format.fprintf
           f
-          "%s{tag=%d"
+          "{%s%s:tag=%d"
           (match mut with
           | Immutable -> "imm"
           | Maybe_mutable -> "")
+          (match k with
+          | Array -> "A"
+          | NotArray -> "NA"
+          | Unknown -> "U")
           t;
         for i = 0 to Array.length a - 1 do
           Format.fprintf f "; %d = %a" i Var.print a.(i)
@@ -818,6 +822,30 @@ let fold_closures_outermost_first { start; blocks; _ } f accu =
   let accu = f None [] (start, []) accu in
   visit blocks start f accu
 
+(* Compute the list of variables containing the return values of each
+   function *)
+let return_values p =
+  fold_closures
+    p
+    (fun name_opt _ (pc, _) rets ->
+      match name_opt with
+      | None -> rets
+      | Some name ->
+          let s =
+            traverse
+              { fold = fold_children }
+              (fun pc s ->
+                let block = Addr.Map.find pc p.blocks in
+                match block.branch with
+                | Return x -> Var.Set.add x s
+                | _ -> s)
+              pc
+              p.blocks
+              Var.Set.empty
+          in
+          Var.Map.add name s rets)
+    Var.Map.empty
+
 let eq p1 p2 =
   p1.start = p2.start
   && Addr.Map.cardinal p1.blocks = Addr.Map.cardinal p2.blocks

compiler/lib/code.mli

@@ -304,6 +304,8 @@ val fold_children_skip_try_body : 'c fold_blocs
 
 val poptraps : block Addr.Map.t -> Addr.t -> Addr.Set.t
 
+val return_values : program -> Var.Set.t Var.Map.t
+
 val traverse :
   fold_blocs_poly -> (Addr.t -> 'c -> 'c) -> Addr.t -> block Addr.Map.t -> 'c -> 'c
 

compiler/lib/config.ml

@@ -101,6 +101,8 @@ module Flag = struct
   let auto_link = o ~name:"auto-link" ~default:true
 
   let es6 = o ~name:"es6" ~default:false
+
+  let load_shapes_auto = o ~name:"load-shapes-auto" ~default:false
 end
 
 module Param = struct

compiler/lib/config.mli

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