Add comments in new router reducer

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

@@ -2,7 +2,7 @@ import React, { useEffect } from 'react'
 import { createFromReadableStream } from 'next/dist/compiled/react-server-dom-webpack'
 import {
   AppRouterContext,
-  AppTreeContext,
+  LayoutRouterContext,
   GlobalLayoutRouterContext,
 } from '../../shared/lib/app-router-context'
 import type {
@@ -99,7 +99,7 @@ export default function AppRouter({
   children: React.ReactNode
   hotReloader?: React.ReactNode
 }) {
-  const [{ tree, cache, pushRef, focusRef, canonicalUrl }, dispatch] =
+  const [{ tree, cache, pushRef, focusAndScrollRef, canonicalUrl }, dispatch] =
     React.useReducer(reducer, {
       tree: initialTree,
       cache: {
@@ -109,7 +109,7 @@ export default function AppRouter({
           typeof window === 'undefined' ? new Map() : initialParallelRoutes,
       },
       pushRef: { pendingPush: false, mpaNavigation: false },
-      focusRef: { focus: false },
+      focusAndScrollRef: { apply: false },
       canonicalUrl:
         initialCanonicalUrl +
         // Hash is read as the initial value for canonicalUrl in the browser
@@ -300,11 +300,11 @@ export default function AppRouter({
           value={{
             changeByServerResponse,
             tree,
-            focusRef,
+            focusAndScrollRef,
           }}
         >
           <AppRouterContext.Provider value={appRouter}>
-            <AppTreeContext.Provider
+            <LayoutRouterContext.Provider
               value={{
                 childNodes: cache.parallelRoutes,
                 tree: tree,
@@ -324,7 +324,7 @@ export default function AppRouter({
                 // HotReloader uses the router tree and router.reload() in order to apply Server Component changes.
                 hotReloader
               }
-            </AppTreeContext.Provider>
+            </LayoutRouterContext.Provider>
           </AppRouterContext.Provider>
         </GlobalLayoutRouterContext.Provider>
       </SearchParamsContext.Provider>

packages/next/client/components/hooks-client.ts

@@ -9,7 +9,7 @@ import {
 } from './hooks-client-context'
 import {
   AppRouterContext,
-  AppTreeContext,
+  LayoutRouterContext,
 } from '../../shared/lib/app-router-context'
 
 /**
@@ -59,7 +59,7 @@ export function usePathname(): string {
 export function useSelectedLayoutSegment(
   parallelRouteKey: string = 'children'
 ): string {
-  const { tree } = useContext(AppTreeContext)
+  const { tree } = useContext(LayoutRouterContext)
 
   const segment = tree[1][parallelRouteKey][0]
 

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

@@ -1,37 +1,51 @@
 import React, { useContext, useEffect, useRef } from 'react'
-import type { ChildProp } from '../../server/app-render'
+import type { ChildProp, Segment } from '../../server/app-render'
 import type { ChildSegmentMap } from '../../shared/lib/app-router-context'
 import type {
   FlightRouterState,
   FlightSegmentPath,
   FlightDataPath,
 } from '../../server/app-render'
 import {
-  AppTreeContext,
+  LayoutRouterContext,
   GlobalLayoutRouterContext,
 } from '../../shared/lib/app-router-context'
 import { fetchServerResponse } from './app-router.client'
 import { matchSegment } from './match-segments'
 
-let infinitePromise: Promise<void> | Error
-
-function equalArray(a: any[], b: any[]) {
+/**
+ * Check if every segment in array a and b matches
+ */
+function equalSegmentPaths(a: Segment[], b: Segment[]) {
+  // Comparing length is a fast path.
   return a.length === b.length && a.every((val, i) => matchSegment(val, b[i]))
 }
 
-function pathMatches(
+/**
+ * Check if flightDataPath matches layoutSegmentPath
+ */
+function segmentPathMatches(
   flightDataPath: FlightDataPath,
   layoutSegmentPath: FlightSegmentPath
 ): boolean {
-  // The last two items are the tree and subTreeData
+  // The last three items are the current segment, tree, and subTreeData
   const pathToLayout = flightDataPath.slice(0, -3)
-  return equalArray(layoutSegmentPath, pathToLayout)
+  return equalSegmentPaths(layoutSegmentPath, pathToLayout)
 }
 
+/**
+ * Used to cache in createInfinitePromise
+ */
+let infinitePromise: Promise<void> | Error
+
+/**
+ * Create a Promise that does not resolve. This is used to suspend when data is not available yet.
+ */
 function createInfinitePromise() {
   if (!infinitePromise) {
+    // Only create the Promise once
     infinitePromise = new Promise((/* resolve */) => {
-      // Note: this is used to debug when the rendering is never updated.
+      // This is used to debug when the rendering is never updated.
       // setTimeout(() => {
       //   infinitePromise = new Error('Infinite promise')
       //   resolve()
@@ -42,18 +56,25 @@ function createInfinitePromise() {
   return infinitePromise
 }
 
+/**
+ * Check if the top of the HTMLElement is in the viewport.
+ */
 function topOfElementInViewport(element: HTMLElement) {
   const rect = element.getBoundingClientRect()
   return rect.top >= 0
 }
 
+/**
+ * InnerLayoutRouter handles rendering the provided segment based on the cache.
+ */
 export function InnerLayoutRouter({
   parallelRouterKey,
   url,
   childNodes,
   childProp,
   segmentPath,
   tree,
+  // TODO-APP: implement `<Offscreen>` when available.
   // isActive,
   path,
   rootLayoutIncluded,
@@ -71,35 +92,48 @@ export function InnerLayoutRouter({
   const {
     changeByServerResponse,
     tree: fullTree,
-    focusRef,
+    focusAndScrollRef,
   } = useContext(GlobalLayoutRouterContext)
-  const focusAndScrollRef = useRef<HTMLDivElement>(null)
+  const focusAndScrollElementRef = useRef<HTMLDivElement>(null)
 
   useEffect(() => {
-    if (focusRef.focus && focusAndScrollRef.current) {
-      focusRef.focus = false
-      focusAndScrollRef.current.focus()
+    // Handle scroll and focus, it's only applied once in the first useEffect that triggers that changed.
+    if (focusAndScrollRef.apply && focusAndScrollElementRef.current) {
+      // State is mutated to ensure that the focus and scroll is applied only once.
+      focusAndScrollRef.apply = false
+      // Set focus on the element
+      focusAndScrollElementRef.current.focus()
       // Only scroll into viewport when the layout is not visible currently.
-      if (!topOfElementInViewport(focusAndScrollRef.current)) {
-        focusAndScrollRef.current.scrollIntoView()
+      if (!topOfElementInViewport(focusAndScrollElementRef.current)) {
+        focusAndScrollElementRef.current.scrollIntoView()
       }
     }
-  }, [focusRef])
+  }, [focusAndScrollRef])
 
+  // Read segment path from the parallel router cache node.
   let childNode = childNodes.get(path)
 
+  // If childProp is available this means it's the Flight / SSR case.
   if (childProp && !childNode) {
+    // Add the segment's subTreeData to the cache.
+    // 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.
     childNodes.set(path, {
       data: null,
       subTreeData: childProp.current,
       parallelRoutes: new Map(),
     })
+    // Mutates the prop in order to clean up the memory associated with the subTreeData as it is now part of the cache.
     childProp.current = null
     // In the above case childNode was set on childNodes, so we have to get it from the cacheNodes again.
     childNode = childNodes.get(path)
   }
 
+  // When childNode is not available during rendering client-side we need to fetch it from the server.
   if (!childNode) {
+    /**
+     * Add refetch marker to router state at the point of the current layout segment.
+     * This ensures the response returned is not further down than the current layout segment.
+     */
     const walkAddRefetch = (
       segmentPathToWalk: FlightSegmentPath | undefined,
       treeToRecreate: FlightRouterState
@@ -145,9 +179,15 @@ export function InnerLayoutRouter({
       return treeToRecreate
     }
 
+    /**
+     * Router state with refetch marker added
+     */
     // TODO-APP: remove ''
     const refetchTree = walkAddRefetch(['', ...segmentPath], fullTree)
 
+    /**
+     * Flight data fetch kicked off during render and put into the cache.
+     */
     const data = fetchServerResponse(new URL(url, location.origin), refetchTree)
     childNodes.set(path, {
       data,
@@ -158,19 +198,23 @@ export function InnerLayoutRouter({
     childNode = childNodes.get(path)
   }
 
-  // In the above case childNode was set on childNodes, so we have to get it from the cacheNodes again.
-  childNode = childNodes.get(path)
-
+  // This case should never happen so it throws an error. It indicates there's a bug in the Next.js.
   if (!childNode) {
     throw new Error('Child node should always exist')
   }
 
+  // This case should never happen so it throws an error. It indicates there's a bug in the Next.js.
   if (childNode.subTreeData && childNode.data) {
     throw new Error('Child node should not have both subTreeData and data')
   }
 
+  // If cache node has a data request we have to readRoot and update the cache.
   if (childNode.data) {
     // TODO-APP: error case
+    /**
+     * Flight response data
+     */
+    // When the data has not resolved yet readRoot will suspend here.
     const flightData = childNode.data.readRoot()
 
     // Handle case when navigating to page in `pages` from `app`
@@ -179,28 +223,36 @@ export function InnerLayoutRouter({
       return null
     }
 
+    /**
+     * If the fast path was triggered.
+     * The fast path is when the returned Flight data path matches the layout segment path, then we can write the data to the cache in render instead of dispatching an action.
+     */
     let fastPath: boolean = false
-    // segmentPath matches what came back from the server. This is the happy path.
+
+    // If there are multiple patches returned in the Flight data we need to dispatch to ensure a single render.
     if (flightData.length === 1) {
       const flightDataPath = flightData[0]
 
-      if (pathMatches(flightDataPath, segmentPath)) {
+      if (segmentPathMatches(flightDataPath, segmentPath)) {
+        // Ensure data is set to null as subTreeData will be set in the cache now.
         childNode.data = null
         // Last item is the subtreeData
         // TODO-APP: routerTreePatch needs to be applied to the tree, handle it in render?
         const [, /* routerTreePatch */ subTreeData] = flightDataPath.slice(-2)
+        // Add subTreeData into the cache
         childNode.subTreeData = subTreeData
+        // This field is required for new items
         childNode.parallelRoutes = new Map()
         fastPath = true
       }
     }
 
+    // When the fast path is not used a new action is dispatched to update the tree and cache.
     if (!fastPath) {
-      // For push we can set data in the cache
-
       // segmentPath from the server does not match the layout's segmentPath
       childNode.data = null
 
+      // setTimeout is used to start a new transition during render, this is an intentional hack around React.
       setTimeout(() => {
         // @ts-ignore startTransition exists
         React.startTransition(() => {
@@ -213,13 +265,15 @@ export function InnerLayoutRouter({
     }
   }
 
-  // TODO-APP: double check users can't return null in a component that will kick in here
+  // If cache node has no subTreeData and no data request we have to infinitely suspend as the data will likely flow in from another place.
+  // TODO-APP: double check users can't return null in a component that will kick in here.
   if (!childNode.subTreeData) {
     throw createInfinitePromise()
   }
 
   const subtree = (
-    <AppTreeContext.Provider
+    // The layout router context narrows down tree and childNodes at each level.
+    <LayoutRouterContext.Provider
       value={{
         tree: tree[1][parallelRouterKey],
         childNodes: childNode.parallelRoutes,
@@ -228,17 +282,21 @@ export function InnerLayoutRouter({
       }}
     >
       {childNode.subTreeData}
-    </AppTreeContext.Provider>
+    </LayoutRouterContext.Provider>
   )
 
-  // Ensure root layout is not wrapped in a div
+  // Ensure root layout is not wrapped in a div as the root layout renders `<html>`
   return rootLayoutIncluded ? (
-    <div ref={focusAndScrollRef}>{subtree}</div>
+    <div ref={focusAndScrollElementRef}>{subtree}</div>
   ) : (
     subtree
   )
 }
 
+/**
+ * Renders suspense boundary with the provided "loading" property as the fallback.
+ * If no loading property is provided it renders the children without a suspense boundary.
+ */
 function LoadingBoundary({
   children,
   loading,
@@ -253,6 +311,10 @@ function LoadingBoundary({
   return <>{children}</>
 }
 
+/**
+ * OuterLayoutRouter handles the current segment as well as <Offscreen> rendering of other segments.
+ * It can be rendered next to each other with a different `parallelRouterKey`, allowing for Parallel routes.
+ */
 export default function OuterLayoutRouter({
   parallelRouterKey,
   segmentPath,
@@ -266,23 +328,34 @@ export default function OuterLayoutRouter({
   loading: React.ReactNode | undefined
   rootLayoutIncluded: boolean
 }) {
-  const { childNodes, tree, url } = useContext(AppTreeContext)
+  const { childNodes, tree, url } = useContext(LayoutRouterContext)
 
+  // Get the current parallelRouter cache node
   let childNodesForParallelRouter = childNodes.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) {
     childNodes.set(parallelRouterKey, new Map())
     childNodesForParallelRouter = childNodes.get(parallelRouterKey)!
   }
 
-  // This relates to the segments in the current router
-  // tree[1].children[0] refers to tree.children.segment in the data format
+  // 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 childPropSegment = Array.isArray(childProp.segment)
     ? childProp.segment[1]
     : childProp.segment
-  const currentChildSegment =
-    (Array.isArray(treeSegment) ? treeSegment[1] : treeSegment) ??
-    childPropSegment
+
+  // 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 currentChildSegment = Array.isArray(treeSegment)
+    ? treeSegment[1]
+    : treeSegment
+
+  /**
+   * 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: string[] = [currentChildSegment]
 
   return (
@@ -294,6 +367,8 @@ export default function OuterLayoutRouter({
         : null} */}
       {preservedSegments.map((preservedSegment) => {
         return (
+          // Loading boundary is render for each segment to ensure they have their own loading state.
+          // The loading boundary is passed to the router during rendering to ensure it can be immediately rendered when suspending on a Flight fetch.
           <LoadingBoundary loading={loading} key={preservedSegment}>
             <InnerLayoutRouter
               parallelRouterKey={parallelRouterKey}