loading shapes

Author: hhugo

compiler/bin-js_of_ocaml/build_fs.ml

@@ -73,7 +73,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.t option) =
+      let (_ : Source_map.t option * 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 chop_extension s = try Filename.chop_extension s with Invalid_argument _ -> s in
     let runtime_files = js_files in
@@ -344,6 +350,7 @@ let options =
       ; bytecode
       ; source_map
       ; keep_unit_names
+      ; shape_files
       }
   in
   let t =
@@ -374,6 +381,7 @@ let options =
       $ output_file
       $ input_file
       $ js_files
+      $ shape_files
       $ keep_unit_names)
   in
   Term.ret t
@@ -573,6 +581,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

@@ -43,7 +43,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:(Option.map ~f:snd source_map) (k, fmt) in
+    let sm, shapes = f ~standalone ~source_map:(Option.map ~f:snd 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
     | None, _ | _, None -> ()
     | Some (output_file, _), Some sm ->
@@ -59,7 +65,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)
@@ -87,6 +92,7 @@ let run
     ; export_file
     ; keep_unit_names
     ; include_runtime
+    ; shape_files
     } =
   let include_cmis = toplevel && not no_cmis in
   let custom_header = common.Jsoo_cmdline.Arg.custom_header in
@@ -97,6 +103,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)
@@ -379,7 +386,7 @@ let run
               let source_map =
                 if not include_runtime
                 then source_map
-                else output_partial_runtime ~standalone ~source_map output
+                else fst (output_partial_runtime ~standalone ~source_map output)
               in
               output_partial cmo code ~standalone ~source_map output)
       | `Cma cma when keep_unit_names ->
@@ -443,9 +450,12 @@ let run
             let source_map =
               if not include_runtime
               then source_map
-              else output_partial_runtime ~standalone ~source_map output
+              else fst (output_partial_runtime ~standalone ~source_map output)
             in
-            List.fold_left cma.lib_units ~init:source_map ~f:(fun source_map cmo ->
+            List.fold_left
+              cma.lib_units
+              ~init:(source_map, StringMap.empty)
+              ~f:(fun (source_map, shapes) cmo ->
                 let t1 = Timer.make () in
                 let code =
                   Parse_bytecode.from_cmo
@@ -462,7 +472,14 @@ let run
                     Timer.print
                     t1
                     (Ocaml_compiler.Cmo_format.name cmo);
-                output_partial cmo ~standalone ~source_map code output)
+                let sm, shapes' =
+                  output_partial cmo ~standalone ~source_map code output
+                in
+                ( sm
+                , StringMap.union
+                    (fun _name s1 s2 -> if Shape.equal s1 s2 then Some s1 else None)
+                    shapes
+                    shapes' ))
           in
           output_gen
             ~standalone:false

compiler/lib/driver.ml

@@ -728,10 +728,10 @@ let full ~standalone ~wrap_with_fun ~profile ~link ~source_map formatter d p =
         (Printf.sprintf "//# shape: %s:%s\n" name (Shape.to_string shape)))
     shapes;
   let () = if times () then Format.eprintf " optimizations : %a@." Timer.print t in
-  emit ~live_vars ~cps_calls prog
+  emit ~live_vars ~cps_calls prog, shapes
 
 let full_no_source_map ~standalone ~wrap_with_fun ~profile ~link formatter d p =
-  let (_ : Source_map.t option) =
+  let (_ : Source_map.t option * Shape.t StringMap.t) =
     full ~standalone ~wrap_with_fun ~profile ~link ~source_map:None 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
 
 val f :
@@ -29,7 +31,7 @@ val f :
   -> Pretty_print.t
   -> Parse_bytecode.Debug.t
   -> Code.program
-  -> Source_map.t option
+  -> Source_map.t option * Shape.t StringMap.t
 
 val f' :
      ?standalone:bool

compiler/lib/parse_bytecode.ml

@@ -590,6 +590,7 @@ module State = struct
     ; handlers : handler list
     ; globals : globals
     ; immutable : Code.Var.Set.t ref
+    ; includes : string list
     }
 
   let fresh_var state loc =
@@ -695,14 +696,15 @@ module State = struct
 
   let pop_handler state = { state with handlers = List.tl state.handlers }
 
-  let initial g immutable =
+  let initial includes g immutable =
     { accu = Unset
     ; stack = []
     ; env = [||]
     ; env_offset = 0
     ; handlers = []
     ; globals = g
     ; immutable
+    ; includes
     }
 
   let rec print_stack f l =
@@ -818,7 +820,7 @@ let get_global state instrs i loc =
         (match g.named_value.(i) with
         | None -> ()
         | Some name -> (
-            match Shape.Store.get ~name with
+            match Shape.Store.load ~name state.includes with
             | None -> ()
             | Some shape -> Shape.State.assign x shape));
 
@@ -2534,9 +2536,9 @@ type one =
   ; debug : Debug.t
   }
 
-let parse_bytecode code globals debug_data =
+let parse_bytecode ~includes code globals debug_data =
   let immutable = ref Code.Var.Set.empty in
-  let state = State.initial globals immutable in
+  let state = State.initial includes globals immutable in
   Code.Var.reset ();
   Shape.State.reset ();
   let blocks = Blocks.analyse debug_data code in
@@ -2748,7 +2750,7 @@ let from_exe
     Ocaml_compiler.Symtable.GlobalMap.iter symbols ~f:(fun id n ->
         globals.named_value.(n) <- Some (Ocaml_compiler.Symtable.Global.name id);
         globals.is_exported.(n) <- true);
-  let p = parse_bytecode code globals debug_data in
+  let p = parse_bytecode ~includes code globals debug_data in
   (* register predefined exception *)
   let body =
     List.fold_left predefined_exceptions ~init:[] ~f:(fun body (i, name) ->
@@ -2872,7 +2874,7 @@ let from_bytes ~prims ~debug (code : bytecode) =
     t
   in
   let globals = make_globals 0 [||] prims in
-  let p = parse_bytecode code globals debug_data in
+  let p = parse_bytecode ~includes:[] code globals debug_data in
   let gdata = Var.fresh_n "global_data" in
   let need_gdata = ref false in
   let find_name i =
@@ -3011,7 +3013,7 @@ module Reloc = struct
     globals
 end
 
-let from_compilation_units ~includes:_ ~include_cmis ~debug_data l =
+let from_compilation_units ~includes ~include_cmis ~debug_data l =
   let reloc = Reloc.create () in
   List.iter l ~f:(fun (compunit, code) -> Reloc.step1 reloc compunit code);
   List.iter l ~f:(fun (compunit, code) -> Reloc.step2 reloc compunit code);
@@ -3020,7 +3022,7 @@ let from_compilation_units ~includes:_ ~include_cmis ~debug_data l =
     let l = List.map l ~f:(fun (_, c) -> Bytes.to_string c) in
     String.concat ~sep:"" l
   in
-  let prog = parse_bytecode code globals debug_data in
+  let prog = parse_bytecode ~includes code globals debug_data in
   let gdata = Var.fresh_n "global_data" in
   let need_gdata = ref false in
   let body =

compiler/lib/shape.ml

@@ -28,6 +28,20 @@ type t =
       ; res : t
       }
 
+type shape = t
+
+let rec equal a b =
+  match a, b with
+  | Top _, Top _ -> true
+  | ( Function { arity = a1; pure = p1; res = r1 }
+    , Function { arity = a2; pure = p2; res = r2 } ) ->
+      a1 = a2 && Bool.(p1 = p2) && equal r1 r2
+  | Block b1, Block b2 -> (
+      try List.for_all2 ~f:equal b1 b2 with Invalid_argument _ -> false)
+  | Top _, (Function _ | Block _)
+  | Function _, (Top _ | Block _)
+  | Block _, (Top _ | Function _) -> false
+
 let rec to_string (shape : t) =
   match shape with
   | Top s -> if true then "N" else Printf.sprintf "N(%s)" s
@@ -43,13 +57,77 @@ module Store = struct
     let hash = Hashtbl.hash
   end)
 
+  let ext = ".jsoo-shape"
+
+  let filename ~dir ~name = Filename.concat dir (name ^ ext)
+
   let t = T.create 17
 
+  let loaded = Hashtbl.create 17
+
   let set ~name shape = T.replace t name shape
 
   let get ~name = T.find_opt t name
 
-  let load ~name:_ _dirs = None
+  let magic = "JsooShape000"
+
+  let load' fn =
+    let ic = open_in_bin fn in
+    let m = really_input_string ic (String.length magic) in
+    if not (String.equal m magic)
+    then failwith (Printf.sprintf "Invalid magic number for shape file %s" fn);
+    let shapes : (string * shape) list = Marshal.from_channel ic in
+    close_in ic;
+    List.iter shapes ~f:(fun (name, shape) -> set ~name shape)
+
+  let load ~name dirs =
+    if T.mem t name
+    then get ~name
+    else
+      match Fs.find_in_path dirs (filename ~dir:"." ~name) with
+      | Some f ->
+          load' f;
+          get ~name
+      | None ->
+          let rec scan : _ -> shape option = function
+            | [] -> None
+            | dir :: xs -> (
+                let l =
+                  Sys.readdir dir
+                  |> Array.to_list
+                  |> List.sort ~cmp:String.compare
+                  |> List.map ~f:(fun n -> Filename.concat dir n)
+                in
+                match
+                  List.find_map l ~f:(fun s ->
+                      if Filename.check_suffix s ext && not (Hashtbl.mem loaded s)
+                      then (
+                        load' s;
+                        Hashtbl.add loaded s ();
+                        match get ~name with
+                        | None -> None
+                        | Some shape -> Some (s, shape))
+                      else None)
+                with
+                | None -> scan xs
+                | Some (fn, shape) ->
+                    Format.eprintf "Shape: %s loaded from %s\n" name fn;
+                    Some shape)
+          in
+          scan dirs
+
+  let save' fn (l : (string * shape) list) =
+    let oc = open_out_bin fn in
+    output_string oc magic;
+    Marshal.to_channel oc l [];
+    close_out oc
+
+  let save ~name ~dir =
+    match get ~name with
+    | None -> failwith (Printf.sprintf "Don't know any shape for %s" name)
+    | Some shape ->
+        let fn = filename ~dir ~name in
+        save' fn [ name, shape ]
 end
 
 module State = struct

compiler/lib/shape.mli

@@ -28,12 +28,22 @@ type t =
 
 val to_string : t -> string
 
+val equal : t -> t -> bool
+
 module Store : sig
+  val ext : string
+
   val set : name:string -> t -> unit
 
   val get : name:string -> t option
 
+  val load' : string -> unit
+
   val load : name:string -> string list -> t option
+
+  val save : name:string -> dir:string -> unit
+
+  val save' : string -> (string * t) list -> unit
 end
 
 module State : sig