Rename variables in LayoutRouter for clarity (#73826)

LayoutRouter has changed significantly since it was originally written
and the structure has become harder to follow. One thing I always find
confusing whenever I'm reading this code is that some of the values
correspond the the *parent* segment, while others correspond to the
*current* segment. So this moves the code around a bit and updates the
names so it's clearer which parts belong to which segments.

While working on this, I noticed a funny thing about how loading
boundaries work that was made more obvious by the revised naming. I've
left a TODO comment to follow up on whether this was intentional.

packages/next/src/client/components/app-router.tsx

@@ -564,14 +564,13 @@ function Router({
 
   const layoutRouterContext = useMemo(() => {
     return {
-      childNodes: cache.parallelRoutes,
-      tree,
+      parentTree: tree,
+      parentCacheNode: cache,
       // Root node always has `url`
       // Provided in AppTreeContext to ensure it can be overwritten in layout-router
       url: canonicalUrl,
-      loading: cache.loading,
     }
-  }, [cache.parallelRoutes, tree, canonicalUrl, cache.loading])
+  }, [tree, cache, canonicalUrl])
 
   const globalLayoutRouterContext = useMemo(() => {
     return {

packages/next/src/client/components/layout-router.tsx

@@ -1,14 +1,13 @@
 'use client'
 
 import type {
-  ChildSegmentMap,
+  CacheNode,
   LazyCacheNode,
   LoadingModuleData,
 } from '../../shared/lib/app-router-context.shared-runtime'
 import type {
   FlightRouterState,
   FlightSegmentPath,
-  Segment,
 } from '../../server/app-render/types'
 import type { ErrorComponent } from './error-boundary'
 import type { FocusAndScrollRef } from './router-reducer/router-reducer-types'
@@ -34,7 +33,6 @@ import { matchSegment } from './match-segments'
 import { handleSmoothScroll } from '../../shared/lib/router/utils/handle-smooth-scroll'
 import { RedirectBoundary } from './redirect-boundary'
 import { HTTPAccessFallbackBoundary } from './http-access-fallback/error-boundary'
-import { getSegmentValue } from './router-reducer/reducers/get-segment-value'
 import { createRouterCacheKey } from './router-reducer/create-router-cache-key'
 import { hasInterceptionRouteInCurrentTree } from './router-reducer/reducers/has-interception-route-in-current-tree'
 
@@ -316,22 +314,15 @@ function ScrollAndFocusHandler({
  * InnerLayoutRouter handles rendering the provided segment based on the cache.
  */
 function InnerLayoutRouter({
-  parallelRouterKey,
-  url,
-  childNodes,
-  segmentPath,
   tree,
-  // TODO-APP: implement `<Offscreen>` when available.
-  // isActive,
-  cacheKey,
+  segmentPath,
+  cacheNode,
+  url,
 }: {
-  parallelRouterKey: string
-  url: string
-  childNodes: ChildSegmentMap
-  segmentPath: FlightSegmentPath
   tree: FlightRouterState
-  isActive: boolean
-  cacheKey: ReturnType<typeof createRouterCacheKey>
+  segmentPath: FlightSegmentPath
+  cacheNode: CacheNode
+  url: string
 }) {
   const context = useContext(GlobalLayoutRouterContext)
   if (!context) {
@@ -340,29 +331,6 @@ function InnerLayoutRouter({
 
   const { changeByServerResponse, tree: fullTree } = context
 
-  // Read segment path from the parallel router cache node.
-  let childNode = childNodes.get(cacheKey)
-
-  // When data is not available during rendering client-side we need to fetch
-  // it from the server.
-  if (childNode === undefined) {
-    const newLazyCacheNode: LazyCacheNode = {
-      lazyData: null,
-      rsc: null,
-      prefetchRsc: null,
-      head: null,
-      prefetchHead: null,
-      parallelRoutes: new Map(),
-      loading: null,
-    }
-
-    /**
-     * Flight data fetch kicked off during render and put into the cache.
-     */
-    childNode = newLazyCacheNode
-    childNodes.set(cacheKey, newLazyCacheNode)
-  }
-
   // `rsc` represents the renderable node for this segment.
 
   // If this segment has a `prefetchRsc`, it's the statically prefetched data.
@@ -371,7 +339,7 @@ function InnerLayoutRouter({
   //
   // If no prefetch data is available, then we go straight to rendering `rsc`.
   const resolvedPrefetchRsc =
-    childNode.prefetchRsc !== null ? childNode.prefetchRsc : childNode.rsc
+    cacheNode.prefetchRsc !== null ? cacheNode.prefetchRsc : cacheNode.rsc
 
   // We use `useDeferredValue` to handle switching between the prefetched and
   // final values. The second argument is returned on initial render, then it
@@ -380,7 +348,7 @@ function InnerLayoutRouter({
   // @ts-expect-error The second argument to `useDeferredValue` is only
   // available in the experimental builds. When its disabled, it will always
   // return `rsc`.
-  const rsc: any = useDeferredValue(childNode.rsc, resolvedPrefetchRsc)
+  const rsc: any = useDeferredValue(cacheNode.rsc, resolvedPrefetchRsc)
 
   // `rsc` is either a React node or a promise for a React node, except we
   // special case `null` to represent that this segment's data is missing. If
@@ -397,15 +365,15 @@ function InnerLayoutRouter({
     // the server and patch the cache.
 
     // Check if there's already a pending request.
-    let lazyData = childNode.lazyData
+    let lazyData = cacheNode.lazyData
     if (lazyData === null) {
       /**
        * Router state with refetch marker added
        */
       // TODO-APP: remove ''
       const refetchTree = walkAddRefetch(['', ...segmentPath], fullTree)
       const includeNextUrl = hasInterceptionRouteInCurrentTree(fullTree)
-      childNode.lazyData = lazyData = fetchServerResponse(
+      cacheNode.lazyData = lazyData = fetchServerResponse(
         new URL(url, location.origin),
         {
           flightRouterState: refetchTree,
@@ -432,11 +400,11 @@ function InnerLayoutRouter({
     // The layout router context narrows down tree and childNodes at each level.
     <LayoutRouterContext.Provider
       value={{
-        tree: tree[1][parallelRouterKey],
-        childNodes: childNode.parallelRoutes,
+        parentTree: tree,
+        parentCacheNode: cacheNode,
+
         // TODO-APP: overriding of url for parallel routes
         url: url,
-        loading: childNode.loading,
       }}
     >
       {resolvedRsc}
@@ -533,86 +501,110 @@ export default function OuterLayoutRouter({
     throw new Error('invariant expected layout router to be mounted')
   }
 
-  const { childNodes, tree, url, loading } = context
+  const { parentTree, parentCacheNode, url } = context
 
-  // Get the current parallelRouter cache node
-  let childNodesForParallelRouter = childNodes.get(parallelRouterKey)
+  // Get the CacheNode for this segment by reading it from the parent segment's
+  // child map.
+  const parentParallelRoutes = parentCacheNode.parallelRoutes
+  let segmentMap = parentParallelRoutes.get(parallelRouterKey)
   // If the parallel router cache node does not exist yet, create it.
   // This writes to the cache when there is no item in the cache yet. It never *overwrites* existing cache items which is why it's safe in concurrent mode.
-  if (!childNodesForParallelRouter) {
-    childNodesForParallelRouter = new Map()
-    childNodes.set(parallelRouterKey, childNodesForParallelRouter)
+  if (!segmentMap) {
+    segmentMap = new Map()
+    parentParallelRoutes.set(parallelRouterKey, segmentMap)
   }
 
   // Get the active segment in the tree
   // The reason arrays are used in the data format is that these are transferred from the server to the browser so it's optimized to save bytes.
-  const treeSegment = tree[1][parallelRouterKey][0]
+  const tree = parentTree[1][parallelRouterKey]
+  const treeSegment = tree[0]
+
+  // The "state" key of a segment is the one passed to React — it represents the
+  // identity of the UI tree. Whenever the state key changes, the tree is
+  // recreated and the state is reset. In the App Router model, search params do
+  // not cause state to be lost, so two segments with the same segment path but
+  // different search params should have the same state key.
+  //
+  // The "cache" key of a segment, however, *does* include the search params, if
+  // it's possible that the segment accessed the search params on the server.
+  // (This only applies to page segments; layout segments cannot access search
+  // params on the server.)
+  const cacheKey = createRouterCacheKey(treeSegment)
+  const stateKey = createRouterCacheKey(treeSegment, true) // no search params
+
+  // Read segment path from the parallel router cache node.
+  let cacheNode = segmentMap.get(cacheKey)
+  if (cacheNode === undefined) {
+    // When data is not available during rendering client-side we need to fetch
+    // it from the server.
+    const newLazyCacheNode: LazyCacheNode = {
+      lazyData: null,
+      rsc: null,
+      prefetchRsc: null,
+      head: null,
+      prefetchHead: null,
+      parallelRoutes: new Map(),
+      loading: null,
+    }
 
-  // If segment is an array it's a dynamic route and we want to read the dynamic route value as the segment to get from the cache.
-  const currentChildSegmentValue = getSegmentValue(treeSegment)
+    // Flight data fetch kicked off during render and put into the cache.
+    cacheNode = newLazyCacheNode
+    segmentMap.set(cacheKey, newLazyCacheNode)
+  }
 
-  /**
-   * Decides which segments to keep rendering, all segments that are not active will be wrapped in `<Offscreen>`.
-   */
-  // TODO-APP: Add handling of `<Offscreen>` when it's available.
-  const preservedSegments: Segment[] = [treeSegment]
+  /*
+    - Error boundary
+      - Only renders error boundary if error component is provided.
+      - Rendered for each segment to ensure they have their own error state.
+    - Loading boundary
+      - Only renders suspense boundary if loading components is provided.
+      - Rendered for each segment to ensure they have their own loading state.
+      - Passed to the router during rendering to ensure it can be immediately rendered when suspending on a Flight fetch.
+  */
+
+  // TODO: The loading module data for a segment is stored on the parent, then
+  // applied to each of that parent segment's parallel route slots. In the
+  // simple case where there's only one parallel route (the `children` slot),
+  // this is no different from if the loading module data where stored on the
+  // child directly. But I'm not sure this actually makes sense when there are
+  // multiple parallel routes. It's not a huge issue because you always have
+  // the option to define a narrower loading boundary for a particular slot. But
+  // this sort of smells like an implementation accident to me.
+  const loadingModuleData = parentCacheNode.loading
 
   return (
-    <>
-      {preservedSegments.map((preservedSegment) => {
-        const preservedSegmentValue = getSegmentValue(preservedSegment)
-        const cacheKey = createRouterCacheKey(preservedSegment)
-
-        return (
-          /*
-            - Error boundary
-              - Only renders error boundary if error component is provided.
-              - Rendered for each segment to ensure they have their own error state.
-            - Loading boundary
-              - Only renders suspense boundary if loading components is provided.
-              - Rendered for each segment to ensure they have their own loading state.
-              - Passed to the router during rendering to ensure it can be immediately rendered when suspending on a Flight fetch.
-          */
-          <TemplateContext.Provider
-            key={createRouterCacheKey(preservedSegment, true)}
-            value={
-              <ScrollAndFocusHandler segmentPath={segmentPath}>
-                <ErrorBoundary
-                  errorComponent={error}
-                  errorStyles={errorStyles}
-                  errorScripts={errorScripts}
-                >
-                  <LoadingBoundary loading={loading}>
-                    <HTTPAccessFallbackBoundary
-                      notFound={notFound}
-                      forbidden={forbidden}
-                      unauthorized={unauthorized}
-                    >
-                      <RedirectBoundary>
-                        <InnerLayoutRouter
-                          parallelRouterKey={parallelRouterKey}
-                          url={url}
-                          tree={tree}
-                          childNodes={childNodesForParallelRouter!}
-                          segmentPath={segmentPath}
-                          cacheKey={cacheKey}
-                          isActive={
-                            currentChildSegmentValue === preservedSegmentValue
-                          }
-                        />
-                      </RedirectBoundary>
-                    </HTTPAccessFallbackBoundary>
-                  </LoadingBoundary>
-                </ErrorBoundary>
-              </ScrollAndFocusHandler>
-            }
+    <TemplateContext.Provider
+      key={stateKey}
+      value={
+        <ScrollAndFocusHandler segmentPath={segmentPath}>
+          <ErrorBoundary
+            errorComponent={error}
+            errorStyles={errorStyles}
+            errorScripts={errorScripts}
           >
-            {templateStyles}
-            {templateScripts}
-            {template}
-          </TemplateContext.Provider>
-        )
-      })}
-    </>
+            <LoadingBoundary loading={loadingModuleData}>
+              <HTTPAccessFallbackBoundary
+                notFound={notFound}
+                forbidden={forbidden}
+                unauthorized={unauthorized}
+              >
+                <RedirectBoundary>
+                  <InnerLayoutRouter
+                    url={url}
+                    tree={tree}
+                    cacheNode={cacheNode}
+                    segmentPath={segmentPath}
+                  />
+                </RedirectBoundary>
+              </HTTPAccessFallbackBoundary>
+            </LoadingBoundary>
+          </ErrorBoundary>
+        </ScrollAndFocusHandler>
+      }
+    >
+      {templateStyles}
+      {templateScripts}
+      {template}
+    </TemplateContext.Provider>
   )
 }

packages/next/src/client/components/navigation.ts

@@ -221,7 +221,7 @@ export function useSelectedLayoutSegments(
   // @ts-expect-error This only happens in `pages`. Type is overwritten in navigation.d.ts
   if (!context) return null
 
-  return getSelectedLayoutSegmentPath(context.tree, parallelRouteKey)
+  return getSelectedLayoutSegmentPath(context.parentTree, parallelRouteKey)
 }
 
 /**

packages/next/src/shared/lib/app-router-context.shared-runtime.ts

@@ -146,10 +146,9 @@ export const AppRouterContext = React.createContext<AppRouterInstance | null>(
   null
 )
 export const LayoutRouterContext = React.createContext<{
-  childNodes: CacheNode['parallelRoutes']
-  tree: FlightRouterState
+  parentTree: FlightRouterState
+  parentCacheNode: CacheNode
   url: string
-  loading: LoadingModuleData | Promise<LoadingModuleData>
 } | null>(null)
 
 export const GlobalLayoutRouterContext = React.createContext<{