Add sourcemap support to wasm binary format

Author: vouillon

compiler/bin-wasm_of_ocaml/compile.ml

@@ -74,8 +74,6 @@ let opt_with action x f =
   | None -> f None
   | Some x -> action x (fun y -> f (Some y))
 
-let output_gen output_file f = Filename.gen_file output_file f
-
 let with_runtime_files ~runtime_wasm_files f =
   let inputs =
     List.map
@@ -163,8 +161,12 @@ let link_and_optimize
   @@ fun runtime_inputs ->
   Binaryen.link
     ~inputs:
-      (({ Binaryen.module_name = "env"; file = runtime_file } :: runtime_inputs)
-      @ List.map ~f:(fun file -> { Binaryen.module_name = "OCaml"; file }) wat_files)
+      ({ Binaryen.module_name = "env"; file = runtime_file; source_map_file = None }
+       :: runtime_inputs
+      @ List.map
+          ~f:(fun (file, source_map_file) ->
+            { Binaryen.module_name = "OCaml"; file; source_map_file })
+          wat_files)
     ~opt_output_sourcemap:opt_temp_sourcemap
     ~output_file:temp_file
     ());
@@ -223,7 +225,7 @@ let link_runtime ~profile runtime_wasm_files output_file =
       ~opt_output_sourcemap:None
       ~inputs:
         (List.map
-           ~f:(fun file -> { Binaryen.module_name = "env"; file })
+           ~f:(fun file -> { Binaryen.module_name = "env"; file; source_map_file = None })
            [ runtime_file; extra_runtime ])
       ~output_file
       ()
@@ -246,7 +248,7 @@ let generate_prelude ~out_file =
       ~deadcode_sentinal
       program
   in
-  Generate.wasm_output ch ~context;
+  Generate.wasm_output ch ~opt_source_map_file:None ~context;
   uinfo.provides
 
 let build_prelude z =
@@ -381,7 +383,7 @@ let run
     | Some p -> p
     | None -> Profile.O1
   in
-  let output (one : Parse_bytecode.one) ~unit_name ch =
+  let output (one : Parse_bytecode.one) ~unit_name ~wat_file ~file ~opt_source_map_file =
     check_debug one;
     let code = one.code in
     let standalone = Option.is_none unit_name in
@@ -399,7 +401,14 @@ let run
         program
     in
     if standalone then Generate.add_start_function ~context toplevel_name;
-    Generate.output ch ~enable_source_maps ~context;
+    let ch = open_out_bin file in
+    Generate.wasm_output ch ~opt_source_map_file ~context;
+    close_out ch;
+    if debug_wat ()
+    then (
+      let ch = open_out_bin wat_file in
+      Generate.output ch ~context;
+      close_out ch);
     if times () then Format.eprintf "compilation: %a@." Timer.print t;
     generated_js
   in
@@ -458,20 +467,28 @@ let run
           else None)
        @@ fun opt_tmp_map_file ->
        let unit_data =
-         (if debug_wat ()
-          then
-            fun f ->
-              f (Filename.concat (Filename.dirname output_file) (unit_name ^ ".wat"))
-          else Fs.with_intermediate_file (Filename.temp_file unit_name ".wat"))
-         @@ fun wat_file ->
+         Fs.with_intermediate_file (Filename.temp_file unit_name ".wasm")
+         @@ fun input_file ->
+         opt_with
+           Fs.with_intermediate_file
+           (if enable_source_maps
+            then Some (Filename.temp_file unit_name ".wasm.map")
+            else None)
+         @@ fun opt_input_sourcemap ->
          let strings, fragments =
-           output_gen wat_file (output code ~unit_name:(Some unit_name))
+           output
+             code
+             ~wat_file:
+               (Filename.concat (Filename.dirname output_file) (unit_name ^ ".wat"))
+             ~unit_name:(Some unit_name)
+             ~file:input_file
+             ~opt_source_map_file:opt_input_sourcemap
          in
          Binaryen.optimize
            ~profile
-           ~opt_input_sourcemap:None
+           ~opt_input_sourcemap
            ~opt_output_sourcemap:opt_tmp_map_file
-           ~input_file:wat_file
+           ~input_file
            ~output_file:tmp_wasm_file
            ();
          { Link.unit_name; unit_info; strings; fragments }
@@ -491,10 +508,8 @@ let run
              ic
          in
          if times () then Format.eprintf "  parsing: %a@." Timer.print t1;
-         (if debug_wat ()
-          then fun f -> f (Filename.chop_extension output_file ^ ".wat")
-          else Fs.with_intermediate_file (Filename.temp_file "code" ".wat"))
-         @@ fun wat_file ->
+         Fs.with_intermediate_file (Filename.temp_file "code" ".wasm")
+         @@ fun input_wasm_file ->
          let dir = Filename.chop_extension output_file ^ ".assets" in
          Link.gen_dir dir
          @@ fun tmp_dir ->
@@ -504,7 +519,19 @@ let run
            then Some (Filename.concat tmp_dir "code.wasm.map")
            else None
          in
-         let generated_js = output_gen wat_file (output code ~unit_name:None) in
+         let opt_source_map_file =
+           if enable_source_maps
+           then Some (Filename.temp_file "code" ".wasm.map")
+           else None
+         in
+         let generated_js =
+           output
+             code
+             ~unit_name:None
+             ~wat_file:(Filename.chop_extension output_file ^ ".wat")
+             ~file:input_wasm_file
+             ~opt_source_map_file
+         in
          let tmp_wasm_file = Filename.concat tmp_dir "code.wasm" in
          let primitives =
            link_and_optimize
@@ -513,7 +540,7 @@ let run
              ~sourcemap_don't_inline_content
              ~opt_sourcemap
              runtime_wasm_files
-             [ wat_file ]
+             [ input_wasm_file, opt_source_map_file ]
              tmp_wasm_file
          in
          let wasm_name =

compiler/lib-wasm/binaryen.ml

@@ -48,6 +48,7 @@ let opt_flag flag v =
 type link_input =
   { module_name : string
   ; file : string
+  ; source_map_file : string option
   }
 
 let link ?options ~inputs ~opt_output_sourcemap ~output_file () =
@@ -57,7 +58,13 @@ let link ?options ~inputs ~opt_output_sourcemap ~output_file () =
        @ Option.value ~default:[] options
        @ List.flatten
            (List.map
-              ~f:(fun { file; module_name } -> [ Filename.quote file; module_name ])
+              ~f:(fun { file; module_name; source_map_file } ->
+                Filename.quote file
+                :: module_name
+                ::
+                (match source_map_file with
+                | None -> []
+                | Some file -> [ "--input-source-map"; Filename.quote file ]))
               inputs)
        @ [ "-o"; Filename.quote output_file ]
        @ opt_flag "--output-source-map" opt_output_sourcemap))

compiler/lib-wasm/binaryen.mli

@@ -19,6 +19,7 @@
 type link_input =
   { module_name : string  (** Name under which the module is imported in other modules *)
   ; file : string  (** File containing the Wasm module *)
+  ; source_map_file : string option
   }
 
 val link :

compiler/lib-wasm/generate.ml

@@ -1315,16 +1315,14 @@ let add_start_function = G.add_start_function
 
 let add_init_function = G.add_init_function
 
-let output ch ~enable_source_maps ~context =
+let output ch ~context =
   let t = Timer.make () in
   let fields = G.output ~context in
-  if enable_source_maps || Debug.find "wat" ()
-  then Wat_output.f ch fields
-  else Wasm_output.f ch fields;
+  Wat_output.f ch fields;
   if times () then Format.eprintf "  output: %a@." Timer.print t
 
-let wasm_output ch ~context =
+let wasm_output ch ~opt_source_map_file ~context =
   let t = Timer.make () in
   let fields = G.output ~context in
-  Wasm_output.f ch fields;
+  Wasm_output.f ch ~opt_source_map_file fields;
   if times () then Format.eprintf "  output: %a@." Timer.print t

compiler/lib-wasm/generate.mli

@@ -33,7 +33,10 @@ val add_start_function : context:Code_generation.context -> Wasm_ast.var -> unit
 
 val add_init_function : context:Code_generation.context -> to_link:string list -> unit
 
-val output :
-  out_channel -> enable_source_maps:bool -> context:Code_generation.context -> unit
+val output : out_channel -> context:Code_generation.context -> unit
 
-val wasm_output : out_channel -> context:Code_generation.context -> unit
+val wasm_output :
+     out_channel
+  -> opt_source_map_file:string option
+  -> context:Code_generation.context
+  -> unit

compiler/lib-wasm/link.ml

@@ -402,7 +402,7 @@ let generate_start_function ~to_link ~out_file =
   @@ fun ch ->
   let context = Generate.start () in
   Generate.add_init_function ~context ~to_link:("prelude" :: to_link);
-  Generate.wasm_output ch ~context;
+  Generate.wasm_output ch ~opt_source_map_file:None ~context;
   if times () then Format.eprintf "    generate start: %a@." Timer.print t1
 
 let output_js js =

compiler/lib-wasm/wasm_output.ml

@@ -62,6 +62,10 @@ module Make (Output : sig
   val byte : t -> int -> unit
 
   val string : t -> string -> unit
+
+  val push_mapping : Source_map.map -> unit
+
+  val get_file_index : string -> int
 end) : sig
   val output_module : Output.t -> module_field list -> unit
 end = struct
@@ -460,6 +464,29 @@ end = struct
     ; current_local_names : int Code.Var.Hashtbl.t
     }
 
+  let last_event = ref None
+
+  let push_no_event ch =
+    if Option.is_some !last_event
+    then (
+      Output.push_mapping (Source_map.Gen { gen_line = 1; gen_col = position ch });
+      last_event := None)
+
+  let push_event ch ~src ~line ~col =
+    match !last_event with
+    | Some (src', line', col') when col = col' && line = line' && String.equal src src' ->
+        ()
+    | _ ->
+        Output.push_mapping
+          (Source_map.Gen_Ori
+             { gen_line = 1
+             ; gen_col = position ch
+             ; ori_source = Output.get_file_index src
+             ; ori_line = line
+             ; ori_col = col
+             });
+        last_event := Some (src, line, col)
+
   let rec output_expression st ch e =
     match e with
     | Const c -> (
@@ -772,7 +799,8 @@ end = struct
         output_byte ch 0x15;
         output_uint ch (Code.Var.Hashtbl.find st.type_names typ)
     | Unreachable -> output_byte ch 0x00
-    | Event _ -> ()
+    | Event Parse_info.{ src = None | Some ""; _ } -> push_no_event ch
+    | Event Parse_info.{ src = Some src; line; col; _ } -> push_event ch ~src ~line ~col
 
   let output_globals ch (st, global_idx, fields) =
     let count =
@@ -1034,7 +1062,8 @@ end = struct
               prerr_endline (Printexc.to_string e);
               prerr_endline backtrace;
               assert false);
-           output_byte ch 0x0B))
+           output_byte ch 0x0B;
+           push_no_event ch))
       ch
       (List.rev l)
 
@@ -1170,7 +1199,9 @@ end = struct
     output_section 0 output_features ch ()
 end
 
-let f ch fields =
+let f ~opt_source_map_file ch fields =
+  let mappings = ref [] in
+  let files = String.Hashtbl.create 16 in
   let module O = Make (struct
     type t = out_channel
 
@@ -1181,5 +1212,27 @@ let f ch fields =
     let byte = output_byte
 
     let string = output_string
+
+    let push_mapping m = mappings := m :: !mappings
+
+    let get_file_index file =
+      try String.Hashtbl.find files file
+      with Not_found ->
+        let pos = String.Hashtbl.length files in
+        String.Hashtbl.add files file pos;
+        pos
   end) in
-  O.output_module ch fields
+  O.output_module ch fields;
+  Option.iter opt_source_map_file ~f:(fun source_map_file ->
+      let hashtbl_to_list htb =
+        String.Hashtbl.fold (fun k v l -> (k, v) :: l) htb []
+        |> List.sort ~cmp:(fun (_, a) (_, b) -> compare a b)
+        |> List.map ~f:fst
+      in
+      let sm =
+        { (Source_map.Standard.empty ~inline_source_content:false) with
+          sources = hashtbl_to_list files
+        ; mappings = Source_map.Mappings.encode (List.rev !mappings)
+        }
+      in
+      Source_map.to_file ~rewrite_paths:false (Standard sm) source_map_file)

compiler/lib-wasm/wasm_output.mli

@@ -16,4 +16,5 @@
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  *)
 
-val f : out_channel -> Wasm_ast.module_field list -> unit
+val f :
+  opt_source_map_file:string option -> out_channel -> Wasm_ast.module_field list -> unit

compiler/lib-wasm/wat_preprocess.ml

@@ -626,7 +626,7 @@ let with_preprocessed_files ~variables ~inputs action =
             if Link.Wasm_binary.check_file ~file then None else Some (Fs.read_file file)
         | Contents contents -> Some contents
       with
-      | None -> cont ({ Binaryen.module_name; file } :: inputs)
+      | None -> cont ({ Binaryen.module_name; file; source_map_file = None } :: inputs)
       | Some contents ->
           let source_file = file in
           Fs.with_intermediate_file (Filename.temp_file module_name ".wat")
@@ -637,7 +637,7 @@ let with_preprocessed_files ~variables ~inputs action =
               (if Link.Wasm_binary.check ~contents
                then contents
                else f ~variables ~filename:source_file ~contents);
-          cont ({ Binaryen.module_name; file } :: inputs))
+          cont ({ Binaryen.module_name; file; source_map_file = None } :: inputs))
     ~init:action
     inputs
     []

compiler/lib/source_map.ml

@@ -574,7 +574,8 @@ module Standard = struct
 
   let to_string m = Yojson.Raw.to_string (json (rewrite_paths m))
 
-  let to_file m file = Yojson.Raw.to_file file (json (rewrite_paths m))
+  let to_file ?rewrite_paths:(rewrite = true) m file =
+    Yojson.Raw.to_file file (json (if rewrite then rewrite_paths m else m))
 
   let invariant
       { version
@@ -715,7 +716,8 @@ module Index = struct
 
   let to_string m = Yojson.Raw.to_string (json (rewrite_paths m))
 
-  let to_file m file = Yojson.Raw.to_file file (json (rewrite_paths m))
+  let to_file ?rewrite_paths:(rewrite = true) m file =
+    Yojson.Raw.to_file file (json (if rewrite then rewrite_paths m else m))
 
   let invariant { version; file = _; sections } =
     if not (version_is_valid version)
@@ -756,10 +758,10 @@ let to_string = function
   | Standard m -> Standard.to_string m
   | Index i -> Index.to_string i
 
-let to_file x f =
+let to_file ?rewrite_paths x f =
   match x with
-  | Standard m -> Standard.to_file m f
-  | Index i -> Index.to_file i f
+  | Standard m -> Standard.to_file ?rewrite_paths m f
+  | Index i -> Index.to_file ?rewrite_paths i f
 
 let invariant = function
   | Standard m -> Standard.invariant m

compiler/lib/source_map.mli

@@ -141,7 +141,7 @@ type t =
 
 val to_string : t -> string
 
-val to_file : t -> string -> unit
+val to_file : ?rewrite_paths:bool -> t -> string -> unit
 
 val of_string : ?tmp_buf:Buffer.t -> string -> t