feat: extended applyToEnvironment and perEnvironmentPlugin (#18544)

Co-authored-by: 翠 / green <green@sapphi.red>
Co-authored-by: Bjorn Lu <bjornlu.dev@gmail.com>

docs/changes/shared-plugins-during-build.md

@@ -59,11 +59,11 @@ function PerEnvironmentCountTransformedModulesPlugin() {
 }
 ```
 
-To simplify this pattern, internally in Vite, we use a `usePerEnvironmentState` helper:
+To simplify this pattern, Vite exports a `perEnvironmentState` helper:
 
 ```js
 function PerEnvironmentCountTransformedModulesPlugin() {
-  const state = usePerEnvironmentState<{ count: number }>(() => ({ count: 0 }))
+  const state = perEnvironmentState<{ count: number }>(() => ({ count: 0 }))
   return {
     name: 'count-transformed-modules',
     perEnvironmentStartEndDuringDev: true,

docs/guide/api-environment-plugins.md

@@ -142,7 +142,8 @@ const UnoCssPlugin = () => {
       // use global hooks normally
     },
     applyToEnvironment(environment) {
-      // return true if this plugin should be active in this environment
+      // return true if this plugin should be active in this environment,
+      // or return a new plugin to replace it.
       // if the hook is not used, the plugin is active in all environments
     },
     resolveId(id, importer) {
@@ -152,6 +153,37 @@ const UnoCssPlugin = () => {
 }
 ```
 
+If a plugin isn't environment aware and has state that isn't keyed on the current environment, the `applyToEnvironment` hook allows to easily make it per-environment.
+
+```js
+import { nonShareablePlugin } from 'non-shareable-plugin'
+
+export default defineConfig({
+  plugins: [
+    {
+      name: 'per-environment-plugin',
+      applyToEnvironment(environment) {
+        return nonShareablePlugin({ outputName: environment.name })
+      },
+    },
+  ],
+})
+```
+
+Vite exports a `perEnvironmentPlugin` helper to simplify these cases where no other hooks are required:
+
+```js
+import { nonShareablePlugin } from 'non-shareable-plugin'
+
+export default defineConfig({
+  plugins: [
+    perEnvironmentPlugin('per-environment-plugin', (environment) =>
+      nonShareablePlugin({ outputName: environment.name }),
+    ),
+  ],
+})
+```
+
 ## Environment in build hooks
 
 In the same way as during dev, plugin hooks also receive the environment instance during build, replacing the `ssr` boolean.

packages/vite/src/node/build.ts

@@ -50,7 +50,7 @@ import {
   normalizePath,
   partialEncodeURIPath,
 } from './utils'
-import { resolveEnvironmentPlugins } from './plugin'
+import { perEnvironmentPlugin, resolveEnvironmentPlugins } from './plugin'
 import { manifestPlugin } from './plugins/manifest'
 import type { Logger } from './logger'
 import { dataURIPlugin } from './plugins/dataUri'
@@ -473,34 +473,26 @@ export async function resolveBuildPlugins(config: ResolvedConfig): Promise<{
   pre: Plugin[]
   post: Plugin[]
 }> {
-  const { commonjsOptions } = config.build
-  const usePluginCommonjs =
-    !Array.isArray(commonjsOptions.include) ||
-    commonjsOptions.include.length !== 0
   return {
     pre: [
       completeSystemWrapPlugin(),
-      /**
-       * environment.config.build.commonjsOptions isn't currently supported
-       * when builder.sharedConfigBuild or builder.sharedPlugins enabled.
-       * To do it, we could inject one commonjs plugin per environment with
-       * an applyToEnvironment hook.
-       */
-      ...(usePluginCommonjs ? [commonjsPlugin(commonjsOptions)] : []),
+      perEnvironmentPlugin('commonjs', (environment) => {
+        const { commonjsOptions } = environment.config.build
+        const usePluginCommonjs =
+          !Array.isArray(commonjsOptions.include) ||
+          commonjsOptions.include.length !== 0
+        return usePluginCommonjs ? commonjsPlugin(commonjsOptions) : false
+      }),
       dataURIPlugin(),
-      /**
-       * environment.config.build.rollupOptions.plugins isn't supported
-       * when builder.sharedConfigBuild or builder.sharedPlugins is enabled.
-       * To do it, we should add all these plugins to the global pipeline, each with
-       * an applyToEnvironment hook. It is similar to letting the user add per
-       * environment plugins giving them a environment.config.plugins option that
-       * we decided against.
-       * For backward compatibility, we are still injecting the rollup plugins
-       * defined in the default root build options.
-       */
-      ...((
-        await asyncFlatten(arraify(config.build.rollupOptions.plugins))
-      ).filter(Boolean) as Plugin[]),
+      perEnvironmentPlugin(
+        'vite:rollup-options-plugins',
+        async (environment) =>
+          (
+            await asyncFlatten(
+              arraify(environment.config.build.rollupOptions.plugins),
+            )
+          ).filter(Boolean) as Plugin[],
+      ),
       ...(config.isWorker ? [webWorkerPostPlugin()] : []),
     ],
     post: [
@@ -1464,13 +1456,12 @@ export class BuildEnvironment extends BaseEnvironment {
     super(name, config, options)
   }
 
-  // TODO: This could be sync, discuss if applyToEnvironment should support async
   async init(): Promise<void> {
     if (this._initiated) {
       return
     }
     this._initiated = true
-    this._plugins = resolveEnvironmentPlugins(this)
+    this._plugins = await resolveEnvironmentPlugins(this)
   }
 }
 

packages/vite/src/node/environment.ts

@@ -14,8 +14,10 @@ export type Environment =
  * Creates a function that hides the complexities of a WeakMap with an initial value
  * to implement object metadata. Used by plugins to implement cross hooks per
  * environment metadata
+ *
+ * @experimental
  */
-export function usePerEnvironmentState<State>(
+export function perEnvironmentState<State>(
   initial: (environment: Environment) => State,
 ): (context: PluginContext) => State {
   const stateMap = new WeakMap<Environment, State>()

packages/vite/src/node/index.ts

@@ -8,6 +8,8 @@ export {
   resolveConfig,
   sortUserPlugins,
 } from './config'
+export { perEnvironmentPlugin } from './plugin'
+export { perEnvironmentState } from './environment'
 export { createServer } from './server'
 export { preview } from './preview'
 export { build, createBuilder } from './build'

packages/vite/src/node/optimizer/scan.ts

@@ -63,7 +63,7 @@ export class ScanEnvironment extends BaseEnvironment {
       return
     }
     this._initiated = true
-    this._plugins = resolveEnvironmentPlugins(this)
+    this._plugins = await resolveEnvironmentPlugins(this)
     this._pluginContainer = await createEnvironmentPluginContainer(
       this,
       this.plugins,

packages/vite/src/node/plugin.ts

@@ -22,7 +22,9 @@ import type { ModuleNode } from './server/mixedModuleGraph'
 import type { HmrContext, HotUpdateOptions } from './server/hmr'
 import type { DevEnvironment } from './server/environment'
 import type { Environment } from './environment'
+import type { PartialEnvironment } from './baseEnvironment'
 import type { PreviewServerHook } from './preview'
+import { arraify, asyncFlatten } from './utils'
 
 /**
  * Vite plugins extends the Rollup plugin interface with a few extra
@@ -203,8 +205,11 @@ export interface Plugin<A = any> extends RollupPlugin<A> {
   /**
    * Define environments where this plugin should be active
    * By default, the plugin is active in all environments
+   * @experimental
    */
-  applyToEnvironment?: (environment: Environment) => boolean
+  applyToEnvironment?: (
+    environment: PartialEnvironment,
+  ) => boolean | Promise<boolean> | PluginOption
   /**
    * Modify vite config before it's resolved. The hook can either mutate the
    * passed-in config directly, or return a partial config object that will be
@@ -324,11 +329,41 @@ type FalsyPlugin = false | null | undefined
 
 export type PluginOption = Thenable<Plugin | FalsyPlugin | PluginOption[]>
 
-export function resolveEnvironmentPlugins(environment: Environment): Plugin[] {
-  return environment
-    .getTopLevelConfig()
-    .plugins.filter(
-      (plugin) =>
-        !plugin.applyToEnvironment || plugin.applyToEnvironment(environment),
-    )
+export async function resolveEnvironmentPlugins(
+  environment: PartialEnvironment,
+): Promise<Plugin[]> {
+  const environmentPlugins: Plugin[] = []
+  for (const plugin of environment.getTopLevelConfig().plugins) {
+    if (plugin.applyToEnvironment) {
+      const applied = await plugin.applyToEnvironment(environment)
+      if (!applied) {
+        continue
+      }
+      if (applied !== true) {
+        environmentPlugins.push(
+          ...((await asyncFlatten(arraify(applied))).filter(
+            Boolean,
+          ) as Plugin[]),
+        )
+        continue
+      }
+    }
+    environmentPlugins.push(plugin)
+  }
+  return environmentPlugins
+}
+
+/**
+ * @experimental
+ */
+export function perEnvironmentPlugin(
+  name: string,
+  applyToEnvironment: (
+    environment: PartialEnvironment,
+  ) => boolean | Promise<boolean> | PluginOption,
+): Plugin {
+  return {
+    name,
+    applyToEnvironment,
+  }
 }

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

@@ -3,7 +3,7 @@ import type { Plugin } from '../plugin'
 import type { ResolvedConfig } from '../config'
 import { CLIENT_ENTRY, ENV_ENTRY } from '../constants'
 import { isObject, normalizePath, resolveHostname } from '../utils'
-import { usePerEnvironmentState } from '../environment'
+import { perEnvironmentState } from '../environment'
 import { replaceDefine, serializeDefine } from './define'
 
 // ids in transform are normalized to unix style
@@ -17,7 +17,7 @@ const normalizedEnvEntry = normalizePath(ENV_ENTRY)
 export function clientInjectionsPlugin(config: ResolvedConfig): Plugin {
   let injectConfigValues: (code: string) => string
 
-  const getDefineReplacer = usePerEnvironmentState((environment) => {
+  const getDefineReplacer = perEnvironmentState((environment) => {
     const userDefine: Record<string, any> = {}
     for (const key in environment.config.define) {
       // import.meta.env.* is handled in `importAnalysis` plugin

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

@@ -18,7 +18,7 @@ import {
   urlRE,
 } from '../utils'
 import type { Environment } from '../environment'
-import { usePerEnvironmentState } from '../environment'
+import { perEnvironmentState } from '../environment'
 import { hasViteIgnoreRE } from './importAnalysis'
 import { workerOrSharedWorkerRE } from './worker'
 
@@ -171,7 +171,7 @@ export function dynamicImportVarsPlugin(config: ResolvedConfig): Plugin {
     extensions: [],
   })
 
-  const getFilter = usePerEnvironmentState((environment: Environment) => {
+  const getFilter = perEnvironmentState((environment: Environment) => {
     const { include, exclude } =
       environment.config.build.dynamicImportVarsOptions
     return createFilter(include, exclude)

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

@@ -32,7 +32,7 @@ import { toOutputFilePathInHtml } from '../build'
 import { resolveEnvPrefix } from '../env'
 import type { Logger } from '../logger'
 import { cleanUrl } from '../../shared/utils'
-import { usePerEnvironmentState } from '../environment'
+import { perEnvironmentState } from '../environment'
 import { getNodeAssetAttributes } from '../assetSource'
 import {
   assetUrlRE,
@@ -330,7 +330,7 @@ export function buildHtmlPlugin(config: ResolvedConfig): Plugin {
   preHooks.push(htmlEnvHook(config))
   postHooks.push(injectNonceAttributeTagHook(config))
   postHooks.push(postImportMapHook())
-  const processedHtml = usePerEnvironmentState(() => new Map<string, string>())
+  const processedHtml = perEnvironmentState(() => new Map<string, string>())
 
   const isExcludedUrl = (url: string) =>
     url[0] === '#' || isExternalUrl(url) || isDataUrl(url)

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

@@ -7,7 +7,7 @@ import type {
 } from 'rollup'
 import type { Plugin } from '../plugin'
 import { normalizePath, sortObjectKeys } from '../utils'
-import { usePerEnvironmentState } from '../environment'
+import { perEnvironmentState } from '../environment'
 import { cssEntriesMap } from './asset'
 
 const endsWithJSRE = /\.[cm]?js$/
@@ -27,7 +27,7 @@ export interface ManifestChunk {
 }
 
 export function manifestPlugin(): Plugin {
-  const getState = usePerEnvironmentState(() => {
+  const getState = perEnvironmentState(() => {
     return {
       manifest: {} as Manifest,
       outputCount: 0,

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

@@ -6,7 +6,7 @@ import type { OutputBundle } from 'rollup'
 import type { Plugin } from '../plugin'
 import type { ResolvedConfig } from '../config'
 import type { Environment } from '../environment'
-import { usePerEnvironmentState } from '../environment'
+import { perEnvironmentState } from '../environment'
 import { isDefined, isInNodeModules, normalizePath } from '../utils'
 import { LogLevels } from '../logger'
 import { withTrailingSlash } from '../../shared/utils'
@@ -40,7 +40,7 @@ export function buildReporterPlugin(config: ResolvedConfig): Plugin {
   const tty = process.stdout.isTTY && !process.env.CI
   const shouldLogInfo = LogLevels[config.logLevel || 'info'] >= LogLevels.info
 
-  const modulesReporter = usePerEnvironmentState((environment: Environment) => {
+  const modulesReporter = perEnvironmentState((environment: Environment) => {
     let hasTransformed = false
     let transformedCount = 0
 
@@ -83,7 +83,7 @@ export function buildReporterPlugin(config: ResolvedConfig): Plugin {
     }
   })
 
-  const chunksReporter = usePerEnvironmentState((environment: Environment) => {
+  const chunksReporter = perEnvironmentState((environment: Environment) => {
     let hasRenderedChunk = false
     let hasCompressChunk = false
     let chunkCount = 0

packages/vite/src/node/server/environment.ts

@@ -167,7 +167,7 @@ export class DevEnvironment extends BaseEnvironment {
       return
     }
     this._initiated = true
-    this._plugins = resolveEnvironmentPlugins(this)
+    this._plugins = await resolveEnvironmentPlugins(this)
     this._pluginContainer = await createEnvironmentPluginContainer(
       this,
       this._plugins,

packages/vite/src/node/ssr/ssrManifestPlugin.ts

@@ -14,11 +14,11 @@ import {
   numberToPos,
   sortObjectKeys,
 } from '../utils'
-import { usePerEnvironmentState } from '../environment'
+import { perEnvironmentState } from '../environment'
 
 export function ssrManifestPlugin(): Plugin {
   // module id => preload assets mapping
-  const getSsrManifest = usePerEnvironmentState(() => {
+  const getSsrManifest = perEnvironmentState(() => {
     return {} as Record<string, string[]>
   })