Source map related code clean-up

vouillon · vouillon · commit b5ae6af4f9de · 2024-11-07T16:15:00.000+01:00
diff --git a/compiler/bin-wasm_of_ocaml/compile.ml b/compiler/bin-wasm_of_ocaml/compile.ml
@@ -234,13 +234,20 @@ let build_js_runtime ~primitives ?runtime_arguments () =
   in
   prelude ^ launcher
 
-let add_source_map sourcemap_don't_inline_content z opt_source_map_file =
-  Option.iter
-    ~f:(fun file ->
-      Zip.add_file z ~name:"source_map.map" ~file;
+let add_source_map sourcemap_don't_inline_content z opt_source_map =
+  let sm =
+    match opt_source_map with
+    | `File opt_file ->
+        Option.map opt_file ~f:(fun file ->
+            Zip.add_file z ~name:"source_map.map" ~file;
+            Source_map.of_file file)
+    | `Source_map sm ->
+        Zip.add_entry z ~name:"source_map.map" ~contents:(Source_map.to_string sm);
+        Some sm
+  in
+  Option.iter sm ~f:(fun sm ->
       if not sourcemap_don't_inline_content
       then
-        let sm = Source_map.of_file file in
         Wasm_source_map.iter_sources sm (fun i j file ->
             if Sys.file_exists file && not (Sys.is_directory file)
             then
@@ -249,7 +256,6 @@ let add_source_map sourcemap_don't_inline_content z opt_source_map_file =
                 z
                 ~name:(Link.source_name i j file)
                 ~contents:(Yojson.Basic.to_string (`String sm))))
-    opt_source_map_file
 
 let run
     { Cmd_arg.common
@@ -488,7 +494,7 @@ let run
          let compile_cmo' z cmo =
            compile_cmo cmo (fun unit_data _ tmp_wasm_file opt_tmp_map_file ->
                Zip.add_file z ~name:"code.wasm" ~file:tmp_wasm_file;
-               add_source_map sourcemap_don't_inline_content z opt_tmp_map_file;
+               add_source_map sourcemap_don't_inline_content z (`File opt_tmp_map_file);
                unit_data)
          in
          let unit_data = [ compile_cmo' z cmo ] in
@@ -499,6 +505,7 @@ let run
          @@ fun tmp_output_file ->
          let z = Zip.open_out tmp_output_file in
          let unit_data =
+           let tmp_buf = Buffer.create 10000 in
            List.fold_right
              ~f:(fun cmo cont l ->
                compile_cmo cmo
@@ -508,26 +515,26 @@ let run
              ~init:(fun l ->
                Fs.with_intermediate_file (Filename.temp_file "wasm" ".wasm")
                @@ fun tmp_wasm_file ->
-               opt_with
-                 Fs.with_intermediate_file
-                 (if enable_source_maps
-                  then Some (Filename.temp_file "wasm" ".map")
-                  else None)
-               @@ fun opt_output_sourcemap_file ->
                let l = List.rev l in
-               Wasm_link.f
-                 (List.map
-                    ~f:(fun (_, _, file, opt_source_map) ->
-                      { Wasm_link.module_name = "OCaml"
-                      ; file
-                      ; code = None
-                      ; opt_source_map = Option.map ~f:(fun f -> `File f) opt_source_map
-                      })
-                    l)
-                 ~output_file:tmp_wasm_file
-                 ~opt_output_sourcemap_file;
+               let source_map =
+                 Wasm_link.f
+                   (List.map
+                      ~f:(fun (_, _, file, opt_source_map) ->
+                        { Wasm_link.module_name = "OCaml"
+                        ; file
+                        ; code = None
+                        ; opt_source_map =
+                            Option.map
+                              ~f:(fun f -> Source_map.Standard.of_file ~tmp_buf f)
+                              opt_source_map
+                        })
+                      l)
+                   ~output_file:tmp_wasm_file
+               in
                Zip.add_file z ~name:"code.wasm" ~file:tmp_wasm_file;
-               add_source_map sourcemap_don't_inline_content z opt_output_sourcemap_file;
+               if enable_source_maps
+               then
+                 add_source_map sourcemap_don't_inline_content z (`Source_map source_map);
                List.map ~f:(fun (unit_data, _, _, _) -> unit_data) l)
              []
          in
diff --git a/compiler/lib-wasm/link.ml b/compiler/lib-wasm/link.ml
@@ -849,63 +849,51 @@ let link ~output_file ~linkall ~enable_source_maps ~files =
   if times () then Format.eprintf "    build JS runtime: %a@." Timer.print t1;
   if times () then Format.eprintf "  emit: %a@." Timer.print t
 
-let opt_with action x f =
-  match x with
-  | None -> f None
-  | Some x -> action x (fun y -> f (Some y))
-
-let rec get_source_map_files files src_index =
+let rec get_source_map_files ~tmp_buf files src_index =
   let z = Zip.open_in files.(!src_index) in
   incr src_index;
-  if Zip.has_entry z ~name:"source_map.map"
-  then
-    let data = Zip.read_entry z ~name:"source_map.map" in
-    let sm = Source_map.Standard.of_string data in
-    if not (Wasm_source_map.is_empty sm)
-    then (
-      let l = ref [] in
-      Wasm_source_map.iter_sources (Standard sm) (fun i j file ->
-          l := source_name i j file :: !l);
-      if not (List.is_empty !l)
-      then z, Array.of_list (List.rev !l)
-      else (
-        Zip.close_in z;
-        get_source_map_files files src_index))
-    else (
-      Zip.close_in z;
-      get_source_map_files files src_index)
-  else get_source_map_files files src_index
-
-let add_source_map files z opt_source_map_file =
-  Option.iter
-    ~f:(fun file ->
-      Zip.add_file z ~name:"source_map.map" ~file;
-      let sm = Source_map.of_file file in
-      let files = Array.of_list files in
-      let src_index = ref 0 in
-      let st = ref None in
-      let finalize () =
+  let l = ref [] in
+  (if Zip.has_entry z ~name:"source_map.map"
+   then
+     let data = Zip.read_entry z ~name:"source_map.map" in
+     let sm = Source_map.Standard.of_string ~tmp_buf data in
+     if not (Wasm_source_map.is_empty sm)
+     then
+       Wasm_source_map.iter_sources (Standard sm) (fun i j file ->
+           l := source_name i j file :: !l));
+  if not (List.is_empty !l)
+  then z, Array.of_list (List.rev !l)
+  else (
+    Zip.close_in z;
+    get_source_map_files ~tmp_buf files src_index)
+
+let add_source_map files z sm =
+  let tmp_buf = Buffer.create 10000 in
+  Zip.add_entry z ~name:"source_map.map" ~contents:(Source_map.to_string sm);
+  let files = Array.of_list files in
+  let src_index = ref 0 in
+  let st = ref None in
+  let finalize () =
+    match !st with
+    | Some (_, (z', _)) -> Zip.close_in z'
+    | None -> ()
+  in
+  Wasm_source_map.iter_sources sm (fun i j file ->
+      let z', files =
         match !st with
-        | Some (_, (z', _)) -> Zip.close_in z'
-        | None -> ()
+        | Some (i', st) when Poly.equal i i' -> st
+        | _ ->
+            let st' = get_source_map_files ~tmp_buf files src_index in
+            finalize ();
+            st := Some (i, st');
+            st'
       in
-      Wasm_source_map.iter_sources sm (fun i j file ->
-          let z', files =
-            match !st with
-            | Some (i', st) when Poly.equal i i' -> st
-            | _ ->
-                let st' = get_source_map_files files src_index in
-                finalize ();
-                st := Some (i, st');
-                st'
-          in
-          if Array.length files > 0 (* Source has source map *)
-          then
-            let name = files.(Option.value ~default:0 j) in
-            if Zip.has_entry z' ~name
-            then Zip.copy_file z' z ~src_name:name ~dst_name:(source_name i j file));
-      finalize ())
-    opt_source_map_file
+      if Array.length files > 0 (* Source has source map *)
+      then
+        let name = files.(Option.value ~default:0 j) in
+        if Zip.has_entry z' ~name
+        then Zip.copy_file z' z ~src_name:name ~dst_name:(source_name i j file));
+  finalize ()
 
 let make_library ~output_file ~enable_source_maps ~files =
   let info =
@@ -936,33 +924,31 @@ let make_library ~output_file ~enable_source_maps ~files =
   @@ fun tmp_output_file ->
   let z = Zip.open_out tmp_output_file in
   add_info z ~build_info ~unit_data ();
-  (*
-- Merge all code files into a single code file (gathering source maps)
-- Copy source files
-*)
   Fs.with_intermediate_file (Filename.temp_file "wasm" ".wasm")
   @@ fun tmp_wasm_file ->
-  opt_with
-    Fs.with_intermediate_file
-    (if enable_source_maps then Some (Filename.temp_file "wasm" ".map") else None)
-  @@ fun opt_output_sourcemap_file ->
-  Wasm_link.f
-    (List.map
-       ~f:(fun file ->
-         let z' = Zip.open_in file in
-         { Wasm_link.module_name = "OCaml"
-         ; file
-         ; code = Some (Zip.read_entry z' ~name:"code.wasm")
-         ; opt_source_map =
-             (if Zip.has_entry z' ~name:"source_map.map"
-              then Some (`Data (Zip.read_entry z' ~name:"source_map.map"))
-              else None)
-         })
-       files)
-    ~output_file:tmp_wasm_file
-    ~opt_output_sourcemap_file;
+  let output_sourcemap =
+    Wasm_link.f
+      (let tmp_buf = Buffer.create 10000 in
+       List.map
+         ~f:(fun file ->
+           let z' = Zip.open_in file in
+           { Wasm_link.module_name = "OCaml"
+           ; file
+           ; code = Some (Zip.read_entry z' ~name:"code.wasm")
+           ; opt_source_map =
+               (if enable_source_maps && Zip.has_entry z' ~name:"source_map.map"
+                then
+                  Some
+                    (Source_map.Standard.of_string
+                       ~tmp_buf
+                       (Zip.read_entry z' ~name:"source_map.map"))
+                else None)
+           })
+         files)
+      ~output_file:tmp_wasm_file
+  in
   Zip.add_file z ~name:"code.wasm" ~file:tmp_wasm_file;
-  add_source_map files z opt_output_sourcemap_file;
+  if enable_source_maps then add_source_map files z output_sourcemap;
   Zip.close_out z
 
 let link ~output_file ~linkall ~mklib ~enable_source_maps ~files =
diff --git a/compiler/lib-wasm/wasm_link.ml b/compiler/lib-wasm/wasm_link.ml
@@ -1856,12 +1856,11 @@ type input =
   { module_name : string
   ; file : string
   ; code : string option
-  ; opt_source_map : [ `File of string | `Data of string ] option
+  ; opt_source_map : Source_map.Standard.t option
   }
 
-let f files ~output_file ~opt_output_sourcemap_file =
+let f files ~output_file =
   let files =
-    let tmp_buf = Buffer.create 10000 in
     Array.map
       ~f:(fun { module_name; file; code; opt_source_map } ->
         let data =
@@ -1870,17 +1869,7 @@ let f files ~output_file ~opt_output_sourcemap_file =
           | Some data -> data
         in
         let contents = Read.open_in file data in
-        { module_name
-        ; file
-        ; contents
-        ; source_map_contents =
-            Option.map
-              ~f:(fun src ->
-                match src with
-                | `File file -> Source_map.Standard.of_file ~tmp_buf file
-                | `Data data -> Source_map.Standard.of_string ~tmp_buf data)
-              opt_source_map
-        })
+        { module_name; file; contents; source_map_contents = opt_source_map })
       (Array.of_list files)
   in
 
@@ -2274,15 +2263,12 @@ let f files ~output_file ~opt_output_sourcemap_file =
     pos_out out_ch + 1 + Buffer.length b
   in
   add_section out_ch ~id:10 code_pieces;
-  Option.iter
-    ~f:(fun file ->
-      Source_map.to_file
-        (Wasm_source_map.concatenate
-           (List.map
-              ~f:(fun (pos, sm) -> pos + code_section_offset, sm)
-              (List.rev !source_maps)))
-        file)
-    opt_output_sourcemap_file;
+  let source_map =
+    Wasm_source_map.concatenate
+      (List.map
+         ~f:(fun (pos, sm) -> pos + code_section_offset, sm)
+         (List.rev !source_maps))
+  in
 
   (* 11: data *)
   ignore
@@ -2429,7 +2415,9 @@ let f files ~output_file ~opt_output_sourcemap_file =
 
   add_section out_ch ~id:0 name_section_buffer;
 
-  close_out out_ch
+  close_out out_ch;
+
+  source_map
 
 (*
 LATER
diff --git a/compiler/lib-wasm/wasm_link.mli b/compiler/lib-wasm/wasm_link.mli
@@ -2,8 +2,7 @@ type input =
   { module_name : string
   ; file : string
   ; code : string option
-  ; opt_source_map : [ `File of string | `Data of string ] option
+  ; opt_source_map : Source_map.Standard.t option
   }
 
-val f :
-  input list -> output_file:string -> opt_output_sourcemap_file:string option -> unit
+val f : input list -> output_file:string -> Source_map.t
