Colocate styles with special entries

packages/next/build/utils.ts

@@ -1095,7 +1095,9 @@ export const collectGenerateParams = async (
 ): Promise<GenerateParams> => {
   if (!Array.isArray(segment)) return generateParams
   const isLayout = !!segment[2]?.layout
-  const mod = await (isLayout ? segment[2]?.layout?.() : segment[2]?.page?.())
+  const mod = await (isLayout
+    ? segment[2]?.layout?.[0]?.()
+    : segment[2]?.page?.[0]?.())
   const config = collectAppConfig(mod)
 
   const result = {

packages/next/build/webpack/loaders/next-app-loader.ts

@@ -21,11 +21,11 @@ const PAGE_SEGMENT = 'page$'
 
 // TODO-APP: check if this can be narrowed.
 type ComponentModule = () => any
+type ModuleReference = [componentModule: ComponentModule, filePath: string]
 export type ComponentsType = {
-  readonly [componentKey in ValueOf<typeof FILE_TYPES>]?: ComponentModule
+  readonly [componentKey in ValueOf<typeof FILE_TYPES>]?: ModuleReference
 } & {
-  readonly layoutOrPagePath?: string
-  readonly page?: ComponentModule
+  readonly page?: ModuleReference
 }
 
 async function createTreeCodeFromPath({
@@ -67,9 +67,10 @@ async function createTreeCodeFromPath({
         if (resolvedPagePath) pages.push(resolvedPagePath)
 
         // Use '' for segment as it's the page. There can't be a segment called '' so this is the safest way to add it.
-        props[parallelKey] = `['', {}, {layoutOrPagePath: ${JSON.stringify(
-          resolvedPagePath
-        )}, page: () => require(${JSON.stringify(resolvedPagePath)})}]`
+        props[parallelKey] = `['', {}, {
+          page: [() => require(${JSON.stringify(
+            resolvedPagePath
+          )}), ${JSON.stringify(resolvedPagePath)}]}]`
         continue
       }
 
@@ -105,11 +106,9 @@ async function createTreeCodeFromPath({
               if (filePath === undefined) {
                 return ''
               }
-              return `${
-                file === FILE_TYPES.layout
-                  ? `layoutOrPagePath: ${JSON.stringify(filePath)},`
-                  : ''
-              }'${file}': () => require(${JSON.stringify(filePath)}),`
+              return `'${file}': [() => require(${JSON.stringify(
+                filePath
+              )}), ${JSON.stringify(filePath)}],`
             })
             .join('\n')}
         }

packages/next/build/webpack/plugins/flight-client-entry-plugin.ts

@@ -179,17 +179,18 @@ export class FlightClientEntryPlugin {
       for (const connection of compilation.moduleGraph.getOutgoingConnections(
         entryModule
       )) {
-        const layoutOrPageDependency = connection.dependency
-        const layoutOrPageRequest = connection.dependency.request
+        // Entry can be any user defined entry files such as layout, page, error, loading, etc.
+        const entryDependency = connection.dependency
+        const entryRequest = connection.dependency.request
 
         const [clientComponentImports] =
           this.collectClientComponentsAndCSSForDependency({
-            layoutOrPageRequest,
+            entryRequest,
             compilation,
-            dependency: layoutOrPageDependency,
+            dependency: entryDependency,
           })
 
-        const isAbsoluteRequest = path.isAbsolute(layoutOrPageRequest)
+        const isAbsoluteRequest = path.isAbsolute(entryRequest)
 
         // Next.js internals are put into a separate entry.
         if (!isAbsoluteRequest) {
@@ -200,8 +201,8 @@ export class FlightClientEntryPlugin {
         }
 
         const relativeRequest = isAbsoluteRequest
-          ? path.relative(compilation.options.context, layoutOrPageRequest)
-          : layoutOrPageRequest
+          ? path.relative(compilation.options.context, entryRequest)
+          : entryRequest
 
         // Replace file suffix as `.js` will be added.
         const bundlePath = relativeRequest.replace(/\.(js|ts)x?$/, '')
@@ -308,14 +309,14 @@ export class FlightClientEntryPlugin {
         for (const connection of compilation.moduleGraph.getOutgoingConnections(
           entryModule
         )) {
-          const layoutOrPageDependency = connection.dependency
-          const layoutOrPageRequest = connection.dependency.request
+          const entryDependency = connection.dependency
+          const entryRequest = connection.dependency.request
 
           const [, cssImports] =
             this.collectClientComponentsAndCSSForDependency({
-              layoutOrPageRequest,
+              entryRequest,
               compilation,
-              dependency: layoutOrPageDependency,
+              dependency: entryDependency,
               clientEntryDependencyMap,
             })
 
@@ -353,12 +354,12 @@ export class FlightClientEntryPlugin {
   }
 
   collectClientComponentsAndCSSForDependency({
-    layoutOrPageRequest,
+    entryRequest,
     compilation,
     dependency,
     clientEntryDependencyMap,
   }: {
-    layoutOrPageRequest: string
+    entryRequest: string
     compilation: any
     dependency: any /* Dependency */
     clientEntryDependencyMap?: Record<string, any>
@@ -397,15 +398,15 @@ export class FlightClientEntryPlugin {
         : mod.resourceResolveData?.path + mod.resourceResolveData?.query
 
       // Ensure module is not walked again if it's already been visited
-      if (!visitedBySegment[layoutOrPageRequest]) {
-        visitedBySegment[layoutOrPageRequest] = new Set()
+      if (!visitedBySegment[entryRequest]) {
+        visitedBySegment[entryRequest] = new Set()
       }
       const storeKey =
         (inClientComponentBoundary ? '0' : '1') + ':' + modRequest
-      if (!modRequest || visitedBySegment[layoutOrPageRequest].has(storeKey)) {
+      if (!modRequest || visitedBySegment[entryRequest].has(storeKey)) {
         return
       }
-      visitedBySegment[layoutOrPageRequest].add(storeKey)
+      visitedBySegment[entryRequest].add(storeKey)
 
       const isClientComponent = isClientComponentModule(mod)
 
@@ -425,9 +426,8 @@ export class FlightClientEntryPlugin {
           }
         }
 
-        serverCSSImports[layoutOrPageRequest] =
-          serverCSSImports[layoutOrPageRequest] || []
-        serverCSSImports[layoutOrPageRequest].push(modRequest)
+        serverCSSImports[entryRequest] = serverCSSImports[entryRequest] || []
+        serverCSSImports[entryRequest].push(modRequest)
       }
 
       // Check if request is for css file.

packages/next/client/components/error-boundary.tsx

@@ -6,6 +6,7 @@ export type ErrorComponent = React.ComponentType<{
 }>
 interface ErrorBoundaryProps {
   errorComponent: ErrorComponent
+  errorStyles?: React.ReactNode | undefined
 }
 
 /**
@@ -32,10 +33,13 @@ class ErrorBoundaryHandler extends React.Component<
   render() {
     if (this.state.error) {
       return (
-        <this.props.errorComponent
-          error={this.state.error}
-          reset={this.reset}
-        />
+        <>
+          {this.props.errorStyles}
+          <this.props.errorComponent
+            error={this.state.error}
+            reset={this.reset}
+          />
+        </>
       )
     }
 
@@ -49,11 +53,15 @@ class ErrorBoundaryHandler extends React.Component<
  */
 export function ErrorBoundary({
   errorComponent,
+  errorStyles,
   children,
 }: ErrorBoundaryProps & { children: React.ReactNode }): JSX.Element {
   if (errorComponent) {
     return (
-      <ErrorBoundaryHandler errorComponent={errorComponent}>
+      <ErrorBoundaryHandler
+        errorComponent={errorComponent}
+        errorStyles={errorStyles}
+      >
         {children}
       </ErrorBoundaryHandler>
     )

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

@@ -257,15 +257,27 @@ export function InnerLayoutRouter({
 function LoadingBoundary({
   children,
   loading,
+  loadingStyles,
   hasLoading,
 }: {
   children: React.ReactNode
   loading?: React.ReactNode
+  loadingStyles?: React.ReactNode
   hasLoading: boolean
 }): JSX.Element {
   if (hasLoading) {
-    // @ts-expect-error TODO-APP: React.Suspense fallback type is wrong
-    return <React.Suspense fallback={loading}>{children}</React.Suspense>
+    return (
+      <React.Suspense
+        fallback={
+          <>
+            {loadingStyles}
+            {loading}
+          </>
+        }
+      >
+        {children}
+      </React.Suspense>
+    )
   }
 
   return <>{children}</>
@@ -322,6 +334,7 @@ function RedirectBoundary({ children }: { children: React.ReactNode }) {
 
 interface NotFoundBoundaryProps {
   notFound?: React.ReactNode
+  notFoundStyles?: React.ReactNode
   children: React.ReactNode
 }
 
@@ -347,6 +360,7 @@ class NotFoundErrorBoundary extends React.Component<
       return (
         <>
           <meta name="robots" content="noindex" />
+          {this.props.notFoundStyles}
           {this.props.notFound}
         </>
       )
@@ -356,9 +370,13 @@ class NotFoundErrorBoundary extends React.Component<
   }
 }
 
-function NotFoundBoundary({ notFound, children }: NotFoundBoundaryProps) {
+function NotFoundBoundary({
+  notFound,
+  notFoundStyles,
+  children,
+}: NotFoundBoundaryProps) {
   return notFound ? (
-    <NotFoundErrorBoundary notFound={notFound}>
+    <NotFoundErrorBoundary notFound={notFound} notFoundStyles={notFoundStyles}>
       {children}
     </NotFoundErrorBoundary>
   ) : (
@@ -375,20 +393,28 @@ export default function OuterLayoutRouter({
   segmentPath,
   childProp,
   error,
+  errorStyles,
+  templateStyles,
   loading,
+  loadingStyles,
   hasLoading,
   template,
   notFound,
+  notFoundStyles,
   rootLayoutIncluded,
 }: {
   parallelRouterKey: string
   segmentPath: FlightSegmentPath
   childProp: ChildProp
   error: ErrorComponent
+  errorStyles: React.ReactNode | undefined
+  templateStyles: React.ReactNode | undefined
   template: React.ReactNode
   loading: React.ReactNode | undefined
+  loadingStyles: React.ReactNode | undefined
   hasLoading: boolean
   notFound: React.ReactNode | undefined
+  notFoundStyles: React.ReactNode | undefined
   rootLayoutIncluded: boolean
 }) {
   const context = useContext(LayoutRouterContext)
@@ -442,9 +468,16 @@ export default function OuterLayoutRouter({
           <TemplateContext.Provider
             key={preservedSegment}
             value={
-              <ErrorBoundary errorComponent={error}>
-                <LoadingBoundary hasLoading={hasLoading} loading={loading}>
-                  <NotFoundBoundary notFound={notFound}>
+              <ErrorBoundary errorComponent={error} errorStyles={errorStyles}>
+                <LoadingBoundary
+                  hasLoading={hasLoading}
+                  loading={loading}
+                  loadingStyles={loadingStyles}
+                >
+                  <NotFoundBoundary
+                    notFound={notFound}
+                    notFoundStyles={notFoundStyles}
+                  >
                     <RedirectBoundary>
                       <InnerLayoutRouter
                         parallelRouterKey={parallelRouterKey}
@@ -467,7 +500,10 @@ export default function OuterLayoutRouter({
               </ErrorBoundary>
             }
           >
-            {template}
+            <>
+              {templateStyles}
+              {template}
+            </>
           </TemplateContext.Provider>
         )
       })}