feat: provide environment in every hook context

package.json

@@ -43,6 +43,7 @@
     "@babel/types": "^7.24.5",
     "@eslint-types/typescript-eslint": "^7.5.0",
     "@rollup/plugin-typescript": "^11.1.6",
+    "@type-challenges/utils": "^0.1.1",
     "@types/babel__core": "^7.20.5",
     "@types/babel__preset-env": "^7.9.6",
     "@types/convert-source-map": "^2.0.3",

packages/vite/package.json

@@ -78,7 +78,7 @@
     "build-types-temp": "tsc --emitDeclarationOnly --outDir temp -p src/node",
     "build-types-roll": "rollup --config rollup.dts.config.ts --configPlugin typescript && rimraf temp",
     "build-types-check": "tsc --project tsconfig.check.json",
-    "typecheck": "tsc --noEmit",
+    "typecheck": "tsc --noEmit && tsc --noEmit -p src/node",
     "lint": "eslint --cache --ext .ts src/**",
     "format": "prettier --write --cache --parser typescript \"src/**/*.ts\"",
     "prepublishOnly": "npm run build"

packages/vite/src/node/__tests_dts__/plugin.ts

@@ -0,0 +1,38 @@
+/**
+ * This is a developement only file for testing types.
+ */
+import type { Plugin as RollupPlugin } from 'rollup'
+import type { Equal, ExpectExtends, ExpectTrue } from '@type-challenges/utils'
+import type { EnvironmentPlugin, PluginContextExtension } from '../plugin'
+import type { ROLLUP_HOOKS } from '../constants'
+import type {
+  GetHookContextMap,
+  NonNeverKeys,
+  RollupPluginHooks,
+} from '../typeUtils'
+
+type EnvironmentPluginHooksContext = GetHookContextMap<EnvironmentPlugin>
+type EnvironmentPluginHooksContextMatched = {
+  [K in keyof EnvironmentPluginHooksContext]: EnvironmentPluginHooksContext[K] extends PluginContextExtension
+    ? never
+    : false
+}
+
+type HooksMissingExtension = NonNeverKeys<EnvironmentPluginHooksContextMatched>
+type HooksMissingInConstans = Exclude<
+  RollupPluginHooks,
+  (typeof ROLLUP_HOOKS)[number]
+>
+
+export type cases = [
+  // Ensure environment plugin hooks are superset of rollup plugin hooks
+  ExpectTrue<ExpectExtends<RollupPlugin, EnvironmentPlugin>>,
+
+  // Ensure all Rollup hooks have Vite's plugin context extension
+  ExpectTrue<Equal<HooksMissingExtension, never>>,
+
+  // Ensure the `ROLLUP_HOOKS` constant is up-to-date
+  ExpectTrue<Equal<HooksMissingInConstans, never>>,
+]
+
+export {}

packages/vite/src/node/build.ts

@@ -24,6 +24,7 @@ import { withTrailingSlash } from '../shared/utils'
 import {
   DEFAULT_ASSETS_INLINE_LIMIT,
   ESBUILD_MODULES_TARGET,
+  ROLLUP_HOOKS,
   VERSION,
 } from './constants'
 import type {
@@ -67,6 +68,7 @@ import { webWorkerPostPlugin } from './plugins/worker'
 import { getHookHandler } from './plugins'
 import { Environment } from './environment'
 import type { Plugin, PluginContext } from './plugin'
+import type { RollupPluginHooks } from './typeUtils'
 
 export interface BuildEnvironmentOptions {
   /**
@@ -1124,23 +1126,30 @@ export function injectEnvironmentToHooks(
   plugin: Plugin,
   environment?: BuildEnvironment,
 ): Plugin {
-  const {
-    buildStart,
-    resolveId,
-    load,
-    transform,
-    generateBundle,
-    renderChunk,
-  } = plugin
-  return {
-    ...plugin,
-    resolveId: wrapEnvironmentResolveId(resolveId, environment),
-    load: wrapEnvironmentLoad(load, environment),
-    transform: wrapEnvironmentTransform(transform, environment),
-    buildStart: wrapEnvironmentHook(buildStart, environment),
-    generateBundle: wrapEnvironmentHook(generateBundle, environment),
-    renderChunk: wrapEnvironmentHook(renderChunk, environment),
+  const { resolveId, load, transform } = plugin
+
+  const clone = { ...plugin }
+
+  for (const hook of Object.keys(clone) as RollupPluginHooks[]) {
+    switch (hook) {
+      case 'resolveId':
+        clone[hook] = wrapEnvironmentResolveId(resolveId, environment)
+        break
+      case 'load':
+        clone[hook] = wrapEnvironmentLoad(load, environment)
+        break
+      case 'transform':
+        clone[hook] = wrapEnvironmentTransform(transform, environment)
+        break
+      default:
+        if (ROLLUP_HOOKS.includes(hook)) {
+          ;(clone as any)[hook] = wrapEnvironmentHook(clone[hook], environment)
+        }
+        break
+    }
   }
+
+  return clone
 }
 
 function wrapEnvironmentResolveId(
@@ -1227,6 +1236,8 @@ function wrapEnvironmentHook<HookName extends keyof Plugin>(
   if (!hook) return
 
   const fn = getHookHandler(hook)
+  if (typeof fn !== 'function') return hook
+
   const handler: Plugin[HookName] = function (
     this: PluginContext,
     ...args: any[]
@@ -1248,14 +1259,8 @@ function injectEnvironmentInContext<Context extends PluginContext>(
   context: Context,
   environment?: BuildEnvironment,
 ) {
-  return new Proxy(context, {
-    get(target, prop, receiver) {
-      if (prop === 'environment') {
-        return environment
-      }
-      return Reflect.get(target, prop, receiver)
-    },
-  })
+  context.environment ??= environment
+  return context
 }
 
 function injectSsrFlag<T extends Record<string, any>>(

packages/vite/src/node/constants.ts

@@ -1,11 +1,40 @@
 import path, { resolve } from 'node:path'
 import { fileURLToPath } from 'node:url'
 import { readFileSync } from 'node:fs'
+import type { RollupPluginHooks } from './typeUtils'
 
 const { version } = JSON.parse(
   readFileSync(new URL('../../package.json', import.meta.url)).toString(),
 )
 
+export const ROLLUP_HOOKS = [
+  'buildStart',
+  'buildEnd',
+  'renderStart',
+  'renderError',
+  'renderChunk',
+  'writeBundle',
+  'generateBundle',
+  'banner',
+  'footer',
+  'augmentChunkHash',
+  'outputOptions',
+  'renderDynamicImport',
+  'resolveFileUrl',
+  'resolveImportMeta',
+  'intro',
+  'outro',
+  'closeBundle',
+  'closeWatcher',
+  'load',
+  'moduleParsed',
+  'watchChange',
+  'resolveDynamicImport',
+  'resolveId',
+  'shouldTransformCachedModule',
+  'transform',
+] satisfies RollupPluginHooks[]
+
 export const VERSION = version as string
 
 export const DEFAULT_MAIN_FIELDS = [

packages/vite/src/node/plugin.ts

@@ -59,17 +59,21 @@ export type PluginEnvironment =
   | BuildEnvironment
   | ScanEnvironment
 
-export interface PluginContext extends RollupPluginContext {
+export interface PluginContextExtension {
   environment?: PluginEnvironment
 }
 
-export interface ResolveIdPluginContext extends RollupPluginContext {
-  environment?: PluginEnvironment
-}
+export interface PluginContext
+  extends RollupPluginContext,
+    PluginContextExtension {}
 
-export interface TransformPluginContext extends RollupTransformPluginContext {
-  environment?: PluginEnvironment
-}
+export interface ResolveIdPluginContext
+  extends RollupPluginContext,
+    PluginContextExtension {}
+
+export interface TransformPluginContext
+  extends RollupTransformPluginContext,
+    PluginContextExtension {}
 
 /**
  * There are two types of plugins in Vite. App plugins and environment plugins.
@@ -87,7 +91,7 @@ type ModifyFunctionContext<Function_, NewContext> = Function_ extends (
   ...parameters: infer Arguments
 ) => infer Return
   ? (this: NewContext, ...parameters: Arguments) => Return
-  : never
+  : Function_
 
 type ModifyObjectHookContext<
   Handler,
@@ -182,13 +186,60 @@ export interface EnvironmentPlugin<A = any> extends RollupPlugin<A> {
     ) => Promise<TransformResult> | TransformResult
   >
 
-  // TODO: abstract to every hook in RollupPlugin?
+  // Extends rollup hooks
+  // ./__tests_dts__/plugin.ts will guard this to ensure we have all hooks
+  augmentChunkHash?: ModifyHookContext<
+    RollupPlugin<A>['augmentChunkHash'],
+    PluginContext
+  >
+  banner?: ModifyHookContext<RollupPlugin<A>['banner'], PluginContext>
+  buildEnd?: ModifyHookContext<RollupPlugin<A>['buildEnd'], PluginContext>
   buildStart?: ModifyHookContext<RollupPlugin<A>['buildStart'], PluginContext>
+  closeBundle?: ModifyHookContext<RollupPlugin<A>['closeBundle'], PluginContext>
+  closeWatcher?: ModifyHookContext<
+    RollupPlugin<A>['closeWatcher'],
+    PluginContext
+  >
+  footer?: ModifyHookContext<RollupPlugin<A>['footer'], PluginContext>
   generateBundle?: ModifyHookContext<
     RollupPlugin<A>['generateBundle'],
     PluginContext
   >
+  intro?: ModifyHookContext<RollupPlugin<A>['intro'], PluginContext>
+  moduleParsed?: ModifyHookContext<
+    RollupPlugin<A>['moduleParsed'],
+    PluginContext
+  >
+  outputOptions?: ModifyHookContext<
+    RollupPlugin<A>['outputOptions'],
+    PluginContext
+  >
+  outro?: ModifyHookContext<RollupPlugin<A>['outro'], PluginContext>
   renderChunk?: ModifyHookContext<RollupPlugin<A>['renderChunk'], PluginContext>
+  renderDynamicImport?: ModifyHookContext<
+    RollupPlugin<A>['renderDynamicImport'],
+    PluginContext
+  >
+  renderError?: ModifyHookContext<RollupPlugin<A>['renderError'], PluginContext>
+  renderStart?: ModifyHookContext<RollupPlugin<A>['renderStart'], PluginContext>
+  resolveDynamicImport?: ModifyHookContext<
+    RollupPlugin<A>['resolveDynamicImport'],
+    PluginContext
+  >
+  resolveFileUrl?: ModifyHookContext<
+    RollupPlugin<A>['resolveFileUrl'],
+    PluginContext
+  >
+  resolveImportMeta?: ModifyHookContext<
+    RollupPlugin<A>['resolveImportMeta'],
+    PluginContext
+  >
+  watchChange?: ModifyHookContext<RollupPlugin<A>['watchChange'], PluginContext>
+  writeBundle?: ModifyHookContext<RollupPlugin<A>['writeBundle'], PluginContext>
+  shouldTransformCachedModule?: ModifyHookContext<
+    RollupPlugin<A>['shouldTransformCachedModule'],
+    PluginContext
+  >
 }
 
 export interface Plugin<A = any> extends EnvironmentPlugin<A> {

packages/vite/src/node/tsconfig.json

@@ -4,10 +4,11 @@
     "./",
     "../module-runner",
     "../dep-types",
+    "./__tests_dts__",
     "../types",
     "constants.ts"
   ],
-  "exclude": ["../**/__tests__"],
+  "exclude": ["../**/__tests__/"],
   "compilerOptions": {
     "lib": ["ESNext", "DOM"],
     "stripInternal": true,

packages/vite/src/node/typeUtils.ts

@@ -0,0 +1,22 @@
+import type {
+  ObjectHook,
+  Plugin as RollupPlugin,
+  PluginContext as RollupPluginContext,
+} from 'rollup'
+
+export type NonNeverKeys<T> = {
+  [K in keyof T]: T[K] extends never ? never : K
+}[keyof T]
+
+export type GetHookContextMap<Plugin> = {
+  [K in keyof Plugin]-?: Plugin[K] extends ObjectHook<infer T, infer B>
+    ? T extends (this: infer This, ...args: any[]) => any
+      ? This extends RollupPluginContext
+        ? This
+        : never
+      : never
+    : never
+}
+
+type RollupPluginHooksContext = GetHookContextMap<RollupPlugin>
+export type RollupPluginHooks = NonNeverKeys<RollupPluginHooksContext>

pnpm-lock.yaml

@@ -33,6 +33,9 @@ importers:
       '@rollup/plugin-typescript':
         specifier: ^11.1.6
         version: 11.1.6(rollup@4.13.0)(tslib@2.6.2)(typescript@5.2.2)
+      '@type-challenges/utils':
+        specifier: ^0.1.1
+        version: 0.1.1
       '@types/babel__core':
         specifier: ^7.20.5
         version: 7.20.5
@@ -4447,6 +4450,10 @@ packages:
   /@tsconfig/node16@1.0.2:
     resolution: {integrity: sha512-eZxlbI8GZscaGS7kkc/trHTT5xgrjH3/1n2JDwusC9iahPKWMRvRjJSAN5mCXviuTGQ/lHnhvv8Q1YTpnfz9gA==}
 
+  /@type-challenges/utils@0.1.1:
+    resolution: {integrity: sha512-A7ljYfBM+FLw+NDyuYvGBJiCEV9c0lPWEAdzfOAkb3JFqfLl0Iv/WhWMMARHiRKlmmiD1g8gz/507yVvHdQUYA==}
+    dev: true
+
   /@types/babel__core@7.20.5:
     resolution: {integrity: sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==}
     dependencies: