Clean up AMP bundle removal

packages/next/build/babel/plugins/next-page-config.ts

@@ -1,7 +1,6 @@
 import { NodePath, PluginObj, types as BabelTypes } from '@babel/core'
 import { PageConfig } from 'next/types'
-
-const STRING_LITERAL_DROP_BUNDLE = '__NEXT_DROP_CLIENT_FILE__'
+import { STRING_LITERAL_DROP_BUNDLE } from '../../../next-server/lib/constants'
 
 // replace program path with just a variable with the drop identifier
 function replaceBundle(path: any, t: typeof BabelTypes): void {
@@ -10,12 +9,8 @@ function replaceBundle(path: any, t: typeof BabelTypes): void {
       [
         t.variableDeclaration('const', [
           t.variableDeclarator(
-            t.identifier('config'),
-            t.assignmentExpression(
-              '=',
-              t.identifier(STRING_LITERAL_DROP_BUNDLE),
-              t.stringLiteral(`${STRING_LITERAL_DROP_BUNDLE} ${Date.now()}`)
-            )
+            t.identifier(STRING_LITERAL_DROP_BUNDLE),
+            t.stringLiteral(`${STRING_LITERAL_DROP_BUNDLE} ${Date.now()}`)
           ),
         ]),
       ],

packages/next/build/index.ts

@@ -572,30 +572,6 @@ export default async function build(dir: string, conf = null): Promise<void> {
             hybridAmpPages.add(page)
           }
 
-          if (workerResult.isAmpOnly) {
-            // ensure all AMP only bundles got removed
-            try {
-              const clientBundle = path.join(
-                distDir,
-                'static',
-                buildId,
-                'pages',
-                actualPage + '.js'
-              )
-              await promises.unlink(clientBundle)
-
-              if (config.experimental.modern) {
-                await promises.unlink(
-                  clientBundle.replace(/\.js$/, '.module.js')
-                )
-              }
-            } catch (err) {
-              if (err.code !== 'ENOENT') {
-                throw err
-              }
-            }
-          }
-
           if (workerResult.hasStaticProps) {
             ssgPages.add(page)
             isSsg = true
@@ -719,19 +695,6 @@ export default async function build(dir: string, conf = null): Promise<void> {
     )
   }
 
-  if (Array.isArray(configs[0].plugins)) {
-    configs[0].plugins.some((plugin: any) => {
-      if (!plugin.ampPages) {
-        return false
-      }
-
-      plugin.ampPages.forEach((pg: any) => {
-        pageInfos.get(pg)!.isAmp = true
-      })
-      return true
-    })
-  }
-
   await writeBuildId(distDir, buildId)
 
   const finalPrerenderRoutes: { [route: string]: SsgRoute } = {}

packages/next/build/utils.ts

@@ -23,6 +23,7 @@ import { isDynamicRoute } from '../next-server/lib/router/utils/is-dynamic'
 import { findPageFile } from '../server/lib/find-page-file'
 import { GetStaticPaths } from 'next/types'
 import { denormalizePagePath } from '../next-server/server/normalize-page-path'
+import { BuildManifest } from '../next-server/server/get-page-files'
 
 const fileGzipStats: { [k: string]: Promise<number> } = {}
 const fsStatGzip = (file: string) => {
@@ -42,7 +43,6 @@ export function collectPages(
 }
 
 export interface PageInfo {
-  isAmp?: boolean
   isHybridAmp?: boolean
   size: number
   totalSize: number
@@ -70,7 +70,7 @@ export async function printTreeView(
     buildId: string
     pagesDir: string
     pageExtensions: string[]
-    buildManifest: BuildManifestShape
+    buildManifest: BuildManifest
     isModern: boolean
     useStatic404: boolean
   }
@@ -141,6 +141,7 @@ export async function printTreeView(
         : '├'
 
     const pageInfo = pageInfos.get(item)
+    const ampFirst = buildManifest.ampFirstPages.includes(item)
 
     messages.push([
       `${symbol} ${
@@ -153,14 +154,14 @@ export async function printTreeView(
           : 'λ'
       } ${item}`,
       pageInfo
-        ? pageInfo.isAmp
+        ? ampFirst
           ? chalk.cyan('AMP')
           : pageInfo.size >= 0
           ? prettyBytes(pageInfo.size)
           : ''
         : '',
       pageInfo
-        ? pageInfo.isAmp
+        ? ampFirst
           ? chalk.cyan('AMP')
           : pageInfo.size >= 0
           ? getPrettySize(pageInfo.totalSize)
@@ -348,7 +349,6 @@ export function printCustomRoutes({
   }
 }
 
-type BuildManifestShape = { pages: { [k: string]: string[] } }
 type ComputeManifestShape = {
   commonFiles: string[]
   uniqueFiles: string[]
@@ -357,14 +357,14 @@ type ComputeManifestShape = {
   sizeCommonFiles: number
 }
 
-let cachedBuildManifest: BuildManifestShape | undefined
+let cachedBuildManifest: BuildManifest | undefined
 
 let lastCompute: ComputeManifestShape | undefined
 let lastComputeModern: boolean | undefined
 let lastComputePageInfo: boolean | undefined
 
 async function computeFromManifest(
-  manifest: BuildManifestShape,
+  manifest: BuildManifest,
   distPath: string,
   isModern: boolean,
   pageInfos?: Map<string, PageInfo>
@@ -383,7 +383,8 @@ async function computeFromManifest(
     if (pageInfos) {
       const pageInfo = pageInfos.get(key)
       // don't include AMP pages since they don't rely on shared bundles
-      if (pageInfo?.isHybridAmp || pageInfo?.isAmp) {
+      // AMP First pages are not under the pageInfos key
+      if (pageInfo?.isHybridAmp) {
         return
       }
     }
@@ -480,7 +481,7 @@ function sum(a: number[]): number {
 export async function getJsPageSizeInKb(
   page: string,
   distPath: string,
-  buildManifest: BuildManifestShape,
+  buildManifest: BuildManifest,
   isModern: boolean
 ): Promise<[number, number]> {
   const data = await computeFromManifest(buildManifest, distPath, isModern)

packages/next/build/webpack/plugins/build-manifest-plugin.ts

@@ -10,6 +10,7 @@ import {
 } from '../../../next-server/lib/constants'
 import { BuildManifest } from '../../../next-server/server/get-page-files'
 import getRouteFromEntrypoint from '../../../next-server/server/get-route-from-entrypoint'
+import { ampFirstEntryNamesMap } from './next-drop-client-page-plugin'
 
 // This function takes the asset map generated in BuildManifestPlugin and creates a
 // reduced version to send to the client.
@@ -63,6 +64,19 @@ export default class BuildManifestPlugin {
           devFiles: [],
           lowPriorityFiles: [],
           pages: { '/_app': [] },
+          ampFirstPages: [],
+        }
+
+        const ampFirstEntryNames = ampFirstEntryNamesMap.get(compilation)
+        if (ampFirstEntryNames) {
+          for (const entryName of ampFirstEntryNames) {
+            const pagePath = getRouteFromEntrypoint(entryName)
+            if (!pagePath) {
+              continue
+            }
+
+            assetMap.ampFirstPages.push(pagePath)
+          }
         }
 
         const mainJsChunk = chunks.find(

packages/next/build/webpack/plugins/next-drop-client-page-plugin.ts

@@ -1,29 +1,94 @@
-import { Compiler, Plugin } from 'webpack'
-import { extname } from 'path'
+import { Compiler, compilation as CompilationType, Plugin } from 'webpack'
+import { STRING_LITERAL_DROP_BUNDLE } from '../../../next-server/lib/constants'
+
+export const ampFirstEntryNamesMap: WeakMap<
+  CompilationType.Compilation,
+  string[]
+> = new WeakMap()
+
+const PLUGIN_NAME = 'DropAmpFirstPagesPlugin'
+
+// Recursively look up the issuer till it ends up at the root
+function findEntryModule(mod: any): CompilationType.Module | null {
+  const queue = new Set([mod])
+  for (const module of queue) {
+    for (const reason of module.reasons) {
+      if (!reason.module) return module
+      queue.add(reason.module)
+    }
+  }
+
+  return null
+}
+
+function handler(parser: any) {
+  function markAsAmpFirst() {
+    const entryModule = findEntryModule(parser.state.module)
+
+    if (!entryModule) {
+      return
+    }
+
+    // @ts-ignore buildInfo exists on Module
+    entryModule.buildInfo.NEXT_ampFirst = true
+  }
+
+  parser.hooks.varDeclarationConst
+    .for(STRING_LITERAL_DROP_BUNDLE)
+    .tap(PLUGIN_NAME, markAsAmpFirst)
+
+  parser.hooks.varDeclarationLet
+    .for(STRING_LITERAL_DROP_BUNDLE)
+    .tap(PLUGIN_NAME, markAsAmpFirst)
+
+  parser.hooks.varDeclaration
+    .for(STRING_LITERAL_DROP_BUNDLE)
+    .tap(PLUGIN_NAME, markAsAmpFirst)
+}
 
 // Prevents outputting client pages when they are not needed
 export class DropClientPage implements Plugin {
   ampPages = new Set()
 
   apply(compiler: Compiler) {
-    compiler.hooks.emit.tap('DropClientPage', (compilation) => {
-      Object.keys(compilation.assets).forEach((assetKey) => {
-        const asset = compilation.assets[assetKey]
+    compiler.hooks.compilation.tap(
+      PLUGIN_NAME,
+      (compilation, { normalModuleFactory }) => {
+        normalModuleFactory.hooks.parser
+          .for('javascript/auto')
+          .tap(PLUGIN_NAME, handler)
 
-        if (asset?._value?.includes?.('__NEXT_DROP_CLIENT_FILE__')) {
-          const cleanAssetKey = assetKey.replace(/\\/g, '/')
-          const page = '/' + cleanAssetKey.split('pages/')[1]
-          const pageNoExt = page.split(extname(page))[0]
+        if (!ampFirstEntryNamesMap.has(compilation)) {
+          ampFirstEntryNamesMap.set(compilation, [])
+        }
 
-          delete compilation.assets[assetKey]
+        const ampFirstEntryNamesItem = ampFirstEntryNamesMap.get(
+          compilation
+        ) as string[]
 
-          // Detect being re-ran through a child compiler and don't re-mark the
-          // page as AMP
-          if (!pageNoExt.endsWith('.module')) {
-            this.ampPages.add(pageNoExt.replace(/\/index$/, '') || '/')
+        compilation.hooks.seal.tap(PLUGIN_NAME, () => {
+          // Remove preparedEntrypoint that has bundle drop marker
+          // This will ensure webpack does not create chunks/bundles for this particular entrypoint
+          for (
+            let i = compilation._preparedEntrypoints.length - 1;
+            i >= 0;
+            i--
+          ) {
+            const entrypoint = compilation._preparedEntrypoints[i]
+            if (entrypoint?.module?.buildInfo?.NEXT_ampFirst) {
+              ampFirstEntryNamesItem.push(entrypoint.name)
+              compilation._preparedEntrypoints.splice(i, 1)
+            }
           }
-        }
-      })
-    })
+
+          for (let i = compilation.entries.length - 1; i >= 0; i--) {
+            const entryModule = compilation.entries[i]
+            if (entryModule?.buildInfo?.NEXT_ampFirst) {
+              compilation.entries.splice(i, 1)
+            }
+          }
+        })
+      }
+    )
   }
 }

packages/next/next-server/lib/constants.ts

@@ -18,6 +18,7 @@ export const CLIENT_PUBLIC_FILES_PATH = 'public'
 export const CLIENT_STATIC_FILES_PATH = 'static'
 export const CLIENT_STATIC_FILES_RUNTIME = 'runtime'
 export const AMP_RENDER_TARGET = '__NEXT_AMP_RENDER_TARGET__'
+export const STRING_LITERAL_DROP_BUNDLE = '__NEXT_DROP_CLIENT_FILE__'
 export const CLIENT_STATIC_FILES_RUNTIME_PATH = `${CLIENT_STATIC_FILES_PATH}/${CLIENT_STATIC_FILES_RUNTIME}`
 // static/runtime/main.js
 export const CLIENT_STATIC_FILES_RUNTIME_MAIN = `${CLIENT_STATIC_FILES_RUNTIME_PATH}/main.js`

packages/next/next-server/server/get-page-files.ts

@@ -8,6 +8,7 @@ export type BuildManifest = {
     '/_app': string[]
     [page: string]: string[]
   }
+  ampFirstPages: string[]
 }
 
 export function getPageFiles(

packages/next/next-server/server/render.tsx

@@ -652,12 +652,17 @@ export async function renderToHTML(
   // the response might be finished on the getInitialProps call
   if (isResSent(res) && !isSSG) return null
 
-  const files = [
-    ...new Set([
-      ...getPageFiles(buildManifest, '/_app'),
-      ...(pathname !== '/_error' ? getPageFiles(buildManifest, pathname) : []),
-    ]),
-  ]
+  // AMP First pages do not have client-side JavaScript files
+  const files = ampState.ampFirst
+    ? []
+    : [
+        ...new Set([
+          ...getPageFiles(buildManifest, '/_app'),
+          ...(pathname !== '/_error'
+            ? getPageFiles(buildManifest, pathname)
+            : []),
+        ]),
+      ]
 
   const renderPage: RenderPage = (
     options: ComponentsEnhancer = {}

packages/next/types/misc.d.ts

@@ -11,6 +11,7 @@ declare module 'styled-jsx/server'
 declare module 'unfetch'
 declare module 'webpack/lib/GraphHelpers'
 declare module 'webpack/lib/DynamicEntryPlugin'
+declare module 'webpack/lib/Entrypoint'
 
 declare module 'next/dist/compiled/amphtml-validator' {
   import m from 'amphtml-validator'