feat: add writeBundle hook (#405)

crates/rolldown/src/bundler/bundler.rs

@@ -75,6 +75,8 @@ impl<T: FileSystem + Default + 'static> Bundler<T> {
 
     let output = self.bundle_up(output_options, true).await?;
 
+    self.plugin_driver.write_bundle(&output.assets).await?;
+
     self.fs.create_dir_all(dir.as_path()).unwrap_or_else(|_| {
       panic!(
         "Could not create directory for output chunks: {:?} \ncwd: {}",

crates/rolldown/src/bundler/plugin_driver/mod.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use rolldown_error::BuildError;
+use rolldown_utils::block_on_spawn_all;
 
 use crate::{
   plugin::{
@@ -78,4 +79,20 @@ impl PluginDriver {
     }
     Ok(())
   }
+
+  #[allow(clippy::unused_async)]
+  pub async fn write_bundle(&self, bundle: &Vec<Output>) -> HookNoopReturn {
+    let result = block_on_spawn_all(self.plugins.iter().map(|plugin| async move {
+      match plugin.write_bundle(&PluginContext::new(), bundle).await {
+        Ok(()) => Ok(()),
+        Err(e) => Err(e),
+      }
+    }));
+
+    for value in result {
+      value?;
+    }
+
+    Ok(())
+  }
 }

crates/rolldown/src/plugin/plugin.rs

@@ -71,6 +71,12 @@ pub trait Plugin: Debug + Send + Sync {
   ) -> HookNoopReturn {
     Ok(())
   }
+
+  // Parallel hook
+  #[allow(clippy::ptr_arg)]
+  async fn write_bundle(&self, _ctx: &PluginContext, _bundle: &Vec<Output>) -> HookNoopReturn {
+    Ok(())
+  }
 }
 
 pub type BoxPlugin = Box<dyn Plugin>;

crates/rolldown_binding/index.d.ts

@@ -19,6 +19,7 @@ export interface PluginOptions {
     chunk: RenderedChunk,
   ) => Promise<undefined | HookRenderChunkOutput>
   generateBundle?: (bundle: Outputs, isWrite: boolean) => Promise<void>
+  writeBundle?: (bundle: Outputs) => Promise<void>
 }
 export interface HookResolveIdArgsOptions {
   isEntry: boolean

crates/rolldown_binding/src/options/input_options/plugin.rs

@@ -49,6 +49,11 @@ pub struct PluginOptions {
   #[serde(skip_deserializing)]
   #[napi(ts_type = "(bundle: Outputs, isWrite: boolean) => Promise<void>")]
   pub generate_bundle: Option<JsFunction>,
+
+  #[derivative(Debug = "ignore")]
+  #[serde(skip_deserializing)]
+  #[napi(ts_type = "(bundle: Outputs) => Promise<void>")]
+  pub write_bundle: Option<JsFunction>,
 }
 
 #[napi_derive::napi(object)]

crates/rolldown_binding/src/options/input_options/plugin_adapter.rs

@@ -18,6 +18,7 @@ pub type TransformCallback = JsCallback<(String, String), Option<SourceResult>>;
 pub type BuildEndCallback = JsCallback<(Option<String>,), ()>;
 pub type RenderChunkCallback = JsCallback<(String, RenderedChunk), Option<HookRenderChunkOutput>>;
 pub type GenerateBundleCallback = JsCallback<(Outputs, bool), Option<HookRenderChunkOutput>>;
+pub type WriteBundleCallback = JsCallback<(Outputs,), ()>;
 
 #[derive(Derivative)]
 #[derivative(Debug)]
@@ -37,6 +38,8 @@ pub struct JsAdapterPlugin {
   render_chunk_fn: Option<RenderChunkCallback>,
   #[derivative(Debug = "ignore")]
   generate_bundle_fn: Option<GenerateBundleCallback>,
+  #[derivative(Debug = "ignore")]
+  write_bundle_fn: Option<WriteBundleCallback>,
 }
 
 impl JsAdapterPlugin {
@@ -49,6 +52,7 @@ impl JsAdapterPlugin {
     let render_chunk_fn = option.render_chunk.as_ref().map(RenderChunkCallback::new).transpose()?;
     let generate_bundle_fn =
       option.generate_bundle.as_ref().map(GenerateBundleCallback::new).transpose()?;
+    let write_bundle_fn = option.write_bundle.as_ref().map(WriteBundleCallback::new).transpose()?;
     Ok(Self {
       name: option.name,
       build_start_fn,
@@ -58,6 +62,7 @@ impl JsAdapterPlugin {
       build_end_fn,
       render_chunk_fn,
       generate_bundle_fn,
+      write_bundle_fn,
     })
   }
 
@@ -175,4 +180,16 @@ impl Plugin for JsAdapterPlugin {
     }
     Ok(())
   }
+
+  #[allow(clippy::redundant_closure_for_method_calls)]
+  async fn write_bundle(
+    &self,
+    _ctx: &rolldown::PluginContext,
+    bundle: &Vec<rolldown::Output>,
+  ) -> rolldown::HookNoopReturn {
+    if let Some(cb) = &self.write_bundle_fn {
+      cb.call_async((bundle.clone().into(),)).await.map_err(|e| e.into_bundle_error())?;
+    }
+    Ok(())
+  }
 }

packages/node/src/options/create-build-plugin-adapter.ts

@@ -7,7 +7,7 @@ import type {
   HookRenderChunkOutput,
   Outputs,
 } from '@rolldown/node-binding'
-import { transformToRollupOutput, unimplemented } from '../utils'
+import { transformToOutputBundle, unimplemented } from '../utils'
 
 // Note: because napi not catch error, so we need to catch error and print error to debugger in adapter.
 export function createBuildPluginAdapter(
@@ -23,6 +23,24 @@ export function createBuildPluginAdapter(
     buildEnd: buildEnd(plugin.buildEnd),
     renderChunk: renderChunk(plugin.renderChunk),
     generateBundle: generateBundle(plugin.generateBundle),
+    writeBundle: writeBundle(plugin.writeBundle),
+  }
+}
+
+function writeBundle(hook: Plugin['writeBundle']) {
+  if (hook) {
+    if (typeof hook !== 'function') {
+      return unimplemented()
+    }
+    return async (outputs: Outputs) => {
+      try {
+        // TODO outputOptions
+        await hook.call({} as any, {} as any, transformToOutputBundle(outputs))
+      } catch (error) {
+        console.error(error)
+        throw error
+      }
+    }
   }
 }
 
@@ -32,15 +50,14 @@ function generateBundle(hook: Plugin['generateBundle']) {
       return unimplemented()
     }
     return async (outputs: Outputs, isWrite: boolean) => {
-      const bundle = Object.fromEntries(
-        transformToRollupOutput(outputs).output.map((item) => [
-          item.fileName,
-          item,
-        ]),
-      )
       try {
         // TODO outputOptions
-        await hook.call({} as any, {} as any, bundle, isWrite)
+        await hook.call(
+          {} as any,
+          {} as any,
+          transformToOutputBundle(outputs),
+          isWrite,
+        )
       } catch (error) {
         console.error(error)
         throw error

packages/node/src/rollup-types.ts

@@ -9,4 +9,5 @@ export type {
   OutputChunk,
   NormalizedInputOptions,
   OutputAsset,
+  OutputBundle,
 } from 'rollup'

packages/node/src/utils/transform-to-rollup-output.ts

@@ -3,6 +3,7 @@ import type {
   RollupOutput,
   OutputChunk as RollupOutputChunk,
   OutputAsset as RollupOutputAsset,
+  OutputBundle,
 } from '../rollup-types'
 import { unimplemented } from '.'
 
@@ -85,3 +86,9 @@ export interface RolldownOutput {
     ...(RolldownOutputChunk | RolldownOutputAsset)[],
   ]
 }
+
+export function transformToOutputBundle(output: Outputs): OutputBundle {
+  return Object.fromEntries(
+    transformToRollupOutput(output).output.map((item) => [item.fileName, item]),
+  )
+}

packages/node/test/cases/plugin/write-bundle/config.ts

@@ -0,0 +1,37 @@
+import type { RollupOptions, RollupOutput } from '@rolldown/node'
+import { expect, vi } from 'vitest'
+import path from 'path'
+import { OutputChunk } from 'rollup'
+
+const entry = path.join(__dirname, './main.js')
+
+const writeBundleFn = vi.fn()
+
+const config: RollupOptions = {
+  input: entry,
+  plugins: [
+    {
+      name: 'test-plugin',
+      writeBundle: (options, bundle) => {
+        writeBundleFn()
+        const chunk = bundle['main.js'] as OutputChunk
+        expect(chunk.code.indexOf('console.log') > -1).toBe(true)
+        expect(chunk.type).toBe('chunk')
+        expect(chunk.fileName).toBe('main.js')
+        expect(chunk.isEntry).toBe(true)
+        expect(chunk.isDynamicEntry).toBe(false)
+        expect(chunk.facadeModuleId).toBe(entry)
+        expect(chunk.exports.length).toBe(0)
+        expect(chunk.moduleIds).toStrictEqual([entry])
+        expect(Object.keys(chunk.modules).length).toBe(1)
+      },
+    },
+  ],
+}
+
+export default {
+  config,
+  afterTest: (output: RollupOutput) => {
+    expect(writeBundleFn).toHaveBeenCalledTimes(1)
+  },
+}

packages/node/test/cases/plugin/write-bundle/main.js

@@ -0,0 +1 @@
+console.log()