Support merging imports in import injector (#16349)

Co-authored-by: Conrad Buck <conartist6@gmail.com>
Co-authored-by: Babel Bot <babel-bot@users.noreply.github.com>

packages/babel-helper-module-imports/src/import-injector.ts

@@ -1,5 +1,11 @@
 import assert from "assert";
-import { numericLiteral, sequenceExpression } from "@babel/types";
+import {
+  identifier,
+  importSpecifier,
+  numericLiteral,
+  sequenceExpression,
+  isImportDeclaration,
+} from "@babel/types";
 import type * as t from "@babel/types";
 import type { NodePath, Scope } from "@babel/traverse";
 import type { File } from "@babel/core";
@@ -431,33 +437,135 @@ export default class ImportInjector {
     importPosition = "before",
     blockHoist = 3,
   ) {
-    const body = this._programPath.get("body");
-
     if (importPosition === "after") {
-      for (let i = body.length - 1; i >= 0; i--) {
-        if (body[i].isImportDeclaration()) {
-          body[i].insertAfter(statements);
-          return;
-        }
-      }
+      if (this._insertStatementsAfter(statements)) return;
     } else {
-      statements.forEach(node => {
-        // @ts-expect-error handle _blockHoist
-        node._blockHoist = blockHoist;
-      });
-
-      const targetPath = body.find(p => {
-        // @ts-expect-error todo(flow->ts): avoid mutations
-        const val = p.node._blockHoist;
-        return Number.isFinite(val) && val < 4;
-      });
-
-      if (targetPath) {
-        targetPath.insertBefore(statements);
-        return;
-      }
+      if (this._insertStatementsBefore(statements, blockHoist)) return;
     }
 
     this._programPath.unshiftContainer("body", statements);
   }
+
+  _insertStatementsBefore(statements: t.Statement[], blockHoist: number) {
+    if (
+      statements.length === 1 &&
+      isImportDeclaration(statements[0]) &&
+      isValueImport(statements[0])
+    ) {
+      const firstImportDecl = this._programPath
+        .get("body")
+        .find((p): p is NodePath<t.ImportDeclaration> => {
+          return p.isImportDeclaration() && isValueImport(p.node);
+        });
+
+      if (
+        firstImportDecl?.node.source.value === statements[0].source.value &&
+        maybeAppendImportSpecifiers(firstImportDecl.node, statements[0])
+      ) {
+        return true;
+      }
+    }
+
+    statements.forEach(node => {
+      // @ts-expect-error handle _blockHoist
+      node._blockHoist = blockHoist;
+    });
+
+    const targetPath = this._programPath.get("body").find(p => {
+      // @ts-expect-error todo(flow->ts): avoid mutations
+      const val = p.node._blockHoist;
+      return Number.isFinite(val) && val < 4;
+    });
+
+    if (targetPath) {
+      targetPath.insertBefore(statements);
+      return true;
+    }
+
+    return false;
+  }
+
+  _insertStatementsAfter(statements: t.Statement[]): boolean {
+    const statementsSet = new Set(statements);
+    const importDeclarations: Map<string, t.ImportDeclaration[]> = new Map();
+
+    for (const statement of statements) {
+      if (isImportDeclaration(statement) && isValueImport(statement)) {
+        const source = statement.source.value;
+        if (!importDeclarations.has(source)) importDeclarations.set(source, []);
+        importDeclarations.get(source).push(statement);
+      }
+    }
+
+    let lastImportPath = null;
+    for (const bodyStmt of this._programPath.get("body")) {
+      if (bodyStmt.isImportDeclaration() && isValueImport(bodyStmt.node)) {
+        lastImportPath = bodyStmt;
+
+        const source = bodyStmt.node.source.value;
+        const newImports = importDeclarations.get(source);
+        if (!newImports) continue;
+
+        for (const decl of newImports) {
+          if (maybeAppendImportSpecifiers(bodyStmt.node, decl)) {
+            statementsSet.delete(decl);
+          }
+        }
+      }
+    }
+
+    if (statementsSet.size === 0) return true;
+
+    if (lastImportPath) lastImportPath.insertAfter(Array.from(statementsSet));
+
+    return !!lastImportPath;
+  }
+}
+
+function isValueImport(node: t.ImportDeclaration) {
+  return node.importKind !== "type" && node.importKind !== "typeof";
+}
+
+function hasNamespaceImport(node: t.ImportDeclaration) {
+  return (
+    (node.specifiers.length === 1 &&
+      node.specifiers[0].type === "ImportNamespaceSpecifier") ||
+    (node.specifiers.length === 2 &&
+      node.specifiers[1].type === "ImportNamespaceSpecifier")
+  );
+}
+
+function hasDefaultImport(node: t.ImportDeclaration) {
+  return (
+    node.specifiers.length > 0 &&
+    node.specifiers[0].type === "ImportDefaultSpecifier"
+  );
+}
+
+function maybeAppendImportSpecifiers(
+  target: t.ImportDeclaration,
+  source: t.ImportDeclaration,
+): boolean {
+  if (!target.specifiers.length) {
+    target.specifiers = source.specifiers;
+    return true;
+  }
+  if (!source.specifiers.length) return true;
+
+  if (hasNamespaceImport(target) || hasNamespaceImport(source)) return false;
+
+  if (hasDefaultImport(source)) {
+    if (hasDefaultImport(target)) {
+      source.specifiers[0] = importSpecifier(
+        source.specifiers[0].local,
+        identifier("default"),
+      );
+    } else {
+      target.specifiers.unshift(source.specifiers.shift());
+    }
+  }
+
+  target.specifiers.push(...source.specifiers);
+
+  return true;
 }

packages/babel-helper-module-imports/test/index.js

@@ -6,19 +6,8 @@ import { ImportInjector } from "../lib/index.js";
 
 const cwd = path.dirname(fileURLToPath(import.meta.url));
 
-function test(sourceType, opts, initializer, inputCode, expectedCode) {
-  if (typeof opts === "function") {
-    expectedCode = inputCode;
-    inputCode = initializer;
-    initializer = opts;
-    opts = null;
-  }
-  if (expectedCode === undefined) {
-    expectedCode = inputCode;
-    inputCode = "";
-  }
-
-  const result = babel.transformSync(inputCode, {
+function transform(sourceType, opts, initializer, inputCode) {
+  return babel.transformSync(inputCode, {
     cwd,
     sourceType,
     filename: "example" + (sourceType === "module" ? ".mjs" : ".js"),
@@ -27,6 +16,9 @@ function test(sourceType, opts, initializer, inputCode, expectedCode) {
     plugins: [
       function ({ types: t }) {
         return {
+          manipulateOptions({ parserOpts }) {
+            parserOpts.plugins.push("typescript");
+          },
           pre(file) {
             file.set("helperGenerator", name =>
               t.memberExpression(
@@ -46,11 +38,26 @@ function test(sourceType, opts, initializer, inputCode, expectedCode) {
         };
       },
     ],
-  });
+  }).code;
+}
+
+function test(sourceType, opts, initializer, inputCode, expectedCode) {
+  if (typeof opts === "function") {
+    expectedCode = inputCode;
+    inputCode = initializer;
+    initializer = opts;
+    opts = null;
+  }
+  if (expectedCode === undefined) {
+    expectedCode = inputCode;
+    inputCode = "";
+  }
 
-  expect(result.code.replace(/\s+/g, " ").trim()).toBe(
-    (expectedCode || "").replace(/\s+/g, " ").trim(),
-  );
+  expect(
+    transform(sourceType, opts, initializer, inputCode)
+      .replace(/\s+/g, " ")
+      .trim(),
+  ).toBe((expectedCode || "").replace(/\s+/g, " ").trim());
 }
 const testScript = test.bind(undefined, "script");
 const testModule = test.bind(undefined, "module");
@@ -1143,5 +1150,124 @@ describe("@babel/helper-module-imports", () => {
         ),
       ).toThrow(`"importPosition": "after" is only supported in modules`);
     });
+
+    describe("imports merging", () => {
+      const opts = { importPosition: "after" };
+      const addNamespace = m => void m.addNamespace("s", opts);
+      const addDefault = m => void m.addDefault("s", opts);
+      const addNamed = m => void m.addNamed("n", "s", opts);
+      const addSideEffect = m => void m.addSideEffect("s", opts);
+
+      it.each`
+        input                                   | operation        | expected
+        ${`import "s"`}                         | ${addNamespace}  | ${`import * as _s from "s";`}
+        ${`import x from "s"`}                  | ${addNamespace}  | ${`import x from "s"; import * as _s from "s";`}
+        ${`import { x } from "s"`}              | ${addNamespace}  | ${`import { x } from "s"; import * as _s from "s";`}
+        ${`import * as x from "s"`}             | ${addNamespace}  | ${`import * as x from "s"; import * as _s from "s";`}
+        ${`import "s"`}                         | ${addNamed}      | ${`import { n as _n } from "s";`}
+        ${`import x from "s"`}                  | ${addNamed}      | ${`import x, { n as _n } from  "s";`}
+        ${`import { x } from "s"`}              | ${addNamed}      | ${`import { x, n as _n } from "s";`}
+        ${`import x, { y } from "s"`}           | ${addNamed}      | ${`import x, { y, n as _n } from  "s";`}
+        ${`import * as x from "s"`}             | ${addNamed}      | ${`import * as x from "s"; import { n as _n } from "s";`}
+        ${`import "s"`}                         | ${addDefault}    | ${`import _default from "s";`}
+        ${`import x from "s"`}                  | ${addDefault}    | ${`import x, { default as _default } from  "s";`}
+        ${`import { x } from "s"`}              | ${addDefault}    | ${`import _default, { x } from "s";`}
+        ${`import x, { y } from "s"`}           | ${addDefault}    | ${`import x, { y, default as _default } from  "s";`}
+        ${`import * as x from "s"`}             | ${addDefault}    | ${`import * as x from "s"; import _default from "s";`}
+        ${`import "s"`}                         | ${addSideEffect} | ${`import "s";`}
+        ${`import x from "s"`}                  | ${addSideEffect} | ${`import x from "s";`}
+        ${`import { x } from "s"`}              | ${addSideEffect} | ${`import { x } from "s";`}
+        ${`import * as x from "s"`}             | ${addSideEffect} | ${`import * as x from "s";`}
+        ${`import "u"; import type T from "s"`} | ${addSideEffect} | ${`import "u"; import "s"; import type T from "s";`}
+      `(
+        "$operation.name works with `$input`",
+        ({ input, operation, expected }) => {
+          const out = transform(
+            "module",
+            { importingInterop: "babel", importedType: "es6" },
+            operation,
+            input,
+          );
+          expect(out.replace(/[\s\n]+/g, " ")).toBe(
+            expected.replace(/[\s\n]+/g, " "),
+          );
+        },
+      );
+
+      describe("ordering", () => {
+        it("should try to merge imports", () => {
+          testModule(
+            { importingInterop: "babel", importedType: "es6" },
+            m => {
+              return babel.types.arrayExpression([
+                m.addNamed("x", "modA", { importPosition: "after" }),
+                m.addNamed("y", "modA", { importPosition: "after" }),
+                m.addNamed("z", "modB", { importPosition: "after" }),
+                m.addNamed("w", "modA", { importPosition: "after" }),
+              ]);
+            },
+            `
+              import { x as _x, y as _y, w as _w } from "modA";
+              import { z as _z } from "modB";
+              [_x, _y, _z, _w];
+            `,
+          );
+        });
+
+        it("with user imports", () => {
+          testModule(
+            { importingInterop: "babel", importedType: "es6" },
+            m => {
+              return babel.types.arrayExpression([
+                m.addNamed("x", "modA", { importPosition: "after" }),
+                m.addNamed("y", "modB", { importPosition: "after" }),
+                m.addNamed("z", "modC", { importPosition: "after" }),
+                m.addNamed("w", "modD", { importPosition: "after" }),
+              ]);
+            },
+            `
+              import { foo } from "modA";
+              import bar from "modB";
+              import * as baz from "modC";
+              import "modD";
+            `,
+            `
+              import { foo, x as _x } from "modA";
+              import bar, { y as _y } from "modB";
+              import * as baz from "modC";
+              import { w as _w } from "modD";
+              import { z as _z } from "modC";
+              [_x, _y, _z, _w];
+            `,
+          );
+        });
+
+        it("with importPosition: before", () => {
+          testModule(
+            { importingInterop: "babel", importedType: "es6" },
+            m => {
+              return babel.types.arrayExpression([
+                m.addNamed("x", "modA", { importPosition: "before" }),
+                m.addNamed("y", "modB", { importPosition: "before" }),
+                m.addNamed("z", "modC", { importPosition: "before" }),
+              ]);
+            },
+            `
+              import { foo } from "modA";
+              import bar from "modB";
+              import * as baz from "modC";
+            `,
+            `
+              import { z as _z } from "modC";
+              import { y as _y } from "modB";
+              import { foo, x as _x } from "modA";
+              import bar from "modB";
+              import * as baz from "modC";
+              [_x, _y, _z];
+            `,
+          );
+        });
+      });
+    });
   });
 });

packages/babel-plugin-proposal-import-wasm-source/test/fixtures/node-esm/import-declaration-multiple/output.mjs

@@ -1,5 +1,4 @@
-import { readFileSync as _readFileSync2 } from "fs";
-import { readFileSync as _readFileSync } from "fs";
+import { readFileSync as _readFileSync, readFileSync as _readFileSync2 } from "fs";
 const s = new WebAssembly.Module(_readFileSync(new URL(import.meta.resolve("./x.wasm")))),
   s2 = new WebAssembly.Module(_readFileSync2(new URL(import.meta.resolve("./x2.wasm"))));
 someBody;

packages/babel-plugin-proposal-json-modules/test/fixtures/node-esm/import-declaration-multiple/output.mjs

@@ -1,5 +1,4 @@
-import { readFileSync as _readFileSync2 } from "fs";
-import { readFileSync as _readFileSync } from "fs";
+import { readFileSync as _readFileSync, readFileSync as _readFileSync2 } from "fs";
 const j = JSON.parse(_readFileSync(new URL(import.meta.resolve("./x.json")))),
   j2 = JSON.parse(_readFileSync2(new URL(import.meta.resolve("./x2.json"))));
 someBody;

packages/babel-plugin-proposal-record-and-tuple/test/fixtures/import-polyfill/custom-name-implies-should-import/output.mjs

@@ -1,5 +1,4 @@
-import { Tuple as _Tuple } from "my-polyfill";
-import { Record as _Record } from "my-polyfill";
+import { Record as _Record, Tuple as _Tuple } from "my-polyfill";
 const r2 = _Record({
   a: _Record({
     b: 456

packages/babel-plugin-proposal-record-and-tuple/test/fixtures/import-polyfill/custom-name/output.mjs

@@ -1,5 +1,4 @@
-import { Tuple as _Tuple } from "my-polyfill";
-import { Record as _Record } from "my-polyfill";
+import { Record as _Record, Tuple as _Tuple } from "my-polyfill";
 const r2 = _Record({
   a: _Record({
     b: 456

packages/babel-plugin-proposal-record-and-tuple/test/fixtures/import-polyfill/default-name/output.mjs

@@ -1,5 +1,4 @@
-import { Tuple as _Tuple } from "@bloomberg/record-tuple-polyfill";
-import { Record as _Record } from "@bloomberg/record-tuple-polyfill";
+import { Record as _Record, Tuple as _Tuple } from "@bloomberg/record-tuple-polyfill";
 const r2 = _Record({
   a: _Record({
     b: 456