Target-specific code (#1655)

Co-authored-by: Olivier Nicole <[email protected]>
Co-authored-by: Jérôme Vouillon <[email protected]>
Co-authored-by: Hugo Heuzard <[email protected]>

Author: 3 people

CHANGES.md

@@ -17,6 +17,8 @@
 * Runtime: allow dynlink of precompiled js with separate compilation (#1676)
 * Lib: Modify Typed_array API for compatibility with WebAssembly
 * Compiler: improved global dead code elimination (#2206)
+* Compiler: add support for the Wasm backend in parts of the pipeline, in
+  prevision for the merge of wasm_of_ocaml
 
 
 ## Bug fixes

compiler/bin-js_of_ocaml/build_fs.ml

@@ -78,7 +78,7 @@ function jsoo_create_file_extern(name,content){
           ~standalone:true
           ~wrap_with_fun:`Iife
           ~link:`Needed
-          pfs_fmt
+          ~formatter:pfs_fmt
           (Parse_bytecode.Debug.create ~include_cmis:false false)
           code
       in

compiler/bin-js_of_ocaml/check_runtime.ml

@@ -43,6 +43,7 @@ let print_groups output l =
               output_string output (Printf.sprintf "%s\n" name)))
 
 let f (runtime_files, bytecode, target_env) =
+  Config.set_target `JavaScript;
   Linker.reset ();
   let runtime_files, builtin =
     List.partition_map runtime_files ~f:(fun name ->

compiler/bin-js_of_ocaml/compile.ml

@@ -90,6 +90,7 @@ let run
     } =
   let include_cmis = toplevel && not no_cmis in
   let custom_header = common.Jsoo_cmdline.Arg.custom_header in
+  Config.set_target `JavaScript;
   Jsoo_cmdline.Arg.eval common;
   Linker.reset ();
   (match output_file with
@@ -186,7 +187,7 @@ let run
     let init_pseudo_fs = fs_external && standalone in
     let sm =
       match output_file with
-      | `Stdout, fmt ->
+      | `Stdout, formatter ->
           let instr =
             List.concat
               [ pseudo_fs_instr `create_file one.debug one.cmis
@@ -201,10 +202,10 @@ let run
             ~link
             ~wrap_with_fun
             ?source_map
-            fmt
+            ~formatter
             one.debug
             code
-      | `File, fmt ->
+      | `File, formatter ->
           let fs_instr1, fs_instr2 =
             match fs_output with
             | None -> pseudo_fs_instr `create_file one.debug one.cmis, []
@@ -225,7 +226,7 @@ let run
               ~link
               ~wrap_with_fun
               ?source_map
-              fmt
+              ~formatter
               one.debug
               code
           in

compiler/bin-js_of_ocaml/link.ml

@@ -150,6 +150,7 @@ let f
     ; mklib
     ; toplevel
     } =
+  Config.set_target `JavaScript;
   Jsoo_cmdline.Arg.eval common;
   Linker.reset ();
   let with_output f =

compiler/lib-dynlink/js_of_ocaml_compiler_dynlink.ml

@@ -26,6 +26,9 @@ let normalize_bytecode code =
       Bytes.to_string b
 
 let () =
+  (match Sys.backend_type with
+  | Sys.Other "js_of_ocaml" -> Config.set_target `JavaScript
+  | Sys.(Native | Bytecode | Other _) -> failwith "Expected backend `js_of_ocaml`");
   let global = J.pure_js_expr "globalThis" in
   Config.Flag.set "use-js-string" (Jsoo_runtime.Sys.Config.use_js_string ());
   Config.Flag.set "effects" (Jsoo_runtime.Sys.Config.effects ());

compiler/lib-runtime-files/gen/gen.ml

@@ -51,6 +51,7 @@ let rec list_product l =
 let bool = [ true; false ]
 
 let () =
+  Js_of_ocaml_compiler.Config.set_target `JavaScript;
   let () = set_binary_mode_out stdout true in
   match Array.to_list Sys.argv with
   | [] -> assert false

compiler/lib/code.ml

@@ -856,6 +856,7 @@ let with_invariant = Debug.find "invariant"
 let check_defs = false
 
 let invariant { blocks; start; _ } =
+  let target = Config.target () in
   if with_invariant ()
   then (
     assert (Addr.Map.mem start blocks);
@@ -870,15 +871,28 @@ let invariant { blocks; start; _ } =
         assert (not (Var.ISet.mem defs x));
         Var.ISet.add defs x)
     in
+    let check_constant = function
+      | NativeInt _ | Int32 _ ->
+          assert (
+            match target with
+            | `Wasm -> true
+            | _ -> false)
+      | String _ | NativeString _ | Float _ | Float_array _ | Int _ | Int64 _
+      | Tuple (_, _, _) -> ()
+    in
+    let check_prim_arg = function
+      | Pc c -> check_constant c
+      | Pv _ -> ()
+    in
     let check_expr = function
       | Apply _ -> ()
       | Block (_, _, _, _) -> ()
       | Field (_, _, _) -> ()
       | Closure (l, cont) ->
           List.iter l ~f:define;
           check_cont cont
-      | Constant _ -> ()
-      | Prim (_, _) -> ()
+      | Constant c -> check_constant c
+      | Prim (_, args) -> List.iter ~f:check_prim_arg args
       | Special _ -> ()
     in
     let check_instr (i, _loc) =

compiler/lib/config.ml

@@ -164,7 +164,7 @@ module Param = struct
     p
       ~name:"tc"
       ~desc:"Set tailcall optimisation"
-      (enum [ "trampoline", TcTrampoline; (* default *) "none", TcNone ])
+      (enum [ "trampoline", TcTrampoline (* default *); "none", TcNone ])
 
   let lambda_lifting_threshold =
     (* When we reach this depth, we start looking for functions to be lifted *)
@@ -183,8 +183,12 @@ end
 
 (****)
 
-let target_ : [ `JavaScript | `Wasm ] ref = ref `JavaScript
+let target_ : [ `JavaScript | `Wasm | `None ] ref = ref `None
 
-let target () = !target_
+let target () =
+  match !target_ with
+  | `None -> failwith "target was not set"
+  | (`JavaScript | `Wasm) as t -> t
 
-let set_target t = target_ := t
+let set_target (t : [ `JavaScript | `Wasm ]) =
+  target_ := (t :> [ `JavaScript | `Wasm | `None ])

compiler/lib/driver.ml

@@ -23,6 +23,14 @@ let debug = Debug.find "main"
 
 let times = Debug.find "times"
 
+type optimized_result =
+  { program : Code.program
+  ; variable_uses : Deadcode.variable_uses
+  ; trampolined_calls : Effects.trampolined_calls
+  ; in_cps : Effects.in_cps
+  ; deadcode_sentinal : Code.Var.t
+  }
+
 type profile =
   | O1
   | O2
@@ -194,14 +202,13 @@ let generate
     ~exported_runtime
     ~wrap_with_fun
     ~warn_on_unhandled_effect
-    ~deadcode_sentinal
-    ((p, live_vars), trampolined_calls, _) =
+    { program; variable_uses; trampolined_calls; deadcode_sentinal; in_cps = _ } =
   if times () then Format.eprintf "Start Generation...@.";
   let should_export = should_export wrap_with_fun in
   Generate.f
-    p
+    program
     ~exported_runtime
-    ~live_vars
+    ~live_vars:variable_uses
     ~trampolined_calls
     ~should_export
     ~warn_on_unhandled_effect
@@ -658,13 +665,19 @@ let configure formatter =
   Code.Var.set_pretty (pretty && not (Config.Flag.shortvar ()));
   Code.Var.set_stable (Config.Flag.stable_var ())
 
-let full ~standalone ~wrap_with_fun ~profile ~link ~source_map formatter d p =
-  let exported_runtime = not standalone in
+let link_and_pack ?(standalone = true) ?(wrap_with_fun = `Iife) ?(link = `No) p =
   let export_runtime =
     match link with
     | `All | `All_from _ -> true
     | `Needed | `No -> false
   in
+  p
+  |> link' ~export_runtime ~standalone ~link
+  |> pack ~wrap_with_fun ~standalone
+  |> coloring
+  |> check_js
+
+let optimize ~profile p =
   let deadcode_sentinal =
     (* If deadcode is disabled, this field is just fresh variable *)
     Code.Var.fresh_n "dummy"
@@ -677,31 +690,31 @@ let full ~standalone ~wrap_with_fun ~profile ~link ~source_map formatter d p =
        | O3 -> o3)
     +> exact_calls ~deadcode_sentinal profile
     +> effects ~deadcode_sentinal
-    +> map_fst (if Config.Flag.effects () then fun x -> x else Generate_closure.f)
+    +> map_fst
+         (match Config.target (), Config.Flag.effects () with
+         | `JavaScript, false -> Generate_closure.f
+         | `JavaScript, true | `Wasm, _ -> Fun.id)
     +> map_fst deadcode'
   in
-  let emit =
-    generate
-      d
-      ~exported_runtime
-      ~wrap_with_fun
-      ~warn_on_unhandled_effect:standalone
-      ~deadcode_sentinal
-    +> link' ~export_runtime ~standalone ~link
-    +> pack ~wrap_with_fun ~standalone
-    +> coloring
-    +> check_js
-    +> output formatter ~source_map ()
-  in
   if times () then Format.eprintf "Start Optimizing...@.";
   let t = Timer.make () in
-  let r = opt p in
+  let (program, variable_uses), trampolined_calls, in_cps = opt p in
   let () = if times () then Format.eprintf " optimizations : %a@." Timer.print t in
-  emit r
+  { program; variable_uses; trampolined_calls; in_cps; deadcode_sentinal }
+
+let full ~standalone ~wrap_with_fun ~profile ~link ~source_map ~formatter d p =
+  let optimized_code = optimize ~profile p in
+  let exported_runtime = not standalone in
+  let emit formatter =
+    generate d ~exported_runtime ~wrap_with_fun ~warn_on_unhandled_effect:standalone
+    +> link_and_pack ~standalone ~wrap_with_fun ~link
+    +> output formatter ~source_map ()
+  in
+  emit formatter optimized_code
 
-let full_no_source_map ~standalone ~wrap_with_fun ~profile ~link formatter d p =
+let full_no_source_map ~formatter ~standalone ~wrap_with_fun ~profile ~link d p =
   let (_ : Source_map.t option) =
-    full ~standalone ~wrap_with_fun ~profile ~link ~source_map:None formatter d p
+    full ~standalone ~wrap_with_fun ~profile ~link ~source_map:None ~formatter d p
   in
   ()
 
@@ -711,22 +724,22 @@ let f
     ?(profile = O1)
     ~link
     ?source_map
-    formatter
+    ~formatter
     d
     p =
-  full ~standalone ~wrap_with_fun ~profile ~link ~source_map formatter d p
+  full ~standalone ~wrap_with_fun ~profile ~link ~source_map ~formatter d p
 
 let f' ?(standalone = true) ?(wrap_with_fun = `Iife) ?(profile = O1) ~link formatter d p =
-  full_no_source_map ~standalone ~wrap_with_fun ~profile ~link formatter d p
+  full_no_source_map ~formatter ~standalone ~wrap_with_fun ~profile ~link d p
 
 let from_string ~prims ~debug s formatter =
   let p, d = Parse_bytecode.from_string ~prims ~debug s in
   full_no_source_map
+    ~formatter
     ~standalone:false
     ~wrap_with_fun:`Anonymous
     ~profile:O1
     ~link:`No
-    formatter
     d
     p
 

compiler/lib/driver.mli

@@ -20,13 +20,23 @@
 
 type profile
 
+type optimized_result =
+  { program : Code.program
+  ; variable_uses : Deadcode.variable_uses
+  ; trampolined_calls : Effects.trampolined_calls
+  ; in_cps : Effects.in_cps
+  ; deadcode_sentinal : Code.Var.t
+  }
+
+val optimize : profile:profile -> Code.program -> optimized_result
+
 val f :
      ?standalone:bool
   -> ?wrap_with_fun:[ `Iife | `Anonymous | `Named of string ]
   -> ?profile:profile
   -> link:[ `All | `All_from of string list | `Needed | `No ]
   -> ?source_map:Source_map.t
-  -> Pretty_print.t
+  -> formatter:Pretty_print.t
   -> Parse_bytecode.Debug.t
   -> Code.program
   -> Source_map.t option
@@ -48,6 +58,13 @@ val from_string :
   -> Pretty_print.t
   -> unit
 
+val link_and_pack :
+     ?standalone:bool
+  -> ?wrap_with_fun:[ `Iife | `Anonymous | `Named of string ]
+  -> ?link:[ `All | `All_from of string list | `Needed | `No ]
+  -> Javascript.statement_list
+  -> Javascript.statement_list
+
 val configure : Pretty_print.t -> unit
 
 val profiles : (int * profile) list