From 6a9a9d7c7092b21dbaa2272e79fa3b88dac83c97 Mon Sep 17 00:00:00 2001
From: Chris Garrett <cgarrett@linkedin.com>
Date: Mon, 22 Feb 2021 11:10:22 -0800
Subject: [PATCH] Make ensureImport lazy

Makes the logic for ensuring the Ember import lazy. This allows it to
handle cases where a different transform has also added the Ember import
already, and so it can use it lazily. Also makes the logic generalized
(in case we ever need to import from a different module in the future).
---
 __tests__/index-test.js |  2 +-
 src/index.js            | 82 ++++++++++++++++++++++++++---------------
 2 files changed, 53 insertions(+), 31 deletions(-)

diff --git a/__tests__/index-test.js b/__tests__/index-test.js
index 9abe526..497b3e2 100644
--- a/__tests__/index-test.js
+++ b/__tests__/index-test.js
@@ -253,7 +253,7 @@ describe('options', () => {
     it(`adds the ember import when used in sub-modules`, () => {
       let input = `import Component from '@ember/component';export default class extends Component {}`;
       let actual = transform(input, [[Plugin, { useEmberModule: true }]]);
-      let expected = `import _Ember from 'ember';\nexport default class extends _Ember.Component {}`;
+      let expected = `import _ember from 'ember';\nexport default class extends _ember.Component {}`;
 
       expect(actual).toEqual(expected);
     });
diff --git a/src/index.js b/src/index.js
index 6e4a452..3160327 100644
--- a/src/index.js
+++ b/src/index.js
@@ -90,39 +90,64 @@ module.exports = function (babel) {
       Program(path, state) {
         let options = state.opts || {};
         let useEmberModule = Boolean(options.useEmberModule);
+        let allAddedImports = {};
 
-        let preexistingEmberImportDeclaration = path
-          .get('body')
-          .filter((n) => n.type === 'ImportDeclaration')
-          .find((n) => n.get('source').get('value').node === 'ember');
-
-        if (
-          // an import was found
-          preexistingEmberImportDeclaration &&
-          // this accounts for `import from 'ember'` without a local identifier
-          preexistingEmberImportDeclaration.node.specifiers.length > 0
-        ) {
-          state.emberIdentifier =
-            preexistingEmberImportDeclaration.node.specifiers[0].local;
-        }
+        state.ensureImport = (exportName, moduleName) => {
+          let addedImports = (allAddedImports[moduleName] =
+            allAddedImports[moduleName] || {});
 
-        state.ensureEmberImport = () => {
-          if (!useEmberModule) {
-            // ensures that we can always assume `state.emberIdentifier` is set
-            state.emberIdentifier = t.identifier('Ember');
-            return;
-          }
+          if (addedImports[exportName]) return addedImports[exportName];
 
-          if (state.emberIdentifier) return;
+          if (
+            exportName === 'default' &&
+            moduleName === 'ember' &&
+            !useEmberModule
+          ) {
+            addedImports[exportName] = t.identifier('Ember');
+            return addedImports[exportName];
+          }
 
-          state.emberIdentifier = path.scope.generateUidIdentifier('Ember');
+          let importDeclarations = path
+            .get('body')
+            .filter((n) => n.type === 'ImportDeclaration');
 
-          let emberImport = t.importDeclaration(
-            [t.importDefaultSpecifier(state.emberIdentifier)],
-            t.stringLiteral('ember')
+          let preexistingImportDeclaration = importDeclarations.find(
+            (n) => n.get('source').get('value').node === moduleName
           );
 
-          path.unshiftContainer('body', emberImport);
+          if (preexistingImportDeclaration) {
+            let importSpecifier = preexistingImportDeclaration
+              .get('specifiers')
+              .find(({ node }) => {
+                return exportName === 'default'
+                  ? t.isImportDefaultSpecifier(node)
+                  : node.imported.name === exportName;
+              });
+
+            if (importSpecifier) {
+              addedImports[exportName] = importSpecifier.node.local;
+            }
+          }
+
+          if (!addedImports[exportName]) {
+            let uid = path.scope.generateUidIdentifier(
+              exportName === 'default' ? moduleName : exportName
+            );
+            addedImports[exportName] = uid;
+
+            let newImportSpecifier =
+              exportName === 'default'
+                ? t.importDefaultSpecifier(uid)
+                : t.importSpecifier(uid, t.identifier(exportName));
+
+            let newImport = t.importDeclaration(
+              [newImportSpecifier],
+              t.stringLiteral(moduleName)
+            );
+            path.unshiftContainer('body', newImport);
+          }
+
+          return addedImports[exportName];
         };
       },
 
@@ -216,9 +241,6 @@ module.exports = function (babel) {
 
             removals.push(specifierPath);
 
-            // ensure that the Ember global is imported if needed
-            state.ensureEmberImport();
-
             if (
               path.scope.bindings[local.name].referencePaths.find(
                 (rp) => rp.parent.type === 'ExportSpecifier'
@@ -268,7 +290,7 @@ module.exports = function (babel) {
                 if (!isTypescriptNode(referencePath.parentPath)) {
                   const memberExpression = getMemberExpressionFor(
                     global,
-                    state.emberIdentifier
+                    state.ensureImport('default', 'ember')
                   );
 
                   try {