Reemit chunks emitted from transform hooks

src/Module.ts

@@ -32,6 +32,7 @@ import ExternalModule from './ExternalModule';
 import Graph from './Graph';
 import {
 	Asset,
+	EmittedChunk,
 	ModuleJSON,
 	RawSourceMap,
 	ResolvedIdMap,
@@ -205,6 +206,7 @@ export default class Module {
 	sourcemapChain!: RawSourceMap[];
 	sources: string[] = [];
 	transformAssets?: Asset[];
+	transformChunks?: EmittedChunk[];
 	usesTopLevelAwait = false;
 
 	private allExportNames?: Set<string>;
@@ -618,6 +620,7 @@ export default class Module {
 			resolvedIds: this.resolvedIds,
 			sourcemapChain: this.sourcemapChain,
 			transformAssets: this.transformAssets,
+			transformChunks: this.transformChunks,
 			transformDependencies: this.transformDependencies
 		};
 	}

src/ModuleLoader.ts

@@ -332,10 +332,13 @@ export class ModuleLoader {
 					!cachedModule.customTransformCache &&
 					cachedModule.originalCode === sourceDescription.code
 				) {
-					// re-emit transform assets
 					if (cachedModule.transformAssets) {
-						for (const asset of cachedModule.transformAssets)
-							this.pluginDriver.emitAsset(asset.name, asset.source);
+						for (const { name, source } of cachedModule.transformAssets)
+							this.pluginDriver.emitAsset(name, source);
+					}
+					if (cachedModule.transformChunks) {
+						for (const { id, options } of cachedModule.transformChunks)
+							this.pluginDriver.emitChunk(id, options);
 					}
 					return cachedModule;
 				}

src/rollup/types.d.ts

@@ -89,6 +89,12 @@ export interface ModuleJSON extends TransformModuleJSON {
 	dependencies: string[];
 	id: string;
 	transformAssets: Asset[] | void;
+	transformChunks: EmittedChunk[] | void;
+}
+
+export interface EmittedChunk {
+	id: string;
+	options: { name?: string } | undefined;
 }
 
 export interface Asset {

src/utils/pluginDriver.ts

@@ -5,6 +5,7 @@ import Graph from '../Graph';
 import Module from '../Module';
 import {
 	EmitAsset,
+	EmitChunk,
 	InputOptions,
 	Plugin,
 	PluginCache,
@@ -28,6 +29,7 @@ type EnsurePromise<T> = Promise<T extends Promise<infer K> ? K : T>;
 
 export interface PluginDriver {
 	emitAsset: EmitAsset;
+	emitChunk: EmitChunk;
 	hasLoadersOrTransforms: boolean;
 	getAssetFileName(assetReferenceId: string): string;
 	hookFirst<H extends keyof PluginHooks, R = ReturnType<PluginHooks[H]>>(
@@ -153,10 +155,7 @@ export function createPluginDriver(
 			emitChunk(id, options) {
 				if (graph.phase > BuildPhase.LOAD_AND_PARSE)
 					this.error(errInvalidRollupPhaseForEmitChunk());
-				return graph.moduleLoader.addEntryModuleAndGetReferenceId({
-					alias: (options && options.name) || null,
-					unresolvedId: id
-				});
+				return pluginDriver.emitChunk(id, options);
 			},
 			error(err): never {
 				if (typeof err === 'string') err = { message: err };
@@ -320,6 +319,12 @@ export function createPluginDriver(
 
 	const pluginDriver: PluginDriver = {
 		emitAsset,
+		emitChunk(id, options) {
+			return graph.moduleLoader.addEntryModuleAndGetReferenceId({
+				alias: (options && options.name) || null,
+				unresolvedId: id
+			});
+		},
 		getAssetFileName: getAssetFileName as (assetId: string) => string,
 		hasLoadersOrTransforms,
 

src/utils/transform.ts

@@ -4,6 +4,7 @@ import Module from '../Module';
 import {
 	Asset,
 	EmitAsset,
+	EmittedChunk,
 	ExistingRawSourceMap,
 	Plugin,
 	PluginCache,
@@ -36,7 +37,8 @@ export default function transform(
 	const originalCode = source.code;
 	let ast = source.ast;
 	let transformDependencies: string[];
-	let assets: Asset[];
+	let emittedAssets: Asset[];
+	const emittedChunks: EmittedChunk[] = [];
 	let customTransformCache = false;
 	let moduleSideEffects: boolean | null = null;
 	let trackedPluginCache: { cache: PluginCache; used: boolean };
@@ -60,7 +62,8 @@ export default function transform(
 			}
 		} else {
 			// assets emitted by transform are transformDependencies
-			if (assets.length) module.transformAssets = assets;
+			if (emittedAssets.length) module.transformAssets = emittedAssets;
+			if (emittedChunks.length) module.transformChunks = emittedChunks;
 
 			if (result && typeof result === 'object' && Array.isArray(result.dependencies)) {
 				// not great, but a useful way to track this without assuming WeakMap
@@ -124,7 +127,10 @@ export default function transform(
 				else trackedPluginCache = trackPluginCache(pluginContext.cache);
 
 				let emitAsset: EmitAsset;
-				({ assets, emitAsset } = createTransformEmitAsset(graph.assetsById, baseEmitAsset));
+				({ assets: emittedAssets, emitAsset } = createTransformEmitAsset(
+					graph.assetsById,
+					baseEmitAsset
+				));
 				return {
 					...pluginContext,
 					cache: trackedPluginCache ? trackedPluginCache.cache : pluginContext.cache,
@@ -143,6 +149,10 @@ export default function transform(
 						return pluginContext.error(err);
 					},
 					emitAsset,
+					emitChunk(id, options) {
+						emittedChunks.push({ id, options });
+						return graph.pluginDriver.emitChunk(id, options);
+					},
 					addWatchFile(id: string) {
 						if (!transformDependencies) transformDependencies = [];
 						transformDependencies.push(id);

test/hooks/index.js

@@ -360,6 +360,58 @@ describe('hooks', () => {
 			});
 	});
 
+	it('caches chunk emission in transform hook', () => {
+		let cache;
+		return rollup
+			.rollup({
+				input: 'input',
+				plugins: [
+					loader({ input: '', chunk: "console.log('chunk');" }),
+					{
+						transform(code, id) {
+							if (id === 'input') {
+								return `export default import.meta.ROLLUP_CHUNK_URL_${this.emitChunk('chunk')};`;
+							}
+						}
+					}
+				]
+			})
+			.then(bundle => {
+				cache = bundle.cache;
+				return bundle.generate({ format: 'es' });
+			})
+			.then(({ output }) => {
+				assert.equal(
+					output[0].code,
+					`var input = new URL('chunk-01406d83.js', import.meta.url).href;\n\nexport default input;\n`
+				);
+				assert.equal(output[1].fileName, 'chunk-01406d83.js');
+				assert.equal(output[1].code, `console.log('chunk');\n`);
+
+				return rollup.rollup({
+					cache,
+					input: 'input',
+					plugins: [
+						loader({ input: '', chunk: "console.log('chunk');" }),
+						{
+							transform() {
+								assert.fail('Should cache transform');
+							}
+						}
+					]
+				});
+			})
+			.then(bundle => bundle.generate({ format: 'es' }))
+			.then(({ output }) => {
+				assert.equal(
+					output[0].code,
+					`var input = new URL('chunk-01406d83.js', import.meta.url).href;\n\nexport default input;\n`
+				);
+				assert.equal(output[1].fileName, 'chunk-01406d83.js');
+				assert.equal(output[1].code, `console.log('chunk');\n`);
+			});
+	});
+
 	it('supports asset emission', () => {
 		return rollup
 			.rollup({