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
 }
 
 /**