From 8e35d9fb87c8db103f18c72dc781c568c13a4e67 Mon Sep 17 00:00:00 2001
From: Zack Tanner <1939140+ztanner@users.noreply.github.com>
Date: Fri, 13 Dec 2024 07:08:52 -0800
Subject: [PATCH] refactor: collectAppPageSegments

---
 .../build/segment-config/app/app-segments.ts  | 57 ++++++++++++-------
 1 file changed, 37 insertions(+), 20 deletions(-)

diff --git a/packages/next/src/build/segment-config/app/app-segments.ts b/packages/next/src/build/segment-config/app/app-segments.ts
index 84c3828717722..115e384bbe54e 100644
--- a/packages/next/src/build/segment-config/app/app-segments.ts
+++ b/packages/next/src/build/segment-config/app/app-segments.ts
@@ -20,7 +20,10 @@ import {
 } from '../../../server/route-modules/checks'
 import { isClientReference } from '../../../lib/client-reference'
 import { getSegmentParam } from '../../../server/app-render/get-segment-param'
-import { getLayoutOrPageModule } from '../../../server/lib/app-dir-module'
+import {
+  getLayoutOrPageModule,
+  type LoaderTree,
+} from '../../../server/lib/app-dir-module'
 import { PAGE_SEGMENT_KEY } from '../../../shared/lib/segment'
 
 type GenerateStaticParams = (options: { params?: Params }) => Promise<Params[]>
@@ -74,17 +77,21 @@ export type AppSegment = {
  * @returns the segments for the app page route module
  */
 async function collectAppPageSegments(routeModule: AppPageRouteModule) {
-  const segments: AppSegment[] = []
+  // We keep track of unique segments, since with parallel routes, it's possible
+  // to see the same segment multiple times.
+  const uniqueSegments = new Map<string, AppSegment>()
+
+  // Queue will store tuples of [loaderTree, currentSegments]
+  type QueueItem = [LoaderTree, AppSegment[]]
+  const queue: QueueItem[] = [[routeModule.userland.loaderTree, []]]
 
-  // Helper function to process a loader tree path
-  async function processLoaderTree(
-    loaderTree: any,
-    currentSegments: AppSegment[] = []
-  ): Promise<void> {
+  while (queue.length > 0) {
+    const [loaderTree, currentSegments] = queue.shift()!
     const [name, parallelRoutes] = loaderTree
-    const { mod: userland, filePath } = await getLayoutOrPageModule(loaderTree)
 
-    const isClientComponent: boolean = userland && isClientReference(userland)
+    // Process current node
+    const { mod: userland, filePath } = await getLayoutOrPageModule(loaderTree)
+    const isClientComponent = userland && isClientReference(userland)
     const isDynamicSegment = /\[.*\]$/.test(name)
     const param = isDynamicSegment ? getSegmentParam(name)?.param : undefined
 
@@ -97,30 +104,40 @@ async function collectAppPageSegments(routeModule: AppPageRouteModule) {
       generateStaticParams: undefined,
     }
 
-    // Only server components can have app segment configurations. If this isn't
-    // an object, then we should skip it. This can happen when parsing the
-    // error components.
+    // Only server components can have app segment configurations
     if (!isClientComponent) {
       attach(segment, userland, routeModule.definition.pathname)
     }
 
-    currentSegments.push(segment)
+    // Create a unique key for the segment
+    const segmentKey = getSegmentKey(segment)
+    if (!uniqueSegments.has(segmentKey)) {
+      uniqueSegments.set(segmentKey, segment)
+    }
 
-    // If this is a page segment, we know we've reached a leaf node associated with the
-    // page we're collecting segments for. We can add the collected segments to our final result.
+    const updatedSegments = [...currentSegments, segment]
+
+    // If this is a page segment, we've reached a leaf node
     if (name === PAGE_SEGMENT_KEY) {
-      segments.push(...currentSegments)
+      // Add all segments in the current path
+      updatedSegments.forEach((seg) => {
+        const key = getSegmentKey(seg)
+        uniqueSegments.set(key, seg)
+      })
     }
 
-    // Recursively process parallel routes
+    // Add all parallel routes to the queue
     for (const parallelRouteKey in parallelRoutes) {
       const parallelRoute = parallelRoutes[parallelRouteKey]
-      await processLoaderTree(parallelRoute, [...currentSegments])
+      queue.push([parallelRoute, updatedSegments])
     }
   }
 
-  await processLoaderTree(routeModule.userland.loaderTree)
-  return segments
+  return Array.from(uniqueSegments.values())
+}
+
+function getSegmentKey(segment: AppSegment) {
+  return `${segment.name}-${segment.filePath ?? ''}-${segment.param ?? ''}`
 }
 
 /**