Use a specific runtime for CPS-based effects (#1920)

* Support for several Wasm runtimes (depending on flags)

The runtimes will either be precompiled for most common flag
combinations, or compiled on the flight for unusual combinations.

* Use a specific runtime for CPS-based effects

Author: vouillon

compiler/bin-wasm_of_ocaml/compile.ml

@@ -83,6 +83,42 @@ let with_runtime_files ~runtime_wasm_files f =
   in
   Wat_preprocess.with_preprocessed_files ~variables:[] ~inputs f
 
+let build_runtime ~runtime_file =
+  (* Keep this variables in sync with gen/gen.ml *)
+  let variables =
+    [ ( "effects"
+      , Wat_preprocess.String
+          (match Config.effects () with
+          | `Jspi -> "jspi"
+          | `Cps -> "cps"
+          | `Disabled | `Double_translation -> assert false) )
+    ]
+  in
+  match
+    List.find_opt Runtime_files.precompiled_runtimes ~f:(fun (flags, _) ->
+        assert (
+          List.length flags = List.length variables
+          && List.for_all2 ~f:(fun (k, _) (k', _) -> String.equal k k') flags variables);
+        Poly.equal flags variables)
+  with
+  | Some (_, contents) -> Fs.write_file ~name:runtime_file ~contents
+  | None ->
+      let inputs =
+        List.map
+          ~f:(fun (module_name, contents) ->
+            { Wat_preprocess.module_name
+            ; file = module_name ^ ".wat"
+            ; source = Contents contents
+            })
+          Runtime_files.wat_files
+      in
+      Runtime.build
+        ~link_options:[ "-g" ]
+        ~opt_options:[ "-g"; "-O2" ]
+        ~variables
+        ~inputs
+        ~output_file:runtime_file
+
 let link_and_optimize
     ~profile
     ~sourcemap_root
@@ -101,7 +137,7 @@ let link_and_optimize
   let enable_source_maps = Option.is_some opt_sourcemap_file in
   Fs.with_intermediate_file (Filename.temp_file "runtime" ".wasm")
   @@ fun runtime_file ->
-  Fs.write_file ~name:runtime_file ~contents:Runtime_files.wasm_runtime;
+  build_runtime ~runtime_file;
   Fs.with_intermediate_file (Filename.temp_file "wasm-merged" ".wasm")
   @@ fun temp_file ->
   opt_with
@@ -147,7 +183,7 @@ let link_and_optimize
 
 let link_runtime ~profile runtime_wasm_files output_file =
   if List.is_empty runtime_wasm_files
-  then Fs.write_file ~name:output_file ~contents:Runtime_files.wasm_runtime
+  then build_runtime ~runtime_file:output_file
   else
     Fs.with_intermediate_file (Filename.temp_file "extra_runtime" ".wasm")
     @@ fun extra_runtime ->
@@ -169,7 +205,7 @@ let link_runtime ~profile runtime_wasm_files output_file =
       ();
     Fs.with_intermediate_file (Filename.temp_file "runtime" ".wasm")
     @@ fun runtime_file ->
-    Fs.write_file ~name:runtime_file ~contents:Runtime_files.wasm_runtime;
+    build_runtime ~runtime_file;
     Binaryen.link
       ~opt_output_sourcemap:None
       ~inputs:

compiler/bin-wasm_of_ocaml/dune

@@ -25,9 +25,10 @@
  (target runtime_files.ml)
  (deps
   gen/gen.exe
-  ../../runtime/wasm/runtime.wasm
   ../../runtime/wasm/runtime.js
-  ../../runtime/wasm/deps.json)
+  ../../runtime/wasm/deps.json
+  (glob_files ../../runtime/wasm/*.wat)
+  (glob_files ../../runtime/wasm/runtime-*.wasm))
  (action
   (with-stdout-to
    %{target}

compiler/bin-wasm_of_ocaml/gen/gen.ml

@@ -1,13 +1,77 @@
 let read_file ic = really_input_string ic (in_channel_length ic)
 
+(* Keep the two variables below in sync with function build_runtime in
+   ../compile.ml *)
+
+let default_flags = []
+
+let interesting_runtimes = [ [ "effects", `S "jspi" ]; [ "effects", `S "cps" ] ]
+
+let name_runtime standard l =
+  let flags =
+    List.filter_map
+      (fun (k, v) ->
+        match v with
+        | `S s -> Some s
+        | `B b -> if b then Some k else None)
+      l
+  in
+  String.concat "-" ("runtime" :: (if standard then [ "standard" ] else flags)) ^ ".wasm"
+
+let print_flags f flags =
+  Format.fprintf
+    f
+    "@[<2>[ %a ]@]"
+    (Format.pp_print_list
+       ~pp_sep:(fun f () -> Format.fprintf f ";@ ")
+       (fun f (k, v) ->
+         Format.fprintf
+           f
+           "@[\"%s\",@ %a@]"
+           k
+           (fun f v ->
+             match v with
+             | `S s -> Format.fprintf f "Wat_preprocess.String \"%s\"" s
+             | `B b ->
+                 Format.fprintf f "Wat_preprocess.Bool %s" (if b then "true" else "false"))
+           v))
+    flags
+
 let () =
   let () = set_binary_mode_out stdout true in
+  Format.printf "open Wasm_of_ocaml_compiler@.";
   Format.printf
-    "let wasm_runtime = \"%s\"@."
+    "let js_runtime = \"%s\"@."
     (String.escaped (read_file (open_in_bin Sys.argv.(1))));
   Format.printf
-    "let js_runtime = \"%s\"@."
+    "let dependencies = \"%s\"@."
     (String.escaped (read_file (open_in_bin Sys.argv.(2))));
+  let wat_files, runtimes =
+    List.partition
+      (fun f -> Filename.check_suffix f ".wat")
+      (Array.to_list (Array.sub Sys.argv 3 (Array.length Sys.argv - 3)))
+  in
   Format.printf
-    "let dependencies = \"%s\"@."
-    (String.escaped (read_file (open_in_bin Sys.argv.(3))))
+    "let wat_files = [%a]@."
+    (Format.pp_print_list (fun f file ->
+         Format.fprintf
+           f
+           "\"%s\", \"%s\"; "
+           Filename.(chop_suffix (basename file) ".wat")
+           (String.escaped (read_file (open_in_bin file)))))
+    wat_files;
+  Format.printf
+    "let precompiled_runtimes = [%a]@."
+    (Format.pp_print_list (fun f (standard, flags) ->
+         let flags = flags @ default_flags in
+         let name = name_runtime standard flags in
+         match List.find_opt (fun file -> Filename.basename file = name) runtimes with
+         | None -> failwith ("Missing runtime " ^ name)
+         | Some file ->
+             Format.fprintf
+               f
+               "%a, \"%s\"; "
+               print_flags
+               flags
+               (String.escaped (read_file (open_in_bin file)))))
+    (List.mapi (fun i flags -> i = 0, flags) interesting_runtimes)

compiler/lib-wasm/gc_target.ml

@@ -1704,18 +1704,6 @@ let post_process_function_body = Initialize_locals.f
 
 let entry_point ~toplevel_fun =
   let code =
-    let* () =
-      match Config.effects () with
-      | `Cps | `Double_translation ->
-          let* f =
-            register_import
-              ~name:"caml_cps_initialize_effects"
-              (Fun { W.params = []; result = [] })
-          in
-          instr (W.CallInstr (f, []))
-      | `Jspi -> return ()
-      | `Disabled -> assert false
-    in
     let* main =
       register_import
         ~name:"caml_main"

compiler/lib-wasm/generate.ml

@@ -1204,9 +1204,6 @@ let init () =
     ]
   in
   Primitive.register "caml_array_of_uniform_array" `Mutable None None;
-  let l =
-    if effects_cps () then ("caml_alloc_stack", "caml_cps_alloc_stack") :: l else l
-  in
   List.iter ~f:(fun (nm, nm') -> Primitive.alias nm nm') l
 
 (* Make sure we can use [br_table] for switches *)

runtime/wasm/dune

@@ -1,10 +1,13 @@
 (install
  (section lib)
  (package wasm_of_ocaml-compiler)
- (files runtime.wasm runtime.js))
+ (files
+  (glob_files *.wat)
+  (glob_files runtime-*.wasm)
+  runtime.js))
 
 (rule
- (target runtime.wasm)
+ (target runtime-standard.wasm)
  (deps
   args
   (glob_files *.wat))
@@ -13,6 +16,21 @@
    ../../compiler/bin-wasm_of_ocaml/wasmoo_link_wasm.exe
    --binaryen=-g
    --binaryen-opt=-O3
+   --set=effects=jspi
+   %{target}
+   %{read-lines:args})))
+
+(rule
+ (target runtime-cps.wasm)
+ (deps
+  args
+  (glob_files *.wat))
+ (action
+  (run
+   ../../compiler/bin-wasm_of_ocaml/wasmoo_link_wasm.exe
+   --binaryen=-g
+   --binaryen-opt=-O3
+   --set=effects=cps
    %{target}
    %{read-lines:args})))