Sourcemap support for wasm

Implement mapping between source and wasm locations.

To work, this requires a version of Binaryen compiled with Jérôme's
patch WebAssembly/binaryen#6372.

Single-stepping can jump around in slightly surprising ways in the OCaml
code, due to the different order of operations in wasm. This could be
improved by modifying Binaryen to support “no location” annotations.
Another future improvement can be to support mapping Wasm identifiers to
OCaml ones.

Co-authored-by: Jérôme Vouillon <jerome.vouillon@gmail.com>

compiler/bin-wasm_of_ocaml/cmd_arg.ml

@@ -28,6 +28,7 @@ type t =
   ; runtime_files : string list
   ; output_file : string * bool
   ; input_file : string
+  ; enable_source_maps : bool
   ; params : (string * string) list
   }
 
@@ -50,11 +51,11 @@ let options =
     Arg.(value & opt (some (enum profile)) None & info [ "opt" ] ~docv:"NUM" ~doc)
   in
   let no_sourcemap =
-    let doc = "Currently ignored (for compatibility with Js_of_ocaml)." in
+    let doc = "Disable sourcemap output." in
     Arg.(value & flag & info [ "no-sourcemap"; "no-source-map" ] ~doc)
   in
   let sourcemap =
-    let doc = "Currently ignored (for compatibility with Js_of_ocaml)." in
+    let doc = "Output source locations in a separate sourcemap file." in
     Arg.(value & flag & info [ "sourcemap"; "source-map" ] ~doc)
   in
   let sourcemap_inline_in_js =
@@ -69,24 +70,42 @@ let options =
       & opt_all (list (pair ~sep:'=' (enum all) string)) []
       & info [ "set" ] ~docv:"PARAM=VALUE" ~doc)
   in
-  let build_t common set_param profile _ _ _ output_file input_file runtime_files =
+  let build_t
+      common
+      set_param
+      profile
+      sourcemap
+      no_sourcemap
+      _
+      output_file
+      input_file
+      runtime_files =
     let chop_extension s = try Filename.chop_extension s with Invalid_argument _ -> s in
     let output_file =
       match output_file with
       | Some s -> s, true
       | None -> chop_extension input_file ^ ".js", false
     in
     let params : (string * string) list = List.flatten set_param in
-    `Ok { common; params; profile; output_file; input_file; runtime_files }
+    let enable_source_maps = (not no_sourcemap) && sourcemap in
+    `Ok
+      { common
+      ; params
+      ; profile
+      ; output_file
+      ; input_file
+      ; runtime_files
+      ; enable_source_maps
+      }
   in
   let t =
     Term.(
       const build_t
       $ Jsoo_cmdline.Arg.t
       $ set_param
       $ profile
-      $ no_sourcemap
       $ sourcemap
+      $ no_sourcemap
       $ sourcemap_inline_in_js
       $ output_file
       $ input_file

compiler/bin-wasm_of_ocaml/cmd_arg.mli

@@ -26,6 +26,7 @@ type t =
   ; runtime_files : string list
   ; output_file : string * bool
   ; input_file : string
+  ; enable_source_maps : bool
   ; params : (string * string) list
   }
 

compiler/bin-wasm_of_ocaml/compile.ml

@@ -84,15 +84,19 @@ let common_binaryen_options () =
   in
   if Config.Flag.pretty () then "-g" :: l else l
 
-let link runtime_files input_file output_file =
+let link ~enable_source_maps runtime_files input_file output_file =
   command
     ("wasm-merge"
     :: (common_binaryen_options ()
        @ List.flatten
            (List.map
               ~f:(fun runtime_file -> [ Filename.quote runtime_file; "env" ])
               runtime_files)
-       @ [ Filename.quote input_file; "exec"; "-o"; Filename.quote output_file ]))
+       @ [ Filename.quote input_file; "exec"; "-o"; Filename.quote output_file ]
+       @
+       if enable_source_maps
+       then [ "--output-source-map"; Filename.quote (output_file ^ ".map") ]
+       else []))
 
 let generate_dependencies primitives =
   Yojson.Basic.to_string
@@ -120,7 +124,7 @@ let filter_unused_primitives primitives usage_file =
    with End_of_file -> ());
   !s
 
-let dead_code_elimination in_file out_file =
+let dead_code_elimination ~enable_source_maps in_file out_file =
   with_intermediate_file (Filename.temp_file "deps" ".json")
   @@ fun deps_file ->
   with_intermediate_file (Filename.temp_file "usage" ".txt")
@@ -130,14 +134,15 @@ let dead_code_elimination in_file out_file =
   command
     ("wasm-metadce"
     :: (common_binaryen_options ()
-       @ [ "--graph-file"
-         ; Filename.quote deps_file
-         ; Filename.quote in_file
-         ; "-o"
-         ; Filename.quote out_file
-         ; ">"
-         ; Filename.quote usage_file
-         ]));
+       @ [ "--graph-file"; Filename.quote deps_file; Filename.quote in_file ]
+       @ (if enable_source_maps
+          then [ "--input-source-map"; Filename.quote (in_file ^ ".map") ]
+          else [])
+       @ [ "-o"; Filename.quote out_file ]
+       @ (if enable_source_maps
+          then [ "--output-source-map"; Filename.quote (out_file ^ ".map") ]
+          else [])
+       @ [ ">"; Filename.quote usage_file ]));
   filter_unused_primitives primitives usage_file
 
 let optimization_options =
@@ -146,29 +151,62 @@ let optimization_options =
    ; [ "-O3"; "--traps-never-happen" ]
   |]
 
-let optimize ~profile in_file out_file =
+let optimize ~profile ?sourcemap_file in_file out_file =
   let level =
     match profile with
     | None -> 1
     | Some p -> fst (List.find ~f:(fun (_, p') -> Poly.equal p p') Driver.profiles)
   in
   command
     ("wasm-opt"
-    :: (common_binaryen_options ()
-       @ optimization_options.(level - 1)
-       @ [ Filename.quote in_file; "-o"; Filename.quote out_file ]))
+     :: (common_binaryen_options ()
+        @ optimization_options.(level - 1)
+        @ [ Filename.quote in_file; "-o"; Filename.quote out_file ])
+    @
+    match sourcemap_file with
+    | Some sourcemap_file ->
+        [ "--input-source-map"
+        ; Filename.quote (in_file ^ ".map")
+        ; "--output-source-map"
+        ; Filename.quote sourcemap_file
+        ; "--output-source-map-url"
+        ; Filename.quote sourcemap_file
+        ]
+    | None -> [])
 
-let link_and_optimize ~profile runtime_wasm_files wat_file output_file =
+let link_and_optimize ~profile ?sourcemap_file runtime_wasm_files wat_file output_file =
+  let sourcemap_file =
+    (* Check that Binaryen supports the necessary sourcemaps options (requires
+       version >= 118) *)
+    match sourcemap_file with
+    | Some _ when Sys.command "wasm-merge -osm foo 2> /dev/null" <> 0 -> None
+    | Some _ | None -> sourcemap_file
+  in
+  let enable_source_maps = Option.is_some sourcemap_file in
   with_intermediate_file (Filename.temp_file "runtime" ".wasm")
   @@ fun runtime_file ->
   write_file runtime_file Wa_runtime.wasm_runtime;
   with_intermediate_file (Filename.temp_file "wasm-merged" ".wasm")
   @@ fun temp_file ->
-  link (runtime_file :: runtime_wasm_files) wat_file temp_file;
+  link ~enable_source_maps (runtime_file :: runtime_wasm_files) wat_file temp_file;
   with_intermediate_file (Filename.temp_file "wasm-dce" ".wasm")
   @@ fun temp_file' ->
-  let primitives = dead_code_elimination temp_file temp_file' in
-  optimize ~profile temp_file' output_file;
+  let primitives = dead_code_elimination ~enable_source_maps temp_file temp_file' in
+  optimize ~profile ?sourcemap_file temp_file' output_file;
+  (* Add source file contents to source map *)
+  Option.iter sourcemap_file ~f:(fun sourcemap_file ->
+      let open Source_map in
+      let source_map, mappings = Source_map_io.of_file_no_mappings sourcemap_file in
+      assert (List.is_empty (Option.value source_map.sources_content ~default:[]));
+      let sources_content =
+        Some
+          (List.map source_map.sources ~f:(fun file ->
+               if Sys.file_exists file && not (Sys.is_directory file)
+               then Some (Fs.read_file file)
+               else None))
+      in
+      let source_map = { source_map with sources_content } in
+      Source_map_io.to_file ?mappings source_map ~file:sourcemap_file);
   primitives
 
 let escape_string s =
@@ -276,7 +314,15 @@ let build_js_runtime primitives (strings, fragments) wasm_file output_file =
     ^ trim_semi (Buffer.contents fragment_buffer)
     ^ String.sub s ~pos:(l + 9) ~len:(String.length s - l - 9))
 
-let run { Cmd_arg.common; profile; runtime_files; input_file; output_file; params } =
+let run
+    { Cmd_arg.common
+    ; profile
+    ; runtime_files
+    ; input_file
+    ; output_file
+    ; enable_source_maps
+    ; params
+    } =
   Jsoo_cmdline.Arg.eval common;
   Wa_generate.init ();
   let output_file = fst output_file in
@@ -316,7 +362,7 @@ let run { Cmd_arg.common; profile; runtime_files; input_file; output_file; param
   let need_debug = Config.Flag.debuginfo () in
   let output (one : Parse_bytecode.one) ~standalone ch =
     let code = one.code in
-    let live_vars, in_cps, p =
+    let live_vars, in_cps, p, debug =
       Driver.f
         ~target:Wasm
         ~standalone
@@ -326,7 +372,7 @@ let run { Cmd_arg.common; profile; runtime_files; input_file; output_file; param
         one.debug
         code
     in
-    let strings = Wa_generate.f ch ~live_vars ~in_cps p in
+    let strings = Wa_generate.f ch ~debug ~live_vars ~in_cps p in
     if times () then Format.eprintf "compilation: %a@." Timer.print t;
     strings
   in
@@ -367,7 +413,13 @@ let run { Cmd_arg.common; profile; runtime_files; input_file; output_file; param
        @@ fun tmp_wasm_file ->
        let strings = output_gen wat_file (output code ~standalone:true) in
        let primitives =
-         link_and_optimize ~profile runtime_wasm_files wat_file tmp_wasm_file
+         link_and_optimize
+           ~profile
+           ?sourcemap_file:
+             (if enable_source_maps then Some (wasm_file ^ ".map") else None)
+           runtime_wasm_files
+           wat_file
+           tmp_wasm_file
        in
        build_js_runtime primitives strings wasm_file output_file
    | `Cmo _ | `Cma _ -> assert false);

compiler/bin-wasm_of_ocaml/dune

@@ -43,5 +43,4 @@
 (install
  (section man)
  (package wasm_of_ocaml-compiler)
- (files
-  wasm_of_ocaml.1))
+ (files wasm_of_ocaml.1))

compiler/lib/driver.ml

@@ -576,7 +576,9 @@ let configure formatter =
 
 type 'a target =
   | JavaScript : Pretty_print.t -> Source_map.t option target
-  | Wasm : (Deadcode.variable_uses * Effects.in_cps * Code.program) target
+  | Wasm
+      : (Deadcode.variable_uses * Effects.in_cps * Code.program * Parse_bytecode.Debug.t)
+        target
 
 let target_flag (type a) (t : a target) =
   match t with
@@ -631,7 +633,7 @@ let full
       source_map
   | Wasm ->
       let (p, live_vars), _, in_cps = r in
-      live_vars, in_cps, p
+      live_vars, in_cps, p, d
 
 let full_no_source_map ~formatter ~standalone ~wrap_with_fun ~profile ~linkall d p =
   let (_ : Source_map.t option) =

compiler/lib/driver.mli

@@ -22,7 +22,9 @@ type profile
 
 type 'a target =
   | JavaScript : Pretty_print.t -> Source_map.t option target
-  | Wasm : (Deadcode.variable_uses * Effects.in_cps * Code.program) target
+  | Wasm
+      : (Deadcode.variable_uses * Effects.in_cps * Code.program * Parse_bytecode.Debug.t)
+        target
 
 val f :
      target:'result target

compiler/lib/generate.ml

@@ -341,11 +341,14 @@ let bool e = J.ECond (e, one, zero)
 
 (****)
 
-let source_location ctx ?force (pc : Code.loc) =
-  match Parse_bytecode.Debug.find_loc ctx.Ctx.debug ?force pc with
+let source_location debug ?force (pc : Code.loc) =
+  match Parse_bytecode.Debug.find_loc debug ?force pc with
   | Some pi -> J.Pi pi
   | None -> J.N
 
+let source_location_ctx ctx ?force (pc : Code.loc) =
+  source_location ctx.Ctx.debug ?force pc
+
 (****)
 
 let float_const f = J.ENum (J.Num.of_float f)
@@ -1240,13 +1243,13 @@ let rec translate_expr ctx queue loc x e level : _ * J.statement_list =
       let (px, cx), queue = access_queue queue x in
       (Mlvalue.Block.field cx n, or_p px mutable_p, queue), []
   | Closure (args, ((pc, _) as cont)) ->
-      let loc = source_location ctx ~force:After (After pc) in
+      let loc = source_location_ctx ctx ~force:After (After pc) in
       let clo = compile_closure ctx cont in
       let clo =
         match clo with
         | (st, x) :: rem ->
             let loc =
-              match x, source_location ctx (Before pc) with
+              match x, source_location_ctx ctx (Before pc) with
               | (J.U | J.N), (J.U | J.N) -> J.U
               | x, (J.U | J.N) -> x
               | (J.U | J.N), x -> x
@@ -1495,14 +1498,14 @@ and translate_instr ctx expr_queue instr =
   let instr, pc = instr in
   match instr with
   | Assign (x, y) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       let (_py, cy), expr_queue = access_queue expr_queue y in
       flush_queue
         expr_queue
         mutator_p
         [ J.Expression_statement (J.EBin (J.Eq, J.EVar (J.V x), cy)), loc ]
   | Let (x, e) -> (
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       let (ce, prop, expr_queue), instrs = translate_expr ctx expr_queue loc x e 0 in
       let keep_name x =
         match Code.Var.get_name x with
@@ -1533,23 +1536,23 @@ and translate_instr ctx expr_queue instr =
             prop
             (instrs @ [ J.variable_declaration [ J.V x, (ce, loc) ], loc ]))
   | Set_field (x, n, y) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       let (_px, cx), expr_queue = access_queue expr_queue x in
       let (_py, cy), expr_queue = access_queue expr_queue y in
       flush_queue
         expr_queue
         mutator_p
         [ J.Expression_statement (J.EBin (J.Eq, Mlvalue.Block.field cx n, cy)), loc ]
   | Offset_ref (x, 1) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       (* FIX: may overflow.. *)
       let (_px, cx), expr_queue = access_queue expr_queue x in
       flush_queue
         expr_queue
         mutator_p
         [ J.Expression_statement (J.EUn (J.IncrA, Mlvalue.Block.field cx 0)), loc ]
   | Offset_ref (x, n) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       (* FIX: may overflow.. *)
       let (_px, cx), expr_queue = access_queue expr_queue x in
       flush_queue
@@ -1558,7 +1561,7 @@ and translate_instr ctx expr_queue instr =
         [ J.Expression_statement (J.EBin (J.PlusEq, Mlvalue.Block.field cx 0, int n)), loc
         ]
   | Array_set (x, y, z) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       let (_px, cx), expr_queue = access_queue expr_queue x in
       let (_py, cy), expr_queue = access_queue expr_queue y in
       let (_pz, cz), expr_queue = access_queue expr_queue z in
@@ -1619,7 +1622,7 @@ and compile_block st queue (pc : Addr.t) loop_stack frontier interm =
                      else (
                        if debug () then Format.eprintf "break;@;}@]@,";
                        body @ [ J.Break_statement None, J.N ])) )
-            , source_location st.ctx (Code.location_of_pc pc) )
+            , source_location_ctx st.ctx (Code.location_of_pc pc) )
           in
           let label = if !lab_used then Some lab else None in
           let for_loop =
@@ -1854,7 +1857,7 @@ and compile_conditional st queue last loop_stack backs frontier interm =
      | Stop -> Format.eprintf "stop;@;"
      | Cond (x, _, _) -> Format.eprintf "@[<hv 2>cond(%a){@;" Code.Var.print x
      | Switch (x, _, _) -> Format.eprintf "@[<hv 2>switch(%a){@;" Code.Var.print x);
-  let loc = source_location st.ctx pc in
+  let loc = source_location_ctx st.ctx pc in
   let res =
     match last with
     | Return x ->