Preserve source map markers across package targets

Author: mununki

compiler/core/js_implementation.ml

@@ -141,17 +141,18 @@ let after_parsing_impl ppf outputprefix (ast : Parsetree.structure) =
       let typedtree_coercion = (typedtree, coercion) in
       print_if ppf Clflags.dump_typedtree
         Printtyped.implementation_with_coercion typedtree_coercion;
-      (if !Js_config.cmi_only then Warnings.check_fatal ()
-       else
-         let lambda, exports =
-           Translmod.transl_implementation modulename typedtree_coercion
-         in
-         let js_program =
-           print_if_pipe ppf Clflags.dump_rawlambda Printlambda.lambda lambda
-           |> Lam_compile_main.compile outputprefix exports
-         in
-         if not !Js_config.cmj_only then
-           Lam_compile_main.lambda_as_module js_program outputprefix);
+      Js_source_map.with_marker_scope (fun () ->
+          if !Js_config.cmi_only then Warnings.check_fatal ()
+          else
+            let lambda, exports =
+              Translmod.transl_implementation modulename typedtree_coercion
+            in
+            let js_program =
+              print_if_pipe ppf Clflags.dump_rawlambda Printlambda.lambda lambda
+              |> Lam_compile_main.compile outputprefix exports
+            in
+            if not !Js_config.cmj_only then
+              Lam_compile_main.lambda_as_module js_program outputprefix);
       process_with_gentype (outputprefix ^ ".cmt"))
 
 let implementation ~parser ppf ?outputprefix fname =

compiler/core/js_source_map.ml

@@ -46,6 +46,15 @@ let comment_of_loc (loc : Location.t) =
       Hashtbl.replace marker_locs id loc;
       Some (marker_prefix ^ string_of_int id)
 
+let with_marker_scope f =
+  let first_marker = !next_marker in
+  Ext_pervasives.finally ()
+    ~clean:(fun () ->
+      for id = first_marker to !next_marker - 1 do
+        Hashtbl.remove marker_locs id
+      done)
+    f
+
 let with_builder builder f =
   let old = !current in
   current := builder;
@@ -182,21 +191,14 @@ let add_mapping builder ~generated_line ~generated_column (loc : Location.t) =
         :: builder.mappings;
       builder.last_generated <- Some (generated_line, generated_column)
 
-let take_marker_loc id =
-  match Hashtbl.find_opt marker_locs id with
-  | None -> None
-  | Some loc ->
-    Hashtbl.remove marker_locs id;
-    Some loc
-
 let mark_comment fmt comment =
   if is_prefix ~prefix:marker_prefix comment then (
     let prefix_len = String.length marker_prefix in
     let id =
       int_of_string
         (String.sub comment prefix_len (String.length comment - prefix_len))
     in
-    (match (!current, take_marker_loc id) with
+    (match (!current, Hashtbl.find_opt marker_locs id) with
     | Some builder, Some loc ->
       let generated_line, generated_column = Ext_pp.position fmt in
       add_mapping builder ~generated_line ~generated_column loc

compiler/core/js_source_map.mli

@@ -3,6 +3,8 @@ type t
 val make :
   generated_file:string -> source_root:string -> sources_content:bool -> t
 
+val with_marker_scope : (unit -> 'a) -> 'a
+
 val with_builder : t option -> (unit -> 'a) -> 'a
 
 val comment_of_loc : Location.t -> string option

tests/build_tests/source_map/input.js

@@ -9,23 +9,28 @@ const { execBuildOrThrow, execClean } = setup(import.meta.dirname);
 
 await execBuildOrThrow();
 
-const jsPath = path.join(import.meta.dirname, "lib", "bs", "src", "Demo.js");
-const mapPath = `${jsPath}.map`;
+for (const filename of ["Demo.cjs", "Demo.mjs"]) {
+  const jsPath = path.join(import.meta.dirname, "lib", "bs", "src", filename);
+  const mapPath = `${jsPath}.map`;
 
-const js = await fs.readFile(jsPath, "utf8");
-assert.match(js, /\/\/# sourceMappingURL=Demo\.js\.map/);
+  const js = await fs.readFile(jsPath, "utf8");
+  assert.match(
+    js,
+    new RegExp(`//# sourceMappingURL=${filename.replace(".", "\\.")}\\.map`),
+  );
 
-const map = JSON.parse(await fs.readFile(mapPath, "utf8"));
-assert.equal(map.version, 3);
-assert.equal(map.file, "Demo.js");
-assert.ok(map.mappings.length > 0, "source map should include mappings");
-assert.ok(
-  map.sources.some(source => source.endsWith("Demo.res")),
-  `source map should include Demo.res, got ${map.sources.join(", ")}`,
-);
-assert.ok(
-  map.sourcesContent.some(content => content.includes("let add = (a, b)")),
-  "source map should include source contents",
-);
+  const map = JSON.parse(await fs.readFile(mapPath, "utf8"));
+  assert.equal(map.version, 3);
+  assert.equal(map.file, filename);
+  assert.ok(map.mappings.length > 0, `${filename}.map should include mappings`);
+  assert.ok(
+    map.sources.some(source => source.endsWith("Demo.res")),
+    `${filename}.map should include Demo.res, got ${map.sources.join(", ")}`,
+  );
+  assert.ok(
+    map.sourcesContent.some(content => content.includes("let add = (a, b)")),
+    `${filename}.map should include source contents`,
+  );
+}
 
 await execClean();

tests/build_tests/source_map/rescript.json

@@ -1,6 +1,18 @@
 {
   "name": "source_map",
   "sources": ["src"],
+  "package-specs": [
+    {
+      "module": "commonjs",
+      "in-source": true,
+      "suffix": ".cjs"
+    },
+    {
+      "module": "esmodule",
+      "in-source": true,
+      "suffix": ".mjs"
+    }
+  ],
   "sourceMap": {
     "enabled": "always",
     "mode": "linked",