From 92952f51d0f6ee8e8004dbbf06b59cf68ca27135 Mon Sep 17 00:00:00 2001
From: Josh Junon <junon@wavetilt.com>
Date: Wed, 9 Dec 2020 15:09:36 +0100
Subject: [PATCH] allow inputSourceMap on .parse() (closes #128, closes #146)

---
 lib/parser.ts   | 15 ++++++--
 lib/printer.ts  |  2 +-
 lib/util.ts     | 10 ++++++
 test/mapping.ts | 91 +++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 115 insertions(+), 3 deletions(-)

diff --git a/lib/parser.ts b/lib/parser.ts
index 60e36776..6ed9cae2 100644
--- a/lib/parser.ts
+++ b/lib/parser.ts
@@ -12,7 +12,8 @@ import { Options } from "./options";
 export function parse(source: string, options?: Partial<Options>) {
   options = normalizeOptions(options);
 
-  const lines = fromString(source, options);
+  const map = options.inputSourceMap ? ({source, inputSourceMap: options.inputSourceMap}) : util.extractInlineSourceMaps(source);
+  const lines = fromString(map.source, options);
 
   const sourceWithoutTabs = lines.toString({
     tabWidth: options.tabWidth,
@@ -120,7 +121,17 @@ export function parse(source: string, options?: Partial<Options>) {
 
   // Return a copy of the original AST so that any changes made may be
   // compared to the original.
-  return new TreeCopier(lines, tokens).copy(file);
+  const copy = new TreeCopier(lines, tokens).copy(file);
+
+  // Attach inline sourcemap metadata when possible
+  Object.defineProperty(copy, "inputSourceMap", {
+    enumerable: false,
+    configurable: false,
+    writable: false,
+    value: map.inputSourceMap
+  });
+
+  return copy;
 }
 
 interface TreeCopierType {
diff --git a/lib/printer.ts b/lib/printer.ts
index ac33d5b9..3a4c3e9f 100644
--- a/lib/printer.ts
+++ b/lib/printer.ts
@@ -143,7 +143,7 @@ const Printer = (function Printer(this: PrinterType, config?: any) {
     return new PrintResult(
       lines.toString(config),
       util.composeSourceMaps(
-        config.inputSourceMap,
+        config.inputSourceMap || ast.inputSourceMap,
         lines.getSourceMap(config.sourceMapName, config.sourceRoot),
       ),
     );
diff --git a/lib/util.ts b/lib/util.ts
index 7bfc8ae5..fa9dee3b 100644
--- a/lib/util.ts
+++ b/lib/util.ts
@@ -348,3 +348,13 @@ export function isTrailingCommaEnabled(options: any, context: any) {
   }
   return !!trailingComma;
 }
+
+export function extractInlineSourceMaps(source: string) {
+  const re = /\s*\/\/(?:@|#) sourceMappingURL=data:application\/json;base64,(\S*)$/m;
+  const map = source.match(re);
+
+  return {
+    source: map ? source.replace(re, '') : source,
+    inputSourceMap: map ? (Buffer.from(map[1], 'base64')).toString() : undefined
+  };
+}
diff --git a/test/mapping.ts b/test/mapping.ts
index ba8f75dd..2fcce669 100644
--- a/test/mapping.ts
+++ b/test/mapping.ts
@@ -169,6 +169,97 @@ describe("source maps", function () {
     });
   });
 
+  it("should extract inline sourcemap from source", function() {
+    function addUseStrict(ast: any) {
+      return recast.visit(ast, {
+        visitFunction: function(path) {
+          path.get("body", "body").unshift(
+            b.expressionStatement(b.literal("use strict"))
+          );
+          this.traverse(path);
+        }
+      });
+    }
+
+    function stripConsole(ast: any) {
+      return recast.visit(ast, {
+        visitCallExpression: function(path) {
+          const node = path.value;
+          if (n.MemberExpression.check(node.callee) &&
+            n.Identifier.check(node.callee.object) &&
+            node.callee.object.name === "console") {
+            n.ExpressionStatement.assert(path.parent.node);
+            path.parent.replace();
+            return false;
+          }
+
+          return;
+        }
+      });
+    }
+
+    const code = [
+      "function add(a, b) {",
+      "  var sum = a + b;",
+      "  console.log(a, b);",
+      "  return sum;",
+      "}"
+    ].join("\n");
+
+    const ast = parse(code, {
+      sourceFileName: "original.js"
+    });
+
+    const useStrictResult = new Printer({
+      sourceMapName: "useStrict.map.json"
+    }).print(addUseStrict(ast));
+
+    const useStrictResultCodeAndMap = [
+        useStrictResult.code,
+        "//# sourceMappingURL=data:application/json;base64,",
+        Buffer.from(JSON.stringify(useStrictResult.map)).toString('base64')
+      ].join('');
+
+    const useStrictAst = parse(useStrictResultCodeAndMap, {
+      sourceFileName: "useStrict.js"
+    });
+
+    const oneStepResult = new Printer({
+      sourceMapName: "oneStep.map.json"
+    }).print(stripConsole(ast));
+
+    const twoStepResult = new Printer({
+      sourceMapName: "twoStep.map.json"
+    }).print(stripConsole(useStrictAst));
+
+    assert.strictEqual(
+      oneStepResult.code,
+      twoStepResult.code
+    );
+
+    const smc1 = new sourceMap.SourceMapConsumer(oneStepResult.map);
+    const smc2 = new sourceMap.SourceMapConsumer(twoStepResult.map);
+
+    smc1.eachMapping(function(mapping) {
+      const pos = {
+        line: mapping.generatedLine,
+        column: mapping.generatedColumn
+      };
+
+      const orig1 = smc1.originalPositionFor(pos);
+      const orig2 = smc2.originalPositionFor(pos);
+
+      // The composition of the source maps generated separately from
+      // the two transforms should be equivalent to the source map
+      // generated from the composition of the two transforms.
+      assert.deepEqual(orig1, orig2);
+
+      // Make sure the two-step source map refers back to the original
+      // source instead of the intermediate source.
+      assert.strictEqual(orig2.source, "original.js");
+    });
+  });
+
   it("should work when a child node becomes null", function () {
     // https://github.com/facebook/regenerator/issues/103
     const code = ["for (var i = 0; false; i++)", "  log(i);"].join(eol);