From 75a9dfac35006088b64c326fe054d8c7e6b9eaa6 Mon Sep 17 00:00:00 2001
From: Jiachi Liu <inbox@huozhi.im>
Date: Tue, 24 Dec 2024 00:28:20 +0100
Subject: [PATCH] refactor: error boundary rendering in app-render (#74259)

---
 .../app-render/create-component-tree.tsx      | 138 ++++++++++--------
 1 file changed, 75 insertions(+), 63 deletions(-)

diff --git a/packages/next/src/server/app-render/create-component-tree.tsx b/packages/next/src/server/app-render/create-component-tree.tsx
index dfbe7b9598b6b..9c0bf21789b62 100644
--- a/packages/next/src/server/app-render/create-component-tree.tsx
+++ b/packages/next/src/server/app-render/create-component-tree.tsx
@@ -209,12 +209,6 @@ async function createComponentTreeInternal({
           injectedJS: injectedJSWithCurrentLayout,
         })
       : []
-  const forbiddenElement = Forbidden ? (
-    <>
-      {forbiddenStyles}
-      <Forbidden />
-    </>
-  ) : undefined
 
   const [Unauthorized, unauthorizedStyles] =
     authInterrupts && unauthorized
@@ -226,6 +220,21 @@ async function createComponentTreeInternal({
           injectedJS: injectedJSWithCurrentLayout,
         })
       : []
+
+  const notFoundElement = NotFound ? (
+    <>
+      {notFoundStyles}
+      <NotFound />
+    </>
+  ) : undefined
+
+  const forbiddenElement = Forbidden ? (
+    <>
+      {forbiddenStyles}
+      <Forbidden />
+    </>
+  ) : undefined
+
   const unauthorizedElement = Unauthorized ? (
     <>
       {unauthorizedStyles}
@@ -411,13 +420,9 @@ async function createComponentTreeInternal({
         const isChildrenRouteKey = parallelRouteKey === 'children'
         const parallelRoute = parallelRoutes[parallelRouteKey]
 
-        const notFoundComponent =
-          NotFound && isChildrenRouteKey ? (
-            <>
-              {notFoundStyles}
-              <NotFound />
-            </>
-          ) : undefined
+        const notFoundComponent = isChildrenRouteKey
+          ? notFoundElement
+          : undefined
 
         const forbiddenComponent = isChildrenRouteKey
           ? forbiddenElement
@@ -709,56 +714,30 @@ async function createComponentTreeInternal({
         // but it's not ideal, as it needlessly invokes the `NotFound` component and renders the `RootLayout` twice.
         // We should instead look into handling the fallback behavior differently in development mode so that it doesn't
         // rely on the `NotFound` behavior.
-        if (NotFound) {
-          const notFoundParallelRouteProps = {
-            children: (
-              <>
-                {notFoundStyles}
-                <NotFound />
-              </>
-            ),
-          }
-          notfoundClientSegment = (
-            <>
-              {layerAssets}
-              <ClientSegmentRoot
-                Component={SegmentComponent}
-                slots={notFoundParallelRouteProps}
-                params={currentParams}
-              />
-            </>
-          )
-        }
-        if (Forbidden) {
-          const forbiddenParallelRouteProps = {
-            children: forbiddenElement,
-          }
-          forbiddenClientSegment = (
-            <>
-              {layerAssets}
-              <ClientSegmentRoot
-                Component={SegmentComponent}
-                slots={forbiddenParallelRouteProps}
-                params={currentParams}
-              />
-            </>
-          )
-        }
-        if (Unauthorized) {
-          const unauthorizedParallelRouteProps = {
-            children: unauthorizedElement,
-          }
-          unauthorizedClientSegment = (
-            <>
-              {layerAssets}
-              <ClientSegmentRoot
-                Component={SegmentComponent}
-                slots={unauthorizedParallelRouteProps}
-                params={currentParams}
-              />
-            </>
-          )
-        }
+        notfoundClientSegment = createErrorBoundaryClientSegmentRoot({
+          ErrorBoundaryComponent: NotFound,
+          errorElement: notFoundElement,
+          ClientSegmentRoot,
+          layerAssets,
+          SegmentComponent,
+          currentParams,
+        })
+        forbiddenClientSegment = createErrorBoundaryClientSegmentRoot({
+          ErrorBoundaryComponent: Forbidden,
+          errorElement: forbiddenElement,
+          ClientSegmentRoot,
+          layerAssets,
+          SegmentComponent,
+          currentParams,
+        })
+        unauthorizedClientSegment = createErrorBoundaryClientSegmentRoot({
+          ErrorBoundaryComponent: Unauthorized,
+          errorElement: unauthorizedElement,
+          ClientSegmentRoot,
+          layerAssets,
+          SegmentComponent,
+          currentParams,
+        })
         if (
           notfoundClientSegment ||
           forbiddenClientSegment ||
@@ -861,3 +840,36 @@ async function MetadataOutlet({
   return null
 }
 MetadataOutlet.displayName = OUTLET_BOUNDARY_NAME
+
+function createErrorBoundaryClientSegmentRoot({
+  ErrorBoundaryComponent,
+  errorElement,
+  ClientSegmentRoot,
+  layerAssets,
+  SegmentComponent,
+  currentParams,
+}: {
+  ErrorBoundaryComponent: React.ComponentType<any> | undefined
+  errorElement: React.ReactNode
+  ClientSegmentRoot: React.ComponentType<any>
+  layerAssets: React.ReactNode
+  SegmentComponent: React.ComponentType<any>
+  currentParams: Params
+}) {
+  if (ErrorBoundaryComponent) {
+    const notFoundParallelRouteProps = {
+      children: errorElement,
+    }
+    return (
+      <>
+        {layerAssets}
+        <ClientSegmentRoot
+          Component={SegmentComponent}
+          slots={notFoundParallelRouteProps}
+          params={currentParams}
+        />
+      </>
+    )
+  }
+  return null
+}