zod-v3-to-v4 1.12.1 → 1.13.0

package/README.md

@@ -135,6 +135,7 @@ pnpm playground:interactive
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/popfendi"><img src="https://avatars.githubusercontent.com/u/84185235?v=4?s=100" width="100px;" alt="popfendi"/><br /><sub><b>popfendi</b></sub></a><br /><a href="https://github.com/nicoespeon/zod-v3-to-v4/issues?q=author%3Apopfendi" title="Bug reports">🐛</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/luis-azevedo-visma"><img src="https://avatars.githubusercontent.com/u/149402066?v=4?s=100" width="100px;" alt="Luís Azevedo"/><br /><sub><b>Luís Azevedo</b></sub></a><br /><a href="#ideas-luis-azevedo-visma" title="Ideas, Planning, & Feedback">🤔</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://florian-lefebvre.dev"><img src="https://avatars.githubusercontent.com/u/69633530?v=4?s=100" width="100px;" alt="Florian Lefebvre"/><br /><sub><b>Florian Lefebvre</b></sub></a><br /><a href="https://github.com/nicoespeon/zod-v3-to-v4/issues?q=author%3Aflorian-lefebvre" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="http://www.robertcooper.me"><img src="https://avatars.githubusercontent.com/u/16786990?v=4?s=100" width="100px;" alt="Robert Cooper"/><br /><sub><b>Robert Cooper</b></sub></a><br /><a href="https://github.com/nicoespeon/zod-v3-to-v4/issues?q=author%3Arobertcoopercode" title="Bug reports">🐛</a></td>
     </tr>
   </tbody>
   <tfoot>

package/dist/collect-imports.js

@@ -25,16 +25,38 @@ export function getZodImport(importDeclaration) {
         ?.getNamedImports()
         .find((namedImport) => namedImport.getName() === "z");
 }
-export function collectZodReferences(importDeclarations) {
-    return importDeclarations.flatMap((importDeclaration) => importDeclaration.getNamedImports().flatMap((namedImport) => {
-        const namedNode = namedImport.getAliasNode() ?? namedImport.getNameNode();
-        if (!namedNode.isKind(SyntaxKind.Identifier)) {
-            return [];
+export function collectZodReferences(sourceFile, importDeclarations) {
+    const identifierMap = buildIdentifierMap(sourceFile);
+    return {
+        zodReferences: importDeclarations.flatMap((importDeclaration) => importDeclaration.getNamedImports().flatMap((namedImport) => {
+            const namedNode = namedImport.getAliasNode() ?? namedImport.getNameNode();
+            if (!namedNode.isKind(SyntaxKind.Identifier)) {
+                return [];
+            }
+            const name = namedNode.getText();
+            const allReferences = identifierMap.get(name) ?? [];
+            return allReferences.filter((node) => node !== namedNode);
+        })),
+        identifierMap,
+    };
+}
+/**
+ * Build a map of identifier name -> all nodes with that name.
+ * This allows O(1) lookups instead of expensive findReferencesAsNodes() calls.
+ */
+function buildIdentifierMap(sourceFile) {
+    const identifierMap = new Map();
+    for (const identifier of sourceFile.getDescendantsOfKind(SyntaxKind.Identifier)) {
+        const name = identifier.getText();
+        const existing = identifierMap.get(name);
+        if (existing) {
+            existing.push(identifier);
+        }
+        else {
+            identifierMap.set(name, [identifier]);
         }
-        return namedNode
-            .findReferencesAsNodes()
-            .filter((node) => node !== namedNode);
-    }));
+    }
+    return identifierMap;
 }
 /**
  * Collect references to variables that are derived from Zod schemas.
@@ -47,7 +69,7 @@ export function collectZodReferences(importDeclarations) {
  * This function will return references to `baseSchema` (like the one used
  * in `constrainedSchema`), so we can also migrate those.
  */
-export function collectDerivedZodSchemaReferences(zodReferences) {
+export function collectDerivedZodSchemaReferences(zodReferences, identifierMap) {
     const derivedReferences = [];
     const visited = new Set();
     function collectFromReferences(refs) {
@@ -67,11 +89,9 @@ export function collectDerivedZodSchemaReferences(zodReferences) {
             if (!nameNode.isKind(SyntaxKind.Identifier)) {
                 continue;
             }
-            // Find all references to this variable (excluding the declaration itself)
-            const variableReferences = nameNode
-                .findReferencesAsNodes()
-                .filter((node) => node !== nameNode)
-                .filter((node) => node.isKind(SyntaxKind.Identifier));
+            // Find all references to this variable
+            const allReferences = identifierMap.get(variableName) ?? [];
+            const variableReferences = allReferences.filter((node) => node !== nameNode);
             derivedReferences.push(...variableReferences);
             // Recursively collect references from these derived schemas
             collectFromReferences(variableReferences);

package/dist/migrate.js

@@ -25,8 +25,8 @@ export function migrateZodV3ToV4(sourceFile, options = {}) {
         }
     });
     // Collect references before modifying imports
-    const zodReferences = collectZodReferences(importDeclarations);
-    const derivedReferences = collectDerivedZodSchemaReferences(zodReferences);
+    const { zodReferences, identifierMap } = collectZodReferences(sourceFile, importDeclarations);
+    const derivedReferences = collectDerivedZodSchemaReferences(zodReferences, identifierMap);
     replaceDeletedTypes(importDeclarations, zodReferences);
     [...zodReferences, ...derivedReferences].forEach((node) => {
         if (node.wasForgotten()) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "zod-v3-to-v4",
-  "version": "1.12.1",
+  "version": "1.13.0",
   "description": "Migrate Zod from v3 to v4",
   "keywords": [
     "zod",