From 81327eb980c308474a586a9cb9c0c5fff10eba34 Mon Sep 17 00:00:00 2001
From: "Eckhardt (Kaizen) Dreyer" <eckhardt.dreyer@gmail.com>
Date: Sat, 3 Aug 2024 04:30:55 +0200
Subject: [PATCH] perf(ssr): do a single-pass over AST with node cache arrays
 (#17812)

---
 packages/vite/src/node/ssr/ssrTransform.ts | 93 +++++++++++++---------
 1 file changed, 57 insertions(+), 36 deletions(-)

diff --git a/packages/vite/src/node/ssr/ssrTransform.ts b/packages/vite/src/node/ssr/ssrTransform.ts
index ac5d520161741a..a03eb7266d9d18 100644
--- a/packages/vite/src/node/ssr/ssrTransform.ts
+++ b/packages/vite/src/node/ssr/ssrTransform.ts
@@ -2,8 +2,12 @@ import path from 'node:path'
 import MagicString from 'magic-string'
 import type { SourceMap } from 'rollup'
 import type {
+  ExportAllDeclaration,
+  ExportDefaultDeclaration,
+  ExportNamedDeclaration,
   Function as FunctionNode,
   Identifier,
+  ImportDeclaration,
   Pattern,
   Property,
   VariableDeclaration,
@@ -130,53 +134,70 @@ async function ssrTransformScript(
     )
   }
 
-  // 1. check all import statements and record id -> importName map
+  const imports: (ImportDeclaration & { start: number; end: number })[] = []
+  const exports: ((
+    | ExportNamedDeclaration
+    | ExportDefaultDeclaration
+    | ExportAllDeclaration
+  ) & { start: number; end: number })[] = []
+
   for (const node of ast.body as Node[]) {
+    if (node.type === 'ImportDeclaration') {
+      imports.push(node)
+    } else if (
+      node.type === 'ExportNamedDeclaration' ||
+      node.type === 'ExportDefaultDeclaration' ||
+      node.type === 'ExportAllDeclaration'
+    ) {
+      exports.push(node)
+    }
+  }
+
+  // 1. check all import statements and record id -> importName map
+  for (const node of imports) {
     // import foo from 'foo' --> foo -> __import_foo__.default
     // import { baz } from 'foo' --> baz -> __import_foo__.baz
     // import * as ok from 'foo' --> ok -> __import_foo__
-    if (node.type === 'ImportDeclaration') {
-      const importId = defineImport(hoistIndex, node.source.value as string, {
-        importedNames: node.specifiers
-          .map((s) => {
-            if (s.type === 'ImportSpecifier')
-              return s.imported.type === 'Identifier'
-                ? s.imported.name
-                : // @ts-expect-error TODO: Estree types don't consider arbitrary module namespace specifiers yet
-                  s.imported.value
-            else if (s.type === 'ImportDefaultSpecifier') return 'default'
-          })
-          .filter(isDefined),
-      })
-      s.remove(node.start, node.end)
-      for (const spec of node.specifiers) {
-        if (spec.type === 'ImportSpecifier') {
-          if (spec.imported.type === 'Identifier') {
-            idToImportMap.set(
-              spec.local.name,
-              `${importId}.${spec.imported.name}`,
-            )
-          } else {
-            idToImportMap.set(
-              spec.local.name,
-              `${importId}[${
-                // @ts-expect-error TODO: Estree types don't consider arbitrary module namespace specifiers yet
-                JSON.stringify(spec.imported.value)
-              }]`,
-            )
-          }
-        } else if (spec.type === 'ImportDefaultSpecifier') {
-          idToImportMap.set(spec.local.name, `${importId}.default`)
+    const importId = defineImport(hoistIndex, node.source.value as string, {
+      importedNames: node.specifiers
+        .map((s) => {
+          if (s.type === 'ImportSpecifier')
+            return s.imported.type === 'Identifier'
+              ? s.imported.name
+              : // @ts-expect-error TODO: Estree types don't consider arbitrary module namespace specifiers yet
+                s.imported.value
+          else if (s.type === 'ImportDefaultSpecifier') return 'default'
+        })
+        .filter(isDefined),
+    })
+    s.remove(node.start, node.end)
+    for (const spec of node.specifiers) {
+      if (spec.type === 'ImportSpecifier') {
+        if (spec.imported.type === 'Identifier') {
+          idToImportMap.set(
+            spec.local.name,
+            `${importId}.${spec.imported.name}`,
+          )
         } else {
-          // namespace specifier
-          idToImportMap.set(spec.local.name, importId)
+          idToImportMap.set(
+            spec.local.name,
+            `${importId}[${
+              // @ts-expect-error TODO: Estree types don't consider arbitrary module namespace specifiers yet
+              JSON.stringify(spec.imported.value)
+            }]`,
+          )
         }
+      } else if (spec.type === 'ImportDefaultSpecifier') {
+        idToImportMap.set(spec.local.name, `${importId}.default`)
+      } else {
+        // namespace specifier
+        idToImportMap.set(spec.local.name, importId)
       }
     }
   }
 
   // 2. check all export statements and define exports
-  for (const node of ast.body as Node[]) {
+  for (const node of exports) {
     // named exports
     if (node.type === 'ExportNamedDeclaration') {
       if (node.declaration) {