From 06d0c09c7ef6089db3f19f7199d1b529a690a05d Mon Sep 17 00:00:00 2001
From: Logan Smyth <loganfsmyth@gmail.com>
Date: Mon, 23 Apr 2018 12:10:20 -0700
Subject: [PATCH] Use the scope visitor to avoid rewriting shadowed bindings.

---
 src/workers/parser/getScopes/index.js         |  3 +
 src/workers/parser/getScopes/visitor.js       |  4 +
 src/workers/parser/mapOriginalExpression.js   | 73 +++++++++++++------
 .../tests/mapOriginalExpression.spec.js       | 13 ++++
 4 files changed, 69 insertions(+), 24 deletions(-)

diff --git a/src/workers/parser/getScopes/index.js b/src/workers/parser/getScopes/index.js
index 70bb3d97e6..134fe31260 100644
--- a/src/workers/parser/getScopes/index.js
+++ b/src/workers/parser/getScopes/index.js
@@ -5,6 +5,7 @@
 // @flow
 
 import {
+  buildScopeList,
   parseSourceScopes,
   type SourceScope,
   type ParsedScope,
@@ -42,6 +43,8 @@ export function clearScopes() {
   parsedScopesCache = new Map();
 }
 
+export { buildScopeList };
+
 /**
  * Searches all scopes and their bindings at the specific location.
  */
diff --git a/src/workers/parser/getScopes/visitor.js b/src/workers/parser/getScopes/visitor.js
index 8ccf3476ee..dc3209721b 100644
--- a/src/workers/parser/getScopes/visitor.js
+++ b/src/workers/parser/getScopes/visitor.js
@@ -142,6 +142,10 @@ export function parseSourceScopes(sourceId: SourceId): ?Array<ParsedScope> {
     return null;
   }
 
+  return buildScopeList(ast, sourceId);
+}
+
+export function buildScopeList(ast: Node, sourceId: SourceId) {
   const { global, lexical } = createGlobalScope(ast, sourceId);
 
   const state = {
diff --git a/src/workers/parser/mapOriginalExpression.js b/src/workers/parser/mapOriginalExpression.js
index c5f2fc9f77..34413bb817 100644
--- a/src/workers/parser/mapOriginalExpression.js
+++ b/src/workers/parser/mapOriginalExpression.js
@@ -4,7 +4,9 @@
 
 // @flow
 
+import type { ColumnPosition } from "../../types";
 import { parseScript } from "./utils/ast";
+import { buildScopeList } from "./getScopes";
 import generate from "@babel/generator";
 import * as t from "@babel/types";
 
@@ -28,48 +30,71 @@ function getFirstExpression(ast) {
   return statements[0].expression;
 }
 
+function locationKey(start: ColumnPosition): string {
+  return `${start.line}:${start.column}`;
+}
+
 export default function mapOriginalExpression(
   expression: string,
   mappings: {
     [string]: string | null
   }
 ): string {
-  let didReplace = false;
-
   const ast = parseScript(expression);
-  t.traverse(ast, (node, ancestors) => {
-    const parent = ancestors[ancestors.length - 1];
-    if (!parent) {
-      return;
-    }
+  const scopes = buildScopeList(ast, "");
 
-    const parentNode = parent.node;
+  const nodes = new Map();
 
-    let name = null;
-    if (t.isIdentifier(node) && t.isReferenced(node, parentNode)) {
-      name = node.name;
-    } else if (t.isThisExpression(node)) {
-      name = "this";
-    } else {
-      return;
+  const replacements = new Map();
+
+  // The ref-only global bindings are the ones that are accessed, but not
+  // declared anywhere in the parsed code, meaning they are either global,
+  // or declared somewhere in a scope outside the parsed code, so we
+  // rewrite all of those specifically to avoid rewritting declarations that
+  // shadow outer mappings.
+  for (const name of Object.keys(scopes[0].bindings)) {
+    const { refs } = scopes[0].bindings[name];
+    const mapping = mappings[name];
+    if (
+      !refs.every(ref => ref.type === "ref") ||
+      !mapping ||
+      mapping === name
+    ) {
+      continue;
+    }
+
+    let node = nodes.get(name);
+    if (!node) {
+      node = getFirstExpression(parseScript(mapping));
+      nodes.set(name, node);
     }
 
-    if (mappings.hasOwnProperty(name)) {
-      const mapping = mappings[name];
-      if (mapping && mapping !== name) {
-        const mappingNode = getFirstExpression(parseScript(mapping));
-        replaceNode(ancestors, mappingNode);
+    for (const ref of refs) {
+      let { line, column } = ref.start;
 
-        didReplace = true;
-      }
+      // This shouldn't happen, just keeping Flow happy.
+      if (typeof column !== "number") column = 0;
+
+      replacements.set(locationKey({ line, column}), node);
     }
-  });
+  }
 
-  if (!didReplace) {
+  if (replacements.size === 0) {
     // Avoid the extra code generation work and also avoid potentially
     // reformatting the user's code unnecessarily.
     return expression;
   }
 
+  t.traverse(ast, (node, ancestors) => {
+    if (!t.isIdentifier(node) && !t.isThisExpression(node)) {
+      return;
+    }
+
+    const replacement = replacements.get(locationKey(node.loc.start));
+    if (replacement) {
+      replaceNode(ancestors, t.cloneNode(replacement));
+    }
+  });
+
   return generate(ast).code;
 }
diff --git a/src/workers/parser/tests/mapOriginalExpression.spec.js b/src/workers/parser/tests/mapOriginalExpression.spec.js
index 4fd3a7a81e..9098128cd1 100644
--- a/src/workers/parser/tests/mapOriginalExpression.spec.js
+++ b/src/workers/parser/tests/mapOriginalExpression.spec.js
@@ -44,4 +44,17 @@ describe("mapOriginalExpression", () => {
     });
     expect(generatedExpression).toEqual("a + b");
   });
+
+  it("shadowed bindings", () => {
+    const generatedExpression = mapOriginalExpression(
+      "window.thing = function fn(){ var a; a; b; }; a; b; ",
+      {
+        a: "_a",
+        b: "_b"
+      }
+    );
+    expect(generatedExpression).toEqual(
+      "window.thing = function fn() {\n  var a;\n  a;\n  _b;\n};\n\n_a;\n_b;"
+    );
+  });
 });