Add render prop support to <Main>

packages/next/pages/_document.tsx

@@ -1,8 +1,5 @@
 import React, { Component, ReactElement, ReactNode, useContext } from 'react'
-import {
-  BODY_RENDER_TARGET,
-  OPTIMIZED_FONT_PROVIDERS,
-} from '../shared/lib/constants'
+import { OPTIMIZED_FONT_PROVIDERS } from '../shared/lib/constants'
 import {
   DocumentContext,
   DocumentInitialProps,
@@ -763,13 +760,18 @@ export class Head extends Component<
   }
 }
 
-export function Main() {
-  const { inAmpMode, docComponentsRendered } = useContext(HtmlContext)
-
+export function Main({
+  children,
+}: {
+  children?: (content: JSX.Element) => JSX.Element
+}) {
+  const { inAmpMode, docComponentsRendered, useMainContent } =
+    useContext(HtmlContext)
+  const content = useMainContent(children)
   docComponentsRendered.Main = true
 
-  if (inAmpMode) return <>{BODY_RENDER_TARGET}</>
-  return <div id="__next">{BODY_RENDER_TARGET}</div>
+  if (inAmpMode) return content
+  return <div id="__next">{content}</div>
 }
 
 export class NextScript extends Component<OriginProps> {

packages/next/server/render.tsx

@@ -20,7 +20,6 @@ import { GetServerSideProps, GetStaticProps, PreviewData } from '../types'
 import { isInAmpMode } from '../shared/lib/amp'
 import { AmpStateContext } from '../shared/lib/amp-context'
 import {
-  BODY_RENDER_TARGET,
   SERVER_PROPS_ID,
   STATIC_PROPS_ID,
   STATIC_STATUS_PAGES,
@@ -997,33 +996,55 @@ export async function renderToHTML(
       }
 
       return {
-        bodyResult: piperFromArray([docProps.html]),
+        bodyResult: () => piperFromArray([docProps.html]),
         documentElement: (htmlProps: HtmlProps) => (
           <Document {...htmlProps} {...docProps} />
         ),
+        useMainContent: (fn?: (content: JSX.Element) => JSX.Element) => {
+          if (fn) {
+            throw new Error(
+              'The `children` property is not supported by non-functional custom Document components'
+            )
+          }
+          // @ts-ignore
+          return <next-js-internal-body-render-target />
+        },
         head: docProps.head,
         headTags: await headTags(documentCtx),
         styles: docProps.styles,
       }
     } else {
-      const content =
-        ctx.err && ErrorDebug ? (
-          <ErrorDebug error={ctx.err} />
-        ) : (
-          <AppContainer>
-            <App {...props} Component={Component} router={router} />
-          </AppContainer>
-        )
-
-      const bodyResult = concurrentFeatures
-        ? process.browser
-          ? await renderToReadableStream(content)
-          : await renderToNodeStream(content, generateStaticHTML)
-        : piperFromArray([ReactDOMServer.renderToString(content)])
+      const contentWrappers: Array<(content: JSX.Element) => JSX.Element> = []
+      const bodyResult = async () => {
+        const initialContent =
+          ctx.err && ErrorDebug ? (
+            <ErrorDebug error={ctx.err} />
+          ) : (
+            <AppContainer>
+              <App {...props} Component={Component} router={router} />
+            </AppContainer>
+          )
+        const content = contentWrappers.reduce((innerContent, fn) => {
+          return fn(innerContent)
+        }, initialContent)
+
+        return concurrentFeatures
+          ? process.browser
+            ? await renderToReadableStream(content)
+            : await renderToNodeStream(content, generateStaticHTML)
+          : piperFromArray([ReactDOMServer.renderToString(content)])
+      }
 
       return {
         bodyResult,
         documentElement: () => (Document as any)(),
+        useMainContent: (fn?: (content: JSX.Element) => JSX.Element) => {
+          if (fn) {
+            contentWrappers.push(fn)
+          }
+          // @ts-ignore
+          return <next-js-internal-body-render-target />
+        },
         head,
         headTags: [],
         styles: jsxStyleRegistry.styles(),
@@ -1051,8 +1072,8 @@ export async function renderToHTML(
   }
 
   const hybridAmp = ampState.hybrid
-
   const docComponentsRendered: DocumentProps['docComponentsRendered'] = {}
+
   const {
     assetPrefix,
     buildId,
@@ -1118,6 +1139,7 @@ export async function renderToHTML(
     head: documentResult.head,
     headTags: documentResult.headTags,
     styles: documentResult.styles,
+    useMainContent: documentResult.useMainContent,
     useMaybeDeferContent,
   }
   const documentHTML = ReactDOMServer.renderToStaticMarkup(
@@ -1151,20 +1173,20 @@ export async function renderToHTML(
     }
   }
 
-  const renderTargetIdx = documentHTML.indexOf(BODY_RENDER_TARGET)
+  const [renderTargetPrefix, renderTargetSuffix] = documentHTML.split(
+    /<next-js-internal-body-render-target><\/next-js-internal-body-render-target>/
+  )
   const prefix: Array<string> = []
   prefix.push('<!DOCTYPE html>')
-  prefix.push(documentHTML.substring(0, renderTargetIdx))
+  prefix.push(renderTargetPrefix)
   if (inAmpMode) {
     prefix.push('<!-- __NEXT_DATA__ -->')
   }
 
   let pipers: Array<NodeWritablePiper> = [
     piperFromArray(prefix),
-    documentResult.bodyResult,
-    piperFromArray([
-      documentHTML.substring(renderTargetIdx + BODY_RENDER_TARGET.length),
-    ]),
+    await documentResult.bodyResult(),
+    piperFromArray([renderTargetSuffix]),
   ]
 
   const postProcessors: Array<((html: string) => Promise<string>) | null> = (

packages/next/shared/lib/constants.ts

@@ -23,7 +23,6 @@ export const BLOCKED_PAGES = ['/_document', '/_app', '/_error']
 export const CLIENT_PUBLIC_FILES_PATH = 'public'
 export const CLIENT_STATIC_FILES_PATH = 'static'
 export const CLIENT_STATIC_FILES_RUNTIME = 'runtime'
-export const BODY_RENDER_TARGET = '__NEXT_BODY_RENDER_TARGET__'
 export const STRING_LITERAL_DROP_BUNDLE = '__NEXT_DROP_CLIENT_FILE__'
 
 // server/middleware-flight-manifest.js

packages/next/shared/lib/utils.ts

@@ -221,6 +221,7 @@ export type HtmlProps = {
   styles?: React.ReactElement[] | React.ReactFragment
   head?: Array<JSX.Element | null>
   useMaybeDeferContent: MaybeDeferContentHook
+  useMainContent: (fn?: (content: JSX.Element) => JSX.Element) => JSX.Element
 }
 
 /**

test/integration/react-18/app/components/context.js

@@ -0,0 +1,2 @@
+import { createContext } from 'react'
+export default createContext(null)

test/integration/react-18/app/pages/_document.js

@@ -1,11 +1,18 @@
 import { Html, Head, Main, NextScript } from 'next/document'
+import Context from '../components/context'
 
 export default function Document() {
   return (
     <Html>
       <Head />
       <body>
-        <Main />
+        <Main>
+          {(content) => (
+            <Context.Provider value="from main render prop">
+              {content}
+            </Context.Provider>
+          )}
+        </Main>
         <NextScript />
       </body>
     </Html>

test/integration/react-18/app/pages/main-render-prop.js

@@ -0,0 +1,7 @@
+import { useContext } from 'react'
+import Context from '../components/context'
+
+export default function MainRenderProp() {
+  const value = useContext(Context)
+  return value
+}

test/integration/react-18/test/index.test.js

@@ -134,6 +134,15 @@ describe('Basics', () => {
       'A React component suspended while rendering, but no fallback UI was specified'
     )
   })
+
+  it('supports render prop as children in <Main />', async () => {
+    const appPort = await findPort()
+    const app = await launchApp(appDir, appPort, { nodeArgs })
+    const html = await renderViaHTTP(appPort, '/main-render-prop')
+    await killApp(app)
+
+    expect(html).toContain('from main render prop')
+  })
 })
 
 describe('Blocking mode', () => {