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 <[email protected]>
OlivierNicole · Mar 15, 2024 · 273c77a · 273c77a
1 parent 5ca926f
commit 273c77a
Show file tree

Hide file tree

Showing 19 changed files with 491 additions and 330 deletions.
diff --git a/compiler/bin-wasm_of_ocaml/cmd_arg.ml b/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,41 @@ 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

diff --git a/compiler/bin-wasm_of_ocaml/cmd_arg.mli b/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
   }
 

diff --git a/compiler/bin-wasm_of_ocaml/compile.ml b/compiler/bin-wasm_of_ocaml/compile.ml
@@ -83,15 +83,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" ])
+              ~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
@@ -119,7 +123,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")
@@ -131,21 +135,27 @@ let dead_code_elimination in_file out_file =
     :: (common_binaryen_options ()
        @ [ "--graph-file"
          ; Filename.quote deps_file
-         ; Filename.quote in_file
-         ; "-o"
-         ; Filename.quote out_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 =
-  [| [ "-O2"; "--skip-pass=inlining-optimizing" ]
+  [| [ "--simplify-locals-notee-nostructure"; "--vacuum"; "--reorder-locals"]
    ; [ "-O2"; "--skip-pass=inlining-optimizing"; "--traps-never-happen" ]
    ; [ "-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
@@ -155,19 +165,48 @@ let optimize ~profile in_file out_file =
     ("wasm-opt"
     :: (common_binaryen_options ()
        @ optimization_options.(level - 1)
-       @ [ Filename.quote in_file; "-o"; Filename.quote out_file ]))
+       @ [ 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 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 =
@@ -274,7 +313,14 @@ let build_js_runtime primitives (strings, fragments) wasm_file output_file =
     ^ trim_semi (Buffer.contents fragment_buffer)
     ^ String.sub s ~pos:(k + 7) ~len:(String.length s - k - 7))
 
-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
@@ -364,7 +410,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);

diff --git a/compiler/lib/driver.ml b/compiler/lib/driver.ml
@@ -618,7 +618,7 @@ let full ~target ~standalone ~wrap_with_fun ~profile ~linkall ~source_map d p =
       source_map, ([], [])
   | `Wasm ch ->
       let (p, live_vars), _, in_cps = r in
-      None, Wa_generate.f ch ~live_vars ~in_cps p
+      None, Wa_generate.f ~debug:d ch ~live_vars ~in_cps p
 
 let full_no_source_map ~target ~standalone ~wrap_with_fun ~profile ~linkall d p =
   let (_ : Source_map.t option * _) =

diff --git a/compiler/lib/generate.ml b/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 ->

diff --git a/compiler/lib/generate.mli b/compiler/lib/generate.mli
@@ -29,3 +29,9 @@ val f :
   -> Javascript.program
 
 val init : unit -> unit
+
+val source_location :
+     Parse_bytecode.Debug.t
+  -> ?force:Parse_bytecode.Debug.force
+  -> Code.loc
+  -> Javascript.location
diff --git a/compiler/lib/link_js.ml b/compiler/lib/link_js.ml
@@ -469,7 +469,7 @@ let link ~output ~linkall ~mklib ~toplevel ~files ~resolve_sourcemap_url ~source
               let s = sourceMappingURL_base64 ^ Base64.encode_exn data in
               Line_writer.write oc s
           | Some file ->
-              Source_map_io.to_file sm file;
+              Source_map_io.to_file sm ~file;
               let s = sourceMappingURL ^ Filename.basename file in
               Line_writer.write oc s));
       if times () then Format.eprintf "  sourcemap: %a@." Timer.print t

diff --git a/compiler/lib/source_map_io.mli b/compiler/lib/source_map_io.mli
@@ -23,6 +23,13 @@ val enabled : bool
 
 val to_string : t -> string
 
-val to_file : t -> string -> unit
-
 val of_string : string -> t
+
+(** Read source map from a file without parsing the mappings (which can be costly). The
+    [mappings] field is returned empty and the raw string is returned alongside the map.
+ *)
+val of_file_no_mappings : string -> t * string option
+
+(** Write to a file. If a string is supplied as [mappings], use it instead of the
+    sourcemap's [mappings]. *)
+val to_file : ?mappings:string -> t -> file:string -> unit