fix merge css plugin to account for css order

packages/next/src/build/webpack/plugins/merge-css-chunks-plugin.ts

@@ -1,21 +1,15 @@
-import type { Chunk, ChunkGraph, Compiler } from 'webpack'
+import type { Chunk, ChunkGraph, Compiler, ChunkGroup } from 'webpack'
 
 const PLUGIN_NAME = 'MergeCssChunksPlugin'
 
 /**
- * Css chunks smaller than this size will be merged with other css chunks.
+ * Merge chunks until they are bigger than the target size.
  */
 const MIN_CSS_CHUNK_SIZE = 30 * 1024
 /**
- * When merging css chunks it will select an number N where the total size is just bigger than this size.
- * Exception: N must be at least 2 even if the size is already bigger than this size.
+ * Avoid merging chunks when they would be bigger than this size.
  */
-const TARGET_CSS_CHUNK_SIZE = 60 * 1024
-/**
- * When an entrypoint has more css chunks than this number it merge the smallest ones to try to stay below that number.
- * Exception: When there are more than twice as much css chunks that larger than MAX_CSS_CHUNK_SIZE it will only half the number of css chunks.
- */
-const MAX_CSS_CHUNKS = 15
+const MAX_CSS_CHUNK_SIZE = 100 * 1024
 
 function isCssChunk(chunkGraph: ChunkGraph, chunk: Chunk): boolean {
   for (const mod of chunkGraph.getChunkModulesIterable(chunk)) {
@@ -24,6 +18,22 @@ function isCssChunk(chunkGraph: ChunkGraph, chunk: Chunk): boolean {
   return false
 }
 
+type CssChunkInfo = {
+  entrypoints: Map<
+    ChunkGroup,
+    {
+      prev: Chunk | null
+      next: Chunk | null
+    }
+  >
+  size: number
+}
+type ChunkGroupItem = {
+  entrypoint: ChunkGroup
+  cssChunks: Set<Chunk>
+  size: number
+}
+
 export class MergeCssChunksPlugin {
   public apply(compiler: Compiler) {
     compiler.hooks.thisCompilation.tap(PLUGIN_NAME, (compilation) => {
@@ -40,38 +50,134 @@ export class MergeCssChunksPlugin {
           once = true
           const chunkGraph = compilation.chunkGraph
           let changed = false
+
+          const cssChunkInfo = new Map<Chunk, CssChunkInfo>()
+          const cssChunksForChunkGroup = new Map<ChunkGroup, Set<Chunk>>()
+          const chunkGroups: ChunkGroupItem[] = []
+
+          // Add all css chunks to the maps and lists
           for (const [, entrypoint] of compilation.entrypoints) {
-            const cssChunks = entrypoint.chunks
-              .filter((chunk) => isCssChunk(chunkGraph, chunk))
-              .map((chunk) => [chunk, chunkGraph.getChunkSize(chunk)] as const)
-            cssChunks.sort((a, b) => b[1] - a[1])
-
-            // We want have at most MAX_CSS_CHUNKS chunks.
-            // When we start merging at startMergingIndex this would half the number of chunks after that index.
-            const startMergingIndex = 2 * MAX_CSS_CHUNKS - cssChunks.length
-
-            // We select small chunks and chunks after the index
-            const selectedCssChunks = cssChunks.filter(
-              ([, size], i) =>
-                size < MIN_CSS_CHUNK_SIZE || i >= startMergingIndex
+            const cssChunks = entrypoint.chunks.filter((chunk) =>
+              isCssChunk(chunkGraph, chunk)
             )
-            while (selectedCssChunks.length >= 2) {
-              const [biggest, bigSize] = selectedCssChunks.shift()!
-              let mergedSize = bigSize
-              do {
-                const [smallest, size] = selectedCssChunks.pop()!
-                if (chunkGraph.canChunksBeIntegrated(biggest, smallest)) {
-                  chunkGraph.integrateChunks(biggest, smallest)
-                  compilation.chunks.delete(smallest)
-                  mergedSize += size
-                  changed = true
+
+            if (cssChunks.length <= 1) continue
+
+            let totalSize = 0
+            const set = new Set<Chunk>()
+
+            function addChunk(
+              chunk: Chunk,
+              prev: Chunk | null,
+              next: Chunk | null
+            ) {
+              let info = cssChunkInfo.get(chunk)
+              if (info === undefined) {
+                info = {
+                  entrypoints: new Map(),
+                  size: chunkGraph.getChunkSize(chunk),
                 }
-              } while (
-                selectedCssChunks.length > 0 &&
-                mergedSize < TARGET_CSS_CHUNK_SIZE
-              )
+                cssChunkInfo.set(chunk, info)
+              }
+              info.entrypoints.set(entrypoint, { prev, next })
+              totalSize += info.size
+              if (info.size < MAX_CSS_CHUNK_SIZE) set.add(chunk)
+            }
+
+            const first = cssChunks[0]
+            const second = cssChunks[1]
+            addChunk(first, null, second)
+            for (let i = 2; i < cssChunks.length; i++) {
+              const prev = cssChunks[i - 2]
+              const current = cssChunks[i - 1]
+              const next = cssChunks[i]
+              addChunk(current, prev, next)
+            }
+            const last = cssChunks[cssChunks.length - 1]
+            addChunk(last, cssChunks[cssChunks.length - 2], null)
+
+            cssChunksForChunkGroup.set(entrypoint, set)
+            chunkGroups.push({ entrypoint, cssChunks: set, size: totalSize })
+          }
+
+          // Sorting for determinism, order isn't that relevant
+          chunkGroups.sort(
+            (a, b) => b.cssChunks.size - a.cssChunks.size || b.size - a.size
+          )
+
+          // Checks if two chunks can be merged without breaking css order and size limits
+          function canBeMerged(a: Chunk, b: Chunk) {
+            const aInfo = cssChunkInfo.get(a)!
+            const bInfo = cssChunkInfo.get(b)!
+            if (
+              aInfo.size > MIN_CSS_CHUNK_SIZE &&
+              bInfo.size > MIN_CSS_CHUNK_SIZE
+            ) {
+              // No need to merge them since they are already big enough
+              return false
+            }
+            if (aInfo.size + bInfo.size > MAX_CSS_CHUNK_SIZE) {
+              // Can't merge than as that would be too big
+              return false
+            }
+            for (const [entrypoint, infoInA] of aInfo.entrypoints) {
+              if (bInfo.entrypoints.has(entrypoint) && infoInA.next !== b) {
+                // Can't merge as they are both in the same entrypoint, but are not siblings
+                // Merging them would break the css order
+                return false
+              }
+            }
+            return chunkGraph.canChunksBeIntegrated(a, b)
+          }
+
+          /**
+           * Update datastructure to reflect that `removedChunk` is being merged into `newChunk`
+           */
+          function updateInfoForMerging(removedChunk: Chunk, newChunk: Chunk) {
+            const info = cssChunkInfo.get(removedChunk)!
+            for (const [entrypoint, entryInfo] of info.entrypoints) {
+              if (entryInfo.prev === newChunk) {
+                const prevInfo = cssChunkInfo.get(entryInfo.prev)!
+                prevInfo.entrypoints.get(entrypoint)!.next = entryInfo.next
+              } else if (entryInfo.prev) {
+                const prevInfo = cssChunkInfo.get(entryInfo.prev)!
+                prevInfo.entrypoints.get(entrypoint)!.next = newChunk
+              }
+              if (entryInfo.next) {
+                const nextInfo = cssChunkInfo.get(entryInfo.next)!
+                nextInfo.entrypoints.get(entrypoint)!.prev = newChunk
+              }
+              cssChunksForChunkGroup.get(entrypoint)!.delete(removedChunk)
             }
+            cssChunkInfo.delete(removedChunk)
+            return info
           }
+
+          // Merge chunks, until nothing more can be merged
+          let changedInIteration = false
+          do {
+            changedInIteration = false
+            for (const { entrypoint, cssChunks } of chunkGroups) {
+              for (const chunk of cssChunks) {
+                const info = cssChunkInfo.get(chunk)!
+                const entryInfo = info.entrypoints.get(entrypoint)!
+                const next = entryInfo.next
+                if (!next) continue
+
+                if (canBeMerged(chunk, next)) {
+                  // `next` will be removed, so we need to update the pointers to it from other chunks
+                  updateInfoForMerging(next, chunk)
+
+                  // Really merge it
+                  chunkGraph.integrateChunks(chunk, next)
+                  compilation.chunks.delete(next)
+                  changed = true
+                  changedInIteration = true
+                }
+              }
+            }
+          } while (changedInIteration)
+
           return changed
         }
       )