Fix import removal for omitted route chunks

Author: markdalgleish

packages/react-router-dev/vite/route-chunks-test.ts

@@ -178,13 +178,58 @@ describe("route chunks", () => {
         "import { shared } from "./shared";
         export const chunk = shared("chunk");"
       `);
-      expect(getChunkedExport(code, "main", {}, ...cache)?.code)
+      expect(omitChunkedExports(code, ["chunk"], {}, ...cache)?.code)
         .toMatchInlineSnapshot(`
         "import { shared } from "./shared";
         export const main = shared("main");"
       `);
     });
 
+    test("shared imports across chunks but not main chunk", () => {
+      const code = dedent`
+        import { shared } from "./shared";
+        export const chunk1 = shared("chunk1");
+        export const chunk2 = shared("chunk2");
+        export const main = "main";
+      `;
+      expect(hasChunkableExport(code, "chunk1", ...cache)).toBe(true);
+      expect(hasChunkableExport(code, "chunk2", ...cache)).toBe(true);
+      expect(hasChunkableExport(code, "main", ...cache)).toBe(true);
+      expect(getChunkedExport(code, "chunk1", {}, ...cache)?.code)
+        .toMatchInlineSnapshot(`
+        "import { shared } from "./shared";
+        export const chunk1 = shared("chunk1");"
+      `);
+      expect(getChunkedExport(code, "chunk2", {}, ...cache)?.code)
+        .toMatchInlineSnapshot(`
+        "import { shared } from "./shared";
+        export const chunk2 = shared("chunk2");"
+      `);
+      expect(
+        omitChunkedExports(code, ["chunk1", "chunk2"], {}, ...cache)?.code
+      ).toMatchInlineSnapshot(`"export const main = "main";"`);
+    });
+
+    test("import with side effect usage", () => {
+      const code = dedent`
+        import { sideEffect } from "./side-effect";
+        sideEffect();
+        export const chunk = "chunk";
+        export const main = "main";
+      `;
+      expect(hasChunkableExport(code, "chunk", ...cache)).toBe(true);
+      expect(hasChunkableExport(code, "main", ...cache)).toBe(true);
+      expect(
+        getChunkedExport(code, "chunk", {}, ...cache)?.code
+      ).toMatchInlineSnapshot(`"export const chunk = "chunk";"`);
+      expect(omitChunkedExports(code, ["chunk"], {}, ...cache)?.code)
+        .toMatchInlineSnapshot(`
+        "import { sideEffect } from "./side-effect";
+        sideEffect();
+        export const main = "main";"
+      `);
+    });
+
     test("re-exports using shared statement", () => {
       const code = dedent`
         export { chunk1, chunk2, main } from "./shared";
@@ -927,6 +972,29 @@ describe("route chunks", () => {
         export const main = getMainMessage();"
       `);
     });
+
+    test("shared imports across chunks but not main chunk with shared side effect usage", () => {
+      const code = dedent`
+        import { shared } from "./shared";
+        shared("side-effect");
+        export const chunk1 = shared("chunk1");
+        export const chunk2 = shared("chunk2");
+        export const main = "main";
+      `;
+      expect(hasChunkableExport(code, "chunk1", ...cache)).toBe(false);
+      expect(hasChunkableExport(code, "chunk2", ...cache)).toBe(false);
+      expect(hasChunkableExport(code, "main", ...cache)).toBe(true);
+      expect(getChunkedExport(code, "chunk1", {}, ...cache)).toBeUndefined();
+      expect(getChunkedExport(code, "chunk2", {}, ...cache)).toBeUndefined();
+      expect(omitChunkedExports(code, ["chunk1", "chunk2"], {}, ...cache)?.code)
+        .toMatchInlineSnapshot(`
+        "import { shared } from "./shared";
+        shared("side-effect");
+        export const chunk1 = shared("chunk1");
+        export const chunk2 = shared("chunk2");
+        export const main = "main";"
+      `);
+    });
   });
 
   describe("export dependency analysis", () => {

packages/react-router-dev/vite/route-chunks.ts

@@ -702,22 +702,31 @@ export function omitChunkedExports(
     )}::${JSON.stringify(generateOptions)}`,
     code,
     () => {
+      const isChunkable = (exportName: string): boolean =>
+        hasChunkableExport(code, exportName, cache, cacheKey);
+
+      const isOmitted = (exportName: string): boolean =>
+        exportNames.includes(exportName) && isChunkable(exportName);
+
+      const isRetained = (exportName: string): boolean =>
+        !isOmitted(exportName);
+
       let exportDependencies = getExportDependencies(code, cache, cacheKey);
 
+      let allExportNames = Array.from(exportDependencies.keys());
+      let omittedExportNames = allExportNames.filter(isOmitted);
+      let retainedExportNames = allExportNames.filter(isRetained);
+
       let omittedStatements = new Set<Statement>();
       let omittedExportedVariableDeclarators = new Set<VariableDeclarator>();
 
-      for (let exportName of exportNames) {
-        let dependencies = exportDependencies.get(exportName);
+      for (let omittedExportName of omittedExportNames) {
+        let dependencies = exportDependencies.get(omittedExportName);
 
-        // If the export is not chunkable then its code will still remain in the
-        // main chunk, so we need to keep its top level statements.
-        if (
-          !dependencies ||
-          !hasChunkableExport(code, exportName, cache, cacheKey)
-        ) {
-          continue;
-        }
+        invariant(
+          dependencies,
+          `Expected dependencies for ${omittedExportName}`
+        );
 
         // Now that we know the export is chunkable, add all of its top level
         // non-module statements to the set of statements to be omitted from the
@@ -742,14 +751,6 @@ export function omitChunkedExports(
         omittedExportedVariableDeclarators
       );
 
-      function isChunkable(exportName: string): boolean {
-        return hasChunkableExport(code, exportName, cache, cacheKey);
-      }
-
-      function isOmitted(exportName: string): boolean {
-        return exportNames.includes(exportName) && isChunkable(exportName);
-      }
-
       ast.program.body = ast.program.body
         // Remove top level statements that belong solely to the chunked
         // exports that are being omitted.
@@ -774,27 +775,28 @@ export function omitChunkedExports(
           // Remove import specifiers that are only used by the omitted chunks.
           // This ensures only the necessary imports remain in the main chunk.
           node.specifiers = node.specifiers.filter((specifier) => {
-            // Check the imported identifiers that each export depends on to see
-            // if it includes the specifier's local name.
-            for (let exportName of exportNames) {
-              // If the export is not chunkable then its code will still remain
-              // in the main chunk, so we need to keep its imports.
-              if (!isChunkable(exportName)) {
-                continue;
+            let importedName = specifier.local.name;
+
+            // Keep the import specifier if it's depended on by any of the
+            // retained exports.
+            for (let retainedExportName of retainedExportNames) {
+              let dependencies = exportDependencies.get(retainedExportName);
+              if (dependencies?.importedIdentifierNames?.has(importedName)) {
+                return true;
               }
+            }
 
-              let importedIdentifierNames =
-                exportDependencies.get(exportName)?.importedIdentifierNames;
-
-              // If the import specifier's local name is in the set of imported
-              // identifiers for the chunked export, we filter it out.
-              if (importedIdentifierNames?.has(specifier.local.name)) {
+            // Now that we've bailed out early and kept the import specifier if
+            // any retained exports depend on it, remove the import specifier if
+            // it's depended on by any of the omitted exports.
+            for (let omittedExportName of omittedExportNames) {
+              let dependencies = exportDependencies.get(omittedExportName);
+              if (dependencies?.importedIdentifierNames?.has(importedName)) {
                 return false;
               }
             }
 
-            // If we didn't return false, the specifier is not in the set of
-            // imported identifiers for any chunked export, so we keep it.
+            // Keep the import specifier if it isn't depended on by any export.
             return true;
           });