Compiler: wasmoo support for shapes

Author: hhugo

compiler/bin-wasm_of_ocaml/cmd_arg.ml

@@ -63,6 +63,7 @@ type t =
   ; params : (string * string) list
   ; include_dirs : string list
   ; effects : Config.effects_backend
+  ; shape_files : string list
   }
 
 let options () =
@@ -78,6 +79,10 @@ let options () =
     let doc = "Compile the bytecode program [$(docv)]. " in
     Arg.(required & pos ~rev:true 0 (some string) None & info [] ~docv:"PROGRAM" ~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 profile =
     let doc = "Set optimization profile : [$(docv)]." in
     let profile =
@@ -140,7 +145,8 @@ let options () =
       output_file
       input_file
       runtime_files
-      effects =
+      effects
+      shape_files =
     let chop_extension s = try Filename.chop_extension s with Invalid_argument _ -> s in
     let output_file =
       let ext =
@@ -172,6 +178,7 @@ let options () =
       ; sourcemap_root
       ; sourcemap_don't_inline_content
       ; effects
+      ; shape_files
       }
   in
   let t =
@@ -189,7 +196,8 @@ let options () =
       $ output_file
       $ input_file
       $ runtime_files
-      $ effects)
+      $ effects
+      $ shape_files)
   in
   Term.ret t
 
@@ -270,6 +278,7 @@ let options_runtime_only () =
       ; sourcemap_root
       ; sourcemap_don't_inline_content
       ; effects
+      ; shape_files = []
       }
   in
   let t =

compiler/bin-wasm_of_ocaml/cmd_arg.mli

@@ -32,6 +32,7 @@ type t =
   ; params : (string * string) list
   ; include_dirs : string list
   ; effects : Config.effects_backend
+  ; shape_files : string list
   }
 
 val options : unit -> t Cmdliner.Term.t

compiler/bin-wasm_of_ocaml/compile.ml

@@ -245,7 +245,15 @@ let generate_prelude ~out_file =
   @@ fun ch ->
   let code, uinfo = Parse_bytecode.predefined_exceptions () in
   let profile = Profile.O1 in
-  let Driver.{ program; variable_uses; in_cps; deadcode_sentinal; _ }, global_flow_data =
+  let ( Driver.
+          { program
+          ; variable_uses
+          ; in_cps
+          ; deadcode_sentinal
+          ; shapes = _
+          ; trampolined_calls = _
+          }
+      , global_flow_data ) =
     Driver.optimize_for_wasm ~profile ~shapes:false code
   in
   let context = Generate.start () in
@@ -328,6 +336,16 @@ let add_source_map sourcemap_don't_inline_content z opt_source_map =
                 ~name:(Link.source_name i j file)
                 ~contents:(Yojson.Basic.to_string (`String sm))))
 
+let merge_shape a b =
+  StringMap.union (fun _name s1 s2 -> if Shape.equal s1 s2 then Some s1 else None) a b
+
+let sexp_of_shapes s =
+  StringMap.bindings s
+  |> List.map ~f:(fun (name, shape) ->
+         Sexp.List [ Atom name; Atom (Shape.to_string shape) ])
+
+let string_of_shapes s = Sexp.List (sexp_of_shapes s) |> Sexp.to_string
+
 let run
     { Cmd_arg.common
     ; profile
@@ -341,11 +359,24 @@ let run
     ; sourcemap_root
     ; sourcemap_don't_inline_content
     ; effects
+    ; shape_files
     } =
   Config.set_target `Wasm;
   Jsoo_cmdline.Arg.eval common;
   Config.set_effects_backend effects;
   Generate.init ();
+  List.iter shape_files ~f:(fun s ->
+      let z = Zip.open_in s in
+      if Zip.has_entry z ~name:"shapes.sexp"
+      then
+        let s = Zip.read_entry z ~name:"shapes.sexp" in
+        match Sexp.from_string s with
+        | List l ->
+            List.iter l ~f:(function
+              | Sexp.List [ Atom name; Atom shape ] ->
+                  Shape.Store.set ~name (Shape.of_string shape)
+              | _ -> ())
+        | _ -> ());
   let output_file = fst output_file in
   if debug_mem () then Debug.start_profiling output_file;
   List.iter params ~f:(fun (s, v) -> Config.Param.set s v);
@@ -398,10 +429,18 @@ let run
     check_debug one;
     let code = one.code in
     let standalone = Option.is_none unit_name in
-    let Driver.{ program; variable_uses; in_cps; deadcode_sentinal; _ }, global_flow_data
-        =
-      Driver.optimize_for_wasm ~profile ~shapes:false code
+    let ( Driver.
+            { program
+            ; variable_uses
+            ; in_cps
+            ; deadcode_sentinal
+            ; shapes
+            ; trampolined_calls = _
+            }
+        , global_flow_data ) =
+      Driver.optimize_for_wasm ~profile ~shapes:true code
     in
+    StringMap.iter (fun name shape -> Shape.Store.set ~name shape) shapes;
     let context = Generate.start () in
     let toplevel_name, generated_js =
       Generate.f
@@ -423,7 +462,7 @@ let run
       Generate.output ch ~context;
       close_out ch);
     if times () then Format.eprintf "compilation: %a@." Timer.print t;
-    generated_js
+    generated_js, shapes
   in
   (if runtime_only
    then (
@@ -479,7 +518,7 @@ let run
           then Some (Filename.temp_file unit_name ".wasm.map")
           else None)
        @@ fun opt_tmp_map_file ->
-       let unit_data =
+       let unit_data, shapes =
          Fs.with_intermediate_file (Filename.temp_file unit_name ".wasm")
          @@ fun input_file ->
          opt_with
@@ -488,7 +527,7 @@ let run
             then Some (Filename.temp_file unit_name ".wasm.map")
             else None)
          @@ fun opt_input_sourcemap ->
-         let fragments =
+         let fragments, shapes =
            output
              code
              ~wat_file:
@@ -504,9 +543,9 @@ let run
            ~input_file
            ~output_file:tmp_wasm_file
            ();
-         { Link.unit_name; unit_info; fragments }
+         { Link.unit_name; unit_info; fragments }, shapes
        in
-       cont unit_data unit_name tmp_wasm_file opt_tmp_map_file
+       cont unit_data unit_name tmp_wasm_file opt_tmp_map_file shapes
      in
      (match kind with
      | `Exe ->
@@ -537,7 +576,7 @@ let run
            then Some (Filename.temp_file "code" ".wasm.map")
            else None
          in
-         let generated_js =
+         let generated_js, _shapes =
            output
              code
              ~unit_name:None
@@ -601,8 +640,9 @@ let run
          @@ fun tmp_output_file ->
          let z = Zip.open_out tmp_output_file in
          let compile_cmo' z cmo =
-           compile_cmo cmo (fun unit_data _ tmp_wasm_file opt_tmp_map_file ->
+           compile_cmo cmo (fun unit_data _ tmp_wasm_file opt_tmp_map_file shapes ->
                Zip.add_file z ~name:"code.wasm" ~file:tmp_wasm_file;
+               Zip.add_entry z ~name:"shapes.sexp" ~contents:(string_of_shapes shapes);
                add_source_map sourcemap_don't_inline_content z (`File opt_tmp_map_file);
                unit_data)
          in
@@ -618,8 +658,8 @@ let run
            List.fold_right
              ~f:(fun cmo cont l ->
                compile_cmo cmo
-               @@ fun unit_data unit_name tmp_wasm_file opt_tmp_map_file ->
-               cont ((unit_data, unit_name, tmp_wasm_file, opt_tmp_map_file) :: l))
+               @@ fun unit_data unit_name tmp_wasm_file opt_tmp_map_file shapes ->
+               cont ((unit_data, unit_name, tmp_wasm_file, opt_tmp_map_file, shapes) :: l))
              cma.lib_units
              ~init:(fun l ->
                Fs.with_intermediate_file (Filename.temp_file "wasm" ".wasm")
@@ -628,7 +668,7 @@ let run
                let source_map =
                  Wasm_link.f
                    (List.map
-                      ~f:(fun (_, _, file, opt_source_map) ->
+                      ~f:(fun (_, _, file, opt_source_map, _) ->
                         { Wasm_link.module_name = "OCaml"
                         ; file
                         ; code = None
@@ -641,10 +681,17 @@ let run
                    ~output_file:tmp_wasm_file
                in
                Zip.add_file z ~name:"code.wasm" ~file:tmp_wasm_file;
+               let shapes =
+                 List.fold_left
+                   ~init:StringMap.empty
+                   ~f:(fun acc (_, _, _, _, shapes) -> merge_shape acc shapes)
+                   l
+               in
+               Zip.add_entry z ~name:"shapes.sexp" ~contents:(string_of_shapes shapes);
                if enable_source_maps
                then
                  add_source_map sourcemap_don't_inline_content z (`Source_map source_map);
-               List.map ~f:(fun (unit_data, _, _, _) -> unit_data) l)
+               List.map ~f:(fun (unit_data, _, _, _, _) -> unit_data) l)
              []
          in
          Link.add_info z ~build_info:(Build_info.create `Cma) ~unit_data ();

compiler/lib-wasm/link.ml

@@ -345,6 +345,7 @@ let info_from_sexp info =
   let build_info =
     info |> member "build_info" |> mandatory (single Build_info.from_sexp)
   in
+
   let predefined_exceptions =
     info
     |> member "predefined_exceptions"

compiler/lib/driver.ml

@@ -120,6 +120,8 @@ let collects_shapes ~shapes (p : Code.program) =
                   | Utf (Utf8 s) -> s
                 in
                 shapes := StringMap.add name block !shapes
+            | Code.Let (_, Prim (Extern "caml_set_global", [ Pc (String name); Pv block ]))
+              -> shapes := StringMap.add name block !shapes
             | _ -> ()))
       p.blocks;
     let map =

compiler/lib/parse_bytecode.ml

@@ -844,6 +844,9 @@ let get_global state instrs i =
                 let x, state = State.fresh_var state in
                 if debug_parser ()
                 then Format.printf "%a = get_global(%s)@." Var.print x name;
+                (match Shape.Store.load ~name with
+                | None -> ()
+                | Some shape -> Shape.State.assign x shape);
                 ( x
                 , state
                 , Let (x, Prim (Extern "caml_get_global", [ Pc (String name) ])) :: instrs