Compiler: propagate arity across unit boundary (WIP)

Propagate shape information through the flow analysis

Function arity from shapes: take advantage of flow analysis

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,7 @@ type t =
   ; static_env : (string * string) list
   ; wrap_with_fun : [ `Iife | `Named of string | `Anonymous ]
   ; target_env : Target_env.t
+  ; shape_files : string list
   ; (* toplevel *)
     dynlink : bool
   ; linkall : bool
@@ -102,6 +103,10 @@ 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" ] ~docv:"FILE" ~doc)
+  in
   let input_file =
     let doc =
       "Compile the bytecode program [$(docv)]. "
@@ -279,6 +284,7 @@ let options =
       output_file
       input_file
       js_files
+      shape_files
       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 +347,7 @@ let options =
       ; bytecode
       ; source_map
       ; keep_unit_names
+      ; shape_files
       }
   in
   let t =
@@ -371,6 +378,7 @@ let options =
       $ output_file
       $ input_file
       $ js_files
+      $ shape_files
       $ keep_unit_names)
   in
   Term.ret t
@@ -567,6 +575,7 @@ let options_runtime_only =
       ; bytecode = `None
       ; source_map
       ; keep_unit_names = false
+      ; shape_files = []
       }
   in
   let t =

compiler/bin-js_of_ocaml/cmd_arg.mli

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

compiler/bin-js_of_ocaml/compile.ml

@@ -52,7 +52,13 @@ let output_gen ~standalone ~custom_header ~build_info ~source_map output_file f
     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
+    (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 +76,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 +135,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 +163,7 @@ let run
     ; export_file
     ; keep_unit_names
     ; include_runtime
+    ; shape_files
     } =
   let source_map_base = Option.map ~f:snd source_map in
   let source_map =
@@ -172,6 +183,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)
@@ -381,7 +393,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
@@ -427,7 +439,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
@@ -460,12 +472,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
@@ -488,7 +501,7 @@ let run
                (`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
@@ -524,16 +537,16 @@ let run
                 (`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,12 +566,17 @@ 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
             ~standalone:false

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)

compiler/lib/driver.ml

@@ -29,6 +29,7 @@ type optimized_result =
   ; trampolined_calls : Effects.trampolined_calls
   ; in_cps : Effects.in_cps
   ; deadcode_sentinal : Code.Var.t
+  ; shapes : Shape.t StringMap.t
   }
 
 type profile =
@@ -95,7 +96,9 @@ let phi p =
 
 let ( +> ) f g x = g (f x)
 
-let map_fst f (x, y, z) = f x, y, z
+let map_fst4 f (x, y, z, t) = f x, y, z, t
+
+let map_fst3 f (x, y, z) = f x, y, z
 
 let effects ~deadcode_sentinal p =
   if Config.Flag.effects ()
@@ -112,7 +115,7 @@ let effects ~deadcode_sentinal p =
         Deadcode.f p
       else p, live_vars
     in
-    p |> Effects.f ~flow_info:info ~live_vars +> map_fst Lambda_lifting.f)
+    p |> Effects.f ~flow_info:info ~live_vars +> map_fst3 Lambda_lifting.f)
   else
     ( p
     , (Code.Var.Set.empty : Effects.trampolined_calls)
@@ -202,7 +205,13 @@ let generate
     ~exported_runtime
     ~wrap_with_fun
     ~warn_on_unhandled_effect
-    { program; variable_uses; trampolined_calls; deadcode_sentinal; in_cps = _ } =
+    { program
+    ; variable_uses
+    ; trampolined_calls
+    ; deadcode_sentinal
+    ; in_cps = _
+    ; shapes = _
+    } =
   if times () then Format.eprintf "Start Generation...@.";
   let should_export = should_export wrap_with_fun in
   Generate.f
@@ -659,6 +668,30 @@ if (typeof module === 'object' && module.exports) {
   if times () then Format.eprintf "  optimizing: %a@." Timer.print t;
   js
 
+let collects_shapes p =
+  let _, info = Flow.f p in
+  let pure = Pure_fun.f p in
+  let l = ref StringMap.empty in
+  Code.Addr.Map.iter
+    (fun _ block ->
+      List.iter block.Code.body ~f:(fun i ->
+          match i with
+          | Code.Let
+              ( _
+              , Prim
+                  ( Extern "caml_register_global"
+                  , [ _code; Pv block; Pc (NativeString name) ] ) ) ->
+              let shape = Flow.the_shape_of ~pure info block in
+              let name =
+                match name with
+                | Byte s -> s
+                | Utf (Utf8 s) -> s
+              in
+              l := StringMap.add name shape !l
+          | _ -> ()))
+    p.blocks;
+  !l
+
 let configure formatter =
   let pretty = Config.Flag.pretty () in
   Pretty_print.set_compact formatter (not pretty);
@@ -689,18 +722,21 @@ let optimize ~profile p =
        | O2 -> o2
        | O3 -> o3)
     +> exact_calls ~deadcode_sentinal profile
-    +> effects ~deadcode_sentinal
-    +> map_fst
+    +> (fun p -> p, collects_shapes p)
+    +> (fun (p, shapes) ->
+         let p, trampolined_calls, cps = effects ~deadcode_sentinal p in
+         p, trampolined_calls, cps, shapes)
+    +> map_fst4
          (match Config.target (), Config.Flag.effects () with
          | `JavaScript, false -> Generate_closure.f
          | `JavaScript, true | `Wasm, _ -> Fun.id)
-    +> map_fst deadcode'
+    +> map_fst4 deadcode'
   in
   if times () then Format.eprintf "Start Optimizing...@.";
   let t = Timer.make () in
-  let (program, variable_uses), trampolined_calls, in_cps = opt p in
+  let (program, variable_uses), trampolined_calls, in_cps, shapes = opt p in
   let () = if times () then Format.eprintf " optimizations : %a@." Timer.print t in
-  { program; variable_uses; trampolined_calls; in_cps; deadcode_sentinal }
+  { program; variable_uses; trampolined_calls; in_cps; deadcode_sentinal; shapes }
 
 let full ~standalone ~wrap_with_fun ~profile ~link ~source_map ~formatter d p =
   let optimized_code = optimize ~profile p in
@@ -710,10 +746,18 @@ let full ~standalone ~wrap_with_fun ~profile ~link ~source_map ~formatter d p =
     +> link_and_pack ~standalone ~wrap_with_fun ~link
     +> output formatter ~source_map ()
   in
-  emit formatter optimized_code
+  let shapes = optimized_code.shapes in
+  StringMap.iter
+    (fun name shape ->
+      Shape.Store.set ~name shape;
+      Pretty_print.string
+        formatter
+        (Printf.sprintf "//# shape: %s:%s\n" name (Shape.to_string shape)))
+    shapes;
+  emit formatter optimized_code, shapes
 
 let full_no_source_map ~formatter ~standalone ~wrap_with_fun ~profile ~link d p =
-  let (_ : Source_map.info) =
+  let (_ : Source_map.info * _) =
     full ~standalone ~wrap_with_fun ~profile ~link ~source_map:false ~formatter d p
   in
   ()

compiler/lib/driver.mli

@@ -18,6 +18,8 @@
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  *)
 
+open Stdlib
+
 type profile
 
 type optimized_result =
@@ -26,6 +28,7 @@ type optimized_result =
   ; trampolined_calls : Effects.trampolined_calls
   ; in_cps : Effects.in_cps
   ; deadcode_sentinal : Code.Var.t
+  ; shapes : Shape.t StringMap.t
   }
 
 val optimize : profile:profile -> Code.program -> optimized_result
@@ -39,7 +42,7 @@ val f :
   -> formatter:Pretty_print.t
   -> Parse_bytecode.Debug.t
   -> Code.program
-  -> Source_map.info
+  -> Source_map.info * Shape.t StringMap.t
 
 val f' :
      ?standalone:bool