feat(worker): remove worker chunk if it's detected that it's not referenced

packages/vite/src/node/plugins/worker.ts

@@ -1,6 +1,12 @@
 import path from 'node:path'
 import MagicString from 'magic-string'
-import type { PluginContext, RolldownOutput, RollupError } from 'rolldown'
+import type {
+  OutputAsset,
+  OutputChunk,
+  PluginContext,
+  RolldownOutput,
+  RollupError,
+} from 'rolldown'
 import colors from 'picocolors'
 import { type ImportSpecifier, init, parse } from 'es-module-lexer'
 import { viteWebWorkerPostPlugin as nativeWebWorkerPostPlugin } from 'rolldown/experimental'
@@ -32,6 +38,7 @@ type WorkerBundle = {
   entryCode: string
   entryUrlPlaceholder: string
   referencedAssets: Set<string>
+  moduleIds: Set<string>
   watchedFiles: string[]
 }
 
@@ -43,6 +50,12 @@ type WorkerBundleAsset = {
   source: string | Uint8Array
 }
 
+/** The input ID of a worker entry, which identifies its bundle. */
+type WorkerBundleId = string
+
+/** `undefined` identifies the main bundle. */
+type BundleId = WorkerBundleId | undefined
+
 class WorkerOutputCache {
   /**
    * worker bundle information for each input id
@@ -56,13 +69,29 @@ class WorkerOutputCache {
     /* entryFilename */ string
   >()
   private invalidatedBundles = new Set</* inputId */ string>()
+  /**
+   * Worker references grouped by their containing bundle and module.
+   * `referencingModuleId` is the module whose inclusion keeps the reference
+   * live: the worker wrapper module for a `?worker` import, or the importing
+   * module for a `new URL(..., import.meta.url)` worker reference.
+   * `childBundleId` identifies the referenced worker bundle and is used as its
+   * `BundleId` when traversing that bundle's references.
+   */
+  private bundleReferences = new Map<
+    BundleId,
+    Map<
+      /* referencingModuleId */ string,
+      Set</* childBundleId */ WorkerBundleId>
+    >
+  >()
 
   saveWorkerBundle(
     file: string,
     watchedFiles: string[],
     outputEntryFilename: string,
     outputEntryCode: string,
     outputAssets: WorkerBundleAsset[],
+    moduleIds: Set<string>,
     logger: Logger,
   ): WorkerBundle {
     for (const asset of outputAssets) {
@@ -74,6 +103,7 @@ class WorkerOutputCache {
       entryUrlPlaceholder:
         this.generateEntryUrlPlaceholder(outputEntryFilename),
       referencedAssets: new Set(outputAssets.map((asset) => asset.fileName)),
+      moduleIds,
       watchedFiles,
     }
     this.bundles.set(file, bundle)
@@ -126,6 +156,59 @@ class WorkerOutputCache {
         this.assets.delete(asset)
       }
     }
+
+    this.bundleReferences.delete(file)
+  }
+
+  recordReference(
+    parentInputId: BundleId,
+    childBundleId: WorkerBundleId,
+    referencingModuleId: string,
+  ) {
+    let referencesByModule = this.bundleReferences.get(parentInputId)
+    if (!referencesByModule) {
+      referencesByModule = new Map()
+      this.bundleReferences.set(parentInputId, referencesByModule)
+    }
+    let childBundleIds = referencesByModule.get(referencingModuleId)
+    if (!childBundleIds) {
+      childBundleIds = new Set()
+      referencesByModule.set(referencingModuleId, childBundleIds)
+    }
+    childBundleIds.add(childBundleId)
+  }
+
+  getLiveAssetFileNames(mainLiveModuleIds: Set<string>): Set<string> {
+    const liveBundles = new Set<string>()
+    const queue: [BundleId, Set<string>][] = [[undefined, mainLiveModuleIds]]
+    while (queue.length > 0) {
+      const [bundleId, moduleIds] = queue.shift()!
+      const referencesByModule = this.bundleReferences.get(bundleId)
+      if (!referencesByModule) continue
+      for (const moduleId of moduleIds) {
+        const childBundleIds = referencesByModule.get(moduleId)
+        if (!childBundleIds) continue
+        for (const childBundleId of childBundleIds) {
+          if (liveBundles.has(childBundleId)) continue
+          liveBundles.add(childBundleId)
+          const childBundle = this.bundles.get(childBundleId)
+          if (childBundle) {
+            queue.push([childBundleId, childBundle.moduleIds])
+          }
+        }
+      }
+    }
+
+    const liveFileNames = new Set<string>()
+    for (const inputId of liveBundles) {
+      const wb = this.bundles.get(inputId)
+      if (!wb) continue
+      liveFileNames.add(wb.entryFilename)
+      for (const fileName of wb.referencedAssets) {
+        liveFileNames.add(fileName)
+      }
+    }
+    return liveFileNames
   }
 
   getWorkerBundle(file: string) {
@@ -159,6 +242,17 @@ const inlineRE = /[?&]inline\b/
 export const WORKER_FILE_ID = 'worker_file'
 const workerOutputCaches = new WeakMap<ResolvedConfig, WorkerOutputCache>()
 
+export function recordWorkerReference(
+  config: ResolvedConfig,
+  parentInputId: string | undefined,
+  childBundleId: WorkerBundleId,
+  referencingModuleId: string,
+): void {
+  workerOutputCaches
+    .get(config.mainConfig || config)!
+    .recordReference(parentInputId, childBundleId, referencingModuleId)
+}
+
 async function bundleWorkerEntry(
   config: ResolvedConfig,
   id: string,
@@ -267,6 +361,8 @@ async function bundleWorkerEntry(
     await bundle.close()
   }
 
+  const moduleIds = collectIncludedModuleIds(result.output)
+
   const {
     output: [outputChunk, ...outputChunks],
   } = result
@@ -300,6 +396,7 @@ async function bundleWorkerEntry(
       outputChunk.fileName,
       outputChunk.code,
       assets,
+      moduleIds,
       config.logger,
     )
   return newBundleInfo
@@ -452,6 +549,12 @@ export function webWorkerPlugin(config: ResolvedConfig): Plugin {
           if (isWorker && config.bundleChain.at(-1) === cleanUrl(id)) {
             urlCode = 'self.location.href'
           } else if (inlineRE.test(id)) {
+            recordWorkerReference(
+              config,
+              config.bundleChain.at(-1),
+              cleanUrl(id),
+              id,
+            )
             const result = await bundleWorkerEntry(config, id)
             for (const file of result.watchedFiles) {
               this.addWatchFile(file)
@@ -502,6 +605,12 @@ export function webWorkerPlugin(config: ResolvedConfig): Plugin {
               map: { mappings: '' },
             }
           } else {
+            recordWorkerReference(
+              config,
+              config.bundleChain.at(-1),
+              cleanUrl(id),
+              id,
+            )
             const result = await workerFileToUrl(config, id)
             let url: string
             if (
@@ -661,7 +770,16 @@ export function webWorkerPlugin(config: ResolvedConfig): Plugin {
       ) {
         return
       }
-      for (const asset of workerOutputCaches.get(config)!.getAssets()) {
+      const cache = workerOutputCaches.get(config)!
+      const liveModuleIds = collectIncludedModuleIds(Object.values(bundle))
+      // Reference tracking relies on hooks running for every module, which is
+      // not guaranteed when an incremental build reuses cached modules.
+      const liveFileNames =
+        isBuild && !config.build.watch
+          ? cache.getLiveAssetFileNames(liveModuleIds)
+          : undefined
+      for (const asset of cache.getAssets()) {
+        if (liveFileNames && !liveFileNames.has(asset.fileName)) continue
         if (emittedAssets.has(asset.fileName)) continue
         emittedAssets.add(asset.fileName)
 
@@ -696,6 +814,20 @@ export function webWorkerPlugin(config: ResolvedConfig): Plugin {
   }
 }
 
+function collectIncludedModuleIds(
+  outputs: (OutputChunk | OutputAsset)[],
+): Set<string> {
+  const moduleIds = new Set<string>()
+  for (const output of outputs) {
+    if (output.type === 'chunk') {
+      for (const moduleId of output.moduleIds) {
+        moduleIds.add(moduleId)
+      }
+    }
+  }
+  return moduleIds
+}
+
 function isSameContent(a: string | Uint8Array, b: string | Uint8Array) {
   if (typeof a === 'string') {
     if (typeof b === 'string') {

packages/vite/src/node/plugins/workerImportMetaUrl.ts

@@ -14,6 +14,7 @@ import type { WorkerType } from './worker'
 import {
   WORKER_FILE_ID,
   emitWorkerAssetsForBundledDev,
+  recordWorkerReference,
   workerFileToUrl,
 } from './worker'
 import { fileToUrl, toOutputFilePathInJSForBundledDev } from './asset'
@@ -259,6 +260,12 @@ export function workerImportMetaUrlPlugin(config: ResolvedConfig): Plugin {
           } else {
             let builtUrl: string
             if (isBundled) {
+              recordWorkerReference(
+                config,
+                config.bundleChain.at(-1),
+                cleanUrl(file),
+                id,
+              )
               const result = await workerFileToUrl(config, file)
               if (this.environment.config.command === 'serve') {
                 emitWorkerAssetsForBundledDev(this, config)

playground/worker/__tests__/es/worker-es.spec.ts

@@ -141,6 +141,29 @@ describe.runIf(isBuild)('build', () => {
       .poll(() => page.textContent('.nested-worker-constructor'))
       .toMatch('"type":"constructor"')
   })
+
+  test('dead-code-eliminated worker asset is not emitted', () => {
+    const assetsDir = path.resolve(testDir, 'dist/es/assets')
+    const files = fs.readdirSync(assetsDir)
+
+    // dce-test-importer.js is imported from main-module.js but its export is
+    // unused; rolldown tree-shakes the importer, so its `?worker` import
+    // should never reach the output bundle.
+    expect(files.some((f) => f.includes('dce-test-worker'))).toBe(false)
+    // the nested worker is only reachable via the (tree-shaken) parent worker,
+    // so it must not be emitted either.
+    expect(files.some((f) => f.includes('dce-test-nested-worker'))).toBe(false)
+
+    // This parent worker is live, but the module that references its nested
+    // worker is tree-shaken from the parent worker bundle.
+    expect(files.some((f) => f.includes('dce-test-live-worker'))).toBe(true)
+    expect(files.some((f) => f.includes('dce-test-live-nested-worker'))).toBe(
+      false,
+    )
+
+    // sanity: the worker we DO use is still emitted (`my-worker`).
+    expect(files.some((f) => f.includes('my-worker'))).toBe(true)
+  })
 })
 
 test('module worker', async () => {

playground/worker/dce-test-importer.js

@@ -0,0 +1,3 @@
+import DceTestWorker from './dce-test-worker.js?worker'
+
+export const dceTestWorker = DceTestWorker

playground/worker/dce-test-live-nested-worker.js

@@ -0,0 +1 @@
+self.postMessage('dce-test-live-nested-worker should be tree-shaken')

playground/worker/dce-test-live-worker-importer.js

@@ -0,0 +1,3 @@
+import DceTestLiveNestedWorker from './dce-test-live-nested-worker.js?worker'
+
+export { DceTestLiveNestedWorker }

playground/worker/dce-test-live-worker.js

@@ -0,0 +1,3 @@
+import { DceTestLiveNestedWorker as _DceTestLiveNestedWorker } from './dce-test-live-worker-importer.js'
+
+self.postMessage('dce-test-live-worker is live')

playground/worker/dce-test-nested-worker.js

@@ -0,0 +1 @@
+self.postMessage('dce-test-nested-worker should be tree-shaken')

playground/worker/dce-test-worker.js

@@ -0,0 +1,4 @@
+import DceTestNestedWorker from './dce-test-nested-worker.js?worker'
+
+const _nested = new DceTestNestedWorker()
+self.postMessage('dce-test-worker should be tree-shaken')

playground/worker/worker/main-module.js

@@ -1,4 +1,9 @@
 import * as depSelfReferenceUrlWorker from '@vitejs/test-dep-self-reference-url-worker'
+import DceTestLiveWorker from '../dce-test-live-worker.js?worker'
+// imported but never used — the dead-importer module itself is side-effect-free,
+// so rolldown DCEs it along with its `?worker` import. The worker plugin should
+// then skip emitting the dce-test-worker asset (and its nested worker too).
+import { dceTestWorker as _dceTestWorker } from '../dce-test-importer.js'
 import myWorker from '../my-worker.ts?worker'
 import InlineWorker from '../my-worker.ts?worker&inline'
 import InlineSharedWorker from '../my-inline-shared-worker?sharedworker&inline'
@@ -14,6 +19,8 @@ function text(el, text) {
 
 document.querySelector('.mode-true').textContent = mode
 
+new DceTestLiveWorker()
+
 const worker = new myWorker()
 worker.postMessage('ping')
 worker.addEventListener('message', (e) => {