Optimize sourcemap processing

Author: OlivierNicole

compiler/bin-js_of_ocaml/cmd_arg.ml

@@ -318,9 +318,11 @@ let options =
             ; file
             ; sourceroot = sourcemap_root
             ; sources = []
-            ; sources_content = (if sourcemap_don't_inline_content then None else Some [])
+            ; sources_contents =
+                (if sourcemap_don't_inline_content then None
+                else Some (Source_map.Sources_contents.encode []))
             ; names = []
-            ; mappings = []
+            ; mappings = Source_map.Mappings.empty
             } )
       else None
     in
@@ -557,9 +559,11 @@ let options_runtime_only =
             ; file
             ; sourceroot = sourcemap_root
             ; sources = []
-            ; sources_content = (if sourcemap_don't_inline_content then None else Some [])
+            ; sources_contents =
+                (if sourcemap_don't_inline_content then None
+                else Some (Source_map.Sources_contents.encode []))
             ; names = []
-            ; mappings = []
+            ; mappings = Source_map.Mappings.empty
             } )
       else None
     in

compiler/bin-js_of_ocaml/link.ml

@@ -106,9 +106,9 @@ let options =
             ; file
             ; sourceroot = sourcemap_root
             ; sources = []
-            ; sources_content = Some []
+            ; sources_contents = Some (Source_map.Sources_contents.encode [])
             ; names = []
-            ; mappings = []
+            ; mappings = Source_map.Mappings.empty
             } )
       else None
     in

compiler/lib/js_output.ml

@@ -1905,8 +1905,8 @@ let program ?(accept_unnamed_var = false) f ?source_map p =
   let names = Hashtbl.create 17 in
   let contents : string option list ref option =
     match source_map with
-    | None | Some { Source_map.sources_content = None; _ } -> None
-    | Some { Source_map.sources_content = Some _; _ } -> Some (ref [])
+    | None | Some { Source_map.sources_contents = None; _ } -> None
+    | Some { Source_map.sources_contents = Some _; _ } -> Some (ref [])
   in
   let push_mapping, get_file_index, get_name_index, source_map_enabled =
     let source_map_enabled =
@@ -1925,7 +1925,10 @@ let program ?(accept_unnamed_var = false) f ?source_map p =
                 Option.iter contents ~f:(fun r -> r := y :: !r);
                 loop xs ys
           in
-          loop sm.sources (Option.value ~default:[] sm.sources_content);
+          let sources_contents =
+            Option.map ~f:Source_map.Sources_contents.decode sm.sources_contents
+          in
+          loop sm.sources (Option.value ~default:[] sources_contents);
           List.iter sm.Source_map.names ~f:(fun f ->
               Hashtbl.add names f (Hashtbl.length names));
           true
@@ -1979,19 +1982,20 @@ let program ?(accept_unnamed_var = false) f ?source_map p =
     | Some sm ->
         let sources = hashtbl_to_list files in
         let names = hashtbl_to_list names in
-        let sources_content =
-          match contents with
-          | None -> None
-          | Some r -> Some (List.rev !r)
+        let sources_contents =
+          let open Option.Syntax in
+          let* r = contents in
+          Option.return (Source_map.Sources_contents.encode (List.rev !r))
         in
         let sources =
           List.map sources ~f:(fun filename ->
               match Builtins.find filename with
               | None -> filename
               | Some _ -> Filename.concat "/builtin" filename)
         in
+        let sm_mappings = Source_map.Mappings.decode sm.mappings in
         let mappings =
-          List.rev_append_map !temp_mappings sm.mappings ~f:(fun (pos, m) ->
+          List.rev_append_map !temp_mappings sm_mappings ~f:(fun (pos, m) ->
               let gen_line = pos.PP.p_line + 1 in
               let gen_col = pos.PP.p_col in
               match m with
@@ -2006,7 +2010,8 @@ let program ?(accept_unnamed_var = false) f ?source_map p =
                   Source_map.Gen_Ori_Name
                     { gen_line; gen_col; ori_source; ori_line; ori_col; ori_name })
         in
-        Some { sm with Source_map.sources; names; sources_content; mappings }
+        let mappings = Source_map.Mappings.encode mappings in
+        Some { sm with Source_map.sources; names; sources_contents; mappings }
   in
   PP.check f;
   (if stats ()

compiler/lib/link_js.ml

@@ -99,9 +99,9 @@ module Line_writer : sig
 
   val of_channel : out_channel -> t
 
-  val write : ?source:Line_reader.t -> t -> string -> unit
+  val write : ?source:Line_reader.t -> t -> string -> int
 
-  val write_lines : ?source:Line_reader.t -> t -> string -> unit
+  val write_lines : ?source:Line_reader.t -> t -> string -> int
 
   val lnum : t -> int
 end = struct
@@ -134,17 +134,20 @@ end = struct
     output_string t.oc "\n";
     let lnum_off = lnum_off + 1 in
     t.source <- source;
-    t.lnum <- t.lnum + lnum_off
+    t.lnum <- t.lnum + lnum_off;
+    lnum_off
 
   let write_lines ?source t lines =
+    let lnum = t.lnum in
     let l = String.split_on_char ~sep:'\n' lines in
     let rec w = function
       | [ "" ] | [] -> ()
       | s :: xs ->
-          write ?source t s;
+          let _ = write ?source t s in
           w xs
     in
-    w l
+    w l;
+    t.lnum - lnum
 
   let lnum t = t.lnum
 end
@@ -318,12 +321,26 @@ let link ~output ~linkall ~mklib ~toplevel ~files ~resolve_sourcemap_url ~source
       in
       let sm_for_file = ref None in
       let ic = Line_reader.open_ file in
-      let skip ic = Line_reader.drop ic in
-      let reloc = ref [] in
+      let old_line_count = Line_writer.lnum oc in
+      let edits = ref [] in
+      let skip ic =
+        edits := Source_map.Line_edits.Drop :: !edits;
+        Line_reader.drop ic
+      in
       let copy ic oc =
         let line = Line_reader.next ic in
-        Line_writer.write ~source:ic oc line;
-        reloc := (Line_reader.lnum ic, Line_writer.lnum oc) :: !reloc
+        let count = Line_writer.write ~source:ic oc line in
+        if count > 1
+        then edits := Source_map.Line_edits.Add { count = count - 1 } :: !edits;
+        edits := Source_map.Line_edits.Keep :: !edits
+      in
+      let write_line oc str =
+        let count = Line_writer.write oc str in
+        edits := Source_map.Line_edits.(Add { count }) :: !edits
+      in
+      let write_lines oc str =
+        let count = Line_writer.write_lines oc str in
+        edits := Source_map.Line_edits.(Add { count }) :: !edits
       in
       let rec read () =
         match Line_reader.peek ic with
@@ -342,7 +359,7 @@ let link ~output ~linkall ~mklib ~toplevel ~files ~resolve_sourcemap_url ~source
                 if not !build_info_emitted
                 then (
                   let bi = Build_info.with_kind bi (if mklib then `Cma else `Unknown) in
-                  Line_writer.write_lines oc (Build_info.to_string bi);
+                  write_lines oc (Build_info.to_string bi);
                   build_info_emitted := true)
             | Drop -> skip ic
             | Unit ->
@@ -358,7 +375,7 @@ let link ~output ~linkall ~mklib ~toplevel ~files ~resolve_sourcemap_url ~source
                   (if mklib
                    then
                      let u = if linkall then { u with force_link = true } else u in
-                     Line_writer.write_lines oc (Unit_info.to_string u));
+                     write_lines oc (Unit_info.to_string u));
                   let size = ref 0 in
                   while
                     match Line_reader.peek ic with
@@ -402,7 +419,7 @@ let link ~output ~linkall ~mklib ~toplevel ~files ~resolve_sourcemap_url ~source
             read ()
       in
       read ();
-      Line_writer.write oc "";
+      write_line oc "";
       Line_reader.close ic;
       (match is_runtime with
       | None -> ()
@@ -424,10 +441,10 @@ let link ~output ~linkall ~mklib ~toplevel ~files ~resolve_sourcemap_url ~source
             (Parse_bytecode.Debug.create ~include_cmis:false false)
             code;
           let content = Buffer.contents b in
-          Line_writer.write_lines oc content);
+          write_lines oc content);
       (match !sm_for_file with
       | None -> ()
-      | Some x -> sm := (x, !reloc) :: !sm);
+      | Some x -> sm := (x, List.rev !edits, Line_writer.lnum oc - old_line_count) :: !sm);
       match !build_info, build_info_for_file with
       | None, None -> ()
       | Some _, None -> ()
@@ -440,32 +457,55 @@ let link ~output ~linkall ~mklib ~toplevel ~files ~resolve_sourcemap_url ~source
   match source_map with
   | None -> ()
   | Some (file, init_sm) ->
-      let sm =
-        List.rev_map !sm ~f:(fun (sm, reloc) ->
-            let tbl = Hashtbl.create 17 in
-            List.iter reloc ~f:(fun (a, b) -> Hashtbl.add tbl a b);
-            Source_map.filter_map sm ~f:(Hashtbl.find_opt tbl))
+      let sourcemaps_and_line_counts =
+        List.rev_map !sm ~f:(fun (sm, edits, lcount) ->
+            (*
+            (match file with
+            | Some f -> Format.eprintf "file = %s@," f
+            | None -> ());
+            Format.eprintf "@[<v 2>edits =@ %a@,@]" Source_map.Line_edits.pp edits;
+            *)
+            let mappings = sm.Source_map.mappings in
+            (*
+            Format.eprintf "mappings = %s@," sm.Source_map.mappings;
+            *)
+            let mappings = Source_map.Mappings.edit ~strict:false mappings edits in
+            { sm with mappings }, lcount)
       in
-      (match Source_map.merge (init_sm :: sm) with
-      | None -> ()
-      | Some sm -> (
-          (* preserve some info from [init_sm] *)
-          let sm =
-            { sm with
-              version = init_sm.version
-            ; file = init_sm.file
-            ; sourceroot = init_sm.sourceroot
-            }
-          in
-          match file with
-          | None ->
-              let data = Source_map_io.to_string sm in
-              let s = sourceMappingURL_base64 ^ Base64.encode_exn data in
-              Line_writer.write oc s
-          | Some file ->
-              Source_map_io.to_file sm file;
-              let s = sourceMappingURL ^ Filename.basename file in
-              Line_writer.write oc s));
+      let merged_sourcemap =
+        let open Source_map in
+        assert (match init_sm.mappings with Uninterpreted "" -> true | _ -> false);
+        { version = init_sm.version
+        ; file = init_sm.file
+        ; Composite.sections =
+            (let _, sections =
+               List.fold_right
+                 sourcemaps_and_line_counts
+                 ~f:(fun (sm, generated_line_count) (cur_ofs, sections) ->
+                   let offset = Composite.{ gen_line = cur_ofs; gen_column = 0 } in
+                   cur_ofs + generated_line_count, (offset, `Map sm) :: sections)
+                 ~init:(0, [])
+             in
+             List.rev sections)
+        }
+      in
+      (* preserve some info from [init_sm] *)
+      let merged_sourcemap =
+        { merged_sourcemap with
+          sections =
+            List.map merged_sourcemap.sections ~f:(fun (ofs, `Map sm) ->
+                ofs, `Map { sm with sourceroot = init_sm.sourceroot })
+        }
+      in
+      (match file with
+      | None ->
+          let data = Source_map_io.Composite.to_string merged_sourcemap in
+          let s = sourceMappingURL_base64 ^ Base64.encode_exn data in
+          Line_writer.write oc s |> ignore
+      | Some file ->
+          Source_map_io.Composite.to_file merged_sourcemap file;
+          let s = sourceMappingURL ^ Filename.basename file in
+          Line_writer.write oc s |> ignore);
       if times () then Format.eprintf "  sourcemap: %a@." Timer.print t
 
 let link ~output ~linkall ~mklib ~toplevel ~files ~resolve_sourcemap_url ~source_map =