feat: support banner hook (#1419)

komyg · underfin · web-flow · commit 0631292b3c63 · 2024-07-16T13:07:20.000Z
### Description  Fix: #1417 issue. --------- Co-authored-by: underfin <likui.underfin@gmail.com>
diff --git a/crates/rolldown/src/ecmascript/ecma_generator.rs b/crates/rolldown/src/ecmascript/ecma_generator.rs
@@ -15,6 +15,7 @@ use rolldown_common::{
   AssetMeta, ChunkKind, EcmaAssetMeta, ExportsKind, OutputFormat, PreliminaryAsset, RenderedModule,
   WrapKind,
 };
+use rolldown_plugin::HookBannerArgs;
 use rolldown_sourcemap::{ConcatSource, RawSource};
 use rolldown_utils::rayon::{IntoParallelRefIterator, ParallelIterator};
 use rustc_hash::FxHashMap;
@@ -164,12 +165,16 @@ impl Generator for EcmaGenerator {
     };
 
     // add banner
-    if let Some(banner) = ctx.options.banner.as_ref() {
-      if let Some(banner_txt) = banner.call(&rendered_chunk).await? {
-        if !banner_txt.is_empty() {
-          concat_source.add_prepend_source(Box::new(RawSource::new(banner_txt)));
-        }
-      }
+    let banner = match ctx.options.banner.as_ref() {
+      Some(banner) => banner.call(&rendered_chunk).await?,
+      None => None,
+    };
+    if let Some(banner) = ctx
+      .plugin_driver
+      .banner(HookBannerArgs { chunk: &rendered_chunk }, banner.unwrap_or_default())
+      .await?
+    {
+      concat_source.add_prepend_source(Box::new(RawSource::new(banner)));
     }
 
     // Add `use strict` directive if needed. This must come before the banner, because users might use banner to add hashbang.
diff --git a/crates/rolldown/src/stages/generate_stage/render_chunk_to_assets.rs b/crates/rolldown/src/stages/generate_stage/render_chunk_to_assets.rs
@@ -149,6 +149,7 @@ impl<'a> GenerateStage<'a> {
         options: self.options,
         link_output: self.link_output,
         chunk_graph,
+        plugin_driver: self.plugin_driver,
       };
       EcmaGenerator::render_preliminary_assets(&ctx).await
     }))
diff --git a/crates/rolldown/src/types/generator.rs b/crates/rolldown/src/types/generator.rs
@@ -1,4 +1,5 @@
 use rolldown_common::{Chunk, ChunkIdx, NormalizedBundlerOptions, PreliminaryAsset};
+use rolldown_plugin::SharedPluginDriver;
 
 use crate::{chunk_graph::ChunkGraph, stages::link_stage::LinkStageOutput};
 
@@ -8,6 +9,7 @@ pub struct GenerateContext<'a> {
   pub options: &'a NormalizedBundlerOptions,
   pub link_output: &'a LinkStageOutput,
   pub chunk_graph: &'a ChunkGraph,
+  pub plugin_driver: &'a SharedPluginDriver,
 }
 
 pub trait Generator {
diff --git a/crates/rolldown_binding/src/options/plugin/binding_plugin_options.rs b/crates/rolldown_binding/src/options/plugin/binding_plugin_options.rs
@@ -124,6 +124,10 @@ pub struct BindingPluginOptions {
     ts_type = "(ctx: BindingPluginContext, bundle: BindingOutputs) => MaybePromise<VoidNullable>"
   )]
   pub write_bundle: Option<MaybeAsyncJsCallback<(BindingPluginContext, BindingOutputs), ()>>,
+
+  #[serde(skip_deserializing)]
+  #[napi(ts_type = "(ctx: BindingPluginContext, chunk: RenderedChunk) => void")]
+  pub banner: Option<MaybeAsyncJsCallback<(BindingPluginContext, RenderedChunk), Option<String>>>,
 }
 
 impl Debug for BindingPluginOptions {
diff --git a/crates/rolldown_binding/src/options/plugin/js_plugin.rs b/crates/rolldown_binding/src/options/plugin/js_plugin.rs
@@ -162,6 +162,23 @@ impl Plugin for JsPlugin {
     Ok(())
   }
 
+  async fn banner(
+    &self,
+    ctx: &rolldown_plugin::SharedPluginContext,
+    args: &rolldown_plugin::HookBannerArgs,
+  ) -> rolldown_plugin::HookBannerOutputReturn {
+    if let Some(cb) = &self.banner {
+      Ok(
+        cb.await_call((Arc::clone(ctx).into(), args.chunk.clone().into()))
+          .await?
+          .map(TryInto::try_into)
+          .transpose()?,
+      )
+    } else {
+      Ok(None)
+    }
+  }
+
   async fn render_chunk(
     &self,
     ctx: &rolldown_plugin::SharedPluginContext,
diff --git a/crates/rolldown_plugin/src/lib.rs b/crates/rolldown_plugin/src/lib.rs
@@ -11,12 +11,14 @@ pub mod inner {
 
 pub use crate::{
   plugin::{
-    BoxPlugin, HookAugmentChunkHashReturn, HookLoadReturn, HookNoopReturn, HookRenderChunkReturn,
-    HookResolveIdReturn, HookTransformAstReturn, HookTransformReturn, Plugin, SharedPlugin,
+    BoxPlugin, HookAugmentChunkHashReturn, HookBannerOutputReturn, HookLoadReturn, HookNoopReturn,
+    HookRenderChunkReturn, HookResolveIdReturn, HookTransformAstReturn, HookTransformReturn,
+    Plugin, SharedPlugin,
   },
   plugin_context::{PluginContext, SharedPluginContext},
   plugin_driver::{PluginDriver, SharedPluginDriver},
   transform_plugin_context::TransformPluginContext,
+  types::hook_banner_args::HookBannerArgs,
   types::hook_build_end_args::HookBuildEndArgs,
   types::hook_load_args::HookLoadArgs,
   types::hook_load_output::HookLoadOutput,
diff --git a/crates/rolldown_plugin/src/plugin.rs b/crates/rolldown_plugin/src/plugin.rs
@@ -3,12 +3,13 @@ use std::{any::Any, borrow::Cow, fmt::Debug, sync::Arc};
 use super::plugin_context::SharedPluginContext;
 use crate::{
   transform_plugin_context::TransformPluginContext,
+  types::hook_render_error::HookRenderErrorArgs,
   types::{
-    hook_render_error::HookRenderErrorArgs, hook_transform_ast_args::HookTransformAstArgs,
-    hook_transform_output::HookTransformOutput,
+    hook_transform_ast_args::HookTransformAstArgs, hook_transform_output::HookTransformOutput,
   },
-  HookBuildEndArgs, HookLoadArgs, HookLoadOutput, HookRenderChunkArgs, HookRenderChunkOutput,
-  HookResolveDynamicImportArgs, HookResolveIdArgs, HookResolveIdOutput, HookTransformArgs,
+  HookBannerArgs, HookBuildEndArgs, HookLoadArgs, HookLoadOutput, HookRenderChunkArgs,
+  HookRenderChunkOutput, HookResolveDynamicImportArgs, HookResolveIdArgs, HookResolveIdOutput,
+  HookTransformArgs,
 };
 use anyhow::Result;
 use rolldown_common::{ModuleInfo, Output, RollupRenderedChunk};
@@ -21,6 +22,7 @@ pub type HookLoadReturn = Result<Option<HookLoadOutput>>;
 pub type HookNoopReturn = Result<()>;
 pub type HookRenderChunkReturn = Result<Option<HookRenderChunkOutput>>;
 pub type HookAugmentChunkHashReturn = Result<Option<String>>;
+pub type HookBannerOutputReturn = Result<Option<String>>;
 
 #[async_trait::async_trait]
 pub trait Plugin: Any + Debug + Send + Sync + 'static {
@@ -95,6 +97,14 @@ pub trait Plugin: Any + Debug + Send + Sync + 'static {
     Ok(())
   }
 
+  async fn banner(
+    &self,
+    _ctx: &SharedPluginContext,
+    _args: &HookBannerArgs,
+  ) -> HookBannerOutputReturn {
+    Ok(None)
+  }
+
   async fn render_chunk(
     &self,
     _ctx: &SharedPluginContext,
diff --git a/crates/rolldown_plugin/src/plugin_driver/output_hooks.rs b/crates/rolldown_plugin/src/plugin_driver/output_hooks.rs
@@ -1,7 +1,7 @@
 use crate::types::hook_render_error::HookRenderErrorArgs;
-use crate::PluginDriver;
 use crate::{HookAugmentChunkHashReturn, HookNoopReturn, HookRenderChunkArgs};
-use anyhow::Result;
+use crate::{HookBannerArgs, PluginDriver};
+use anyhow::{Ok, Result};
 use rolldown_common::{Output, RollupRenderedChunk};
 use rolldown_sourcemap::SourceMap;
 
@@ -13,6 +13,23 @@ impl PluginDriver {
     Ok(())
   }
 
+  pub async fn banner(
+    &self,
+    args: HookBannerArgs<'_>,
+    mut banner: String,
+  ) -> Result<Option<String>> {
+    for (plugin, ctx) in &self.plugins {
+      if let Some(r) = plugin.banner(ctx, &args).await? {
+        banner.push('\n');
+        banner.push_str(r.as_str());
+      }
+    }
+    if banner.is_empty() {
+      return Ok(None);
+    }
+    Ok(Some(banner))
+  }
+
   pub async fn render_chunk(
     &self,
     mut args: HookRenderChunkArgs<'_>,
diff --git a/crates/rolldown_plugin/src/types/hook_banner_args.rs b/crates/rolldown_plugin/src/types/hook_banner_args.rs
@@ -0,0 +1,6 @@
+use rolldown_common::RollupRenderedChunk;
+
+#[derive(Debug)]
+pub struct HookBannerArgs<'a> {
+  pub chunk: &'a RollupRenderedChunk,
+}
diff --git a/crates/rolldown_plugin/src/types/mod.rs b/crates/rolldown_plugin/src/types/mod.rs
@@ -1,3 +1,4 @@
+pub mod hook_banner_args;
 pub mod hook_build_end_args;
 pub mod hook_load_args;
 pub mod hook_load_output;
diff --git a/packages/rolldown/src/binding.d.ts b/packages/rolldown/src/binding.d.ts
@@ -223,6 +223,7 @@ export interface BindingPluginOptions {
   renderError?: (ctx: BindingPluginContext, error: string) => void
   generateBundle?: (ctx: BindingPluginContext, bundle: BindingOutputs, isWrite: boolean) => MaybePromise<VoidNullable>
   writeBundle?: (ctx: BindingPluginContext, bundle: BindingOutputs) => MaybePromise<VoidNullable>
+  banner?: (ctx: BindingPluginContext, chunk: RenderedChunk) => void
 }
 
 export interface BindingPluginWithIndex {
diff --git a/packages/rolldown/src/plugin/bindingify-output-hooks.ts b/packages/rolldown/src/plugin/bindingify-output-hooks.ts
@@ -147,3 +147,23 @@ export function bindingifyWriteBundle(
     )
   }
 }
+
+export function bindingifyBanner(
+  plugin: Plugin,
+  options: NormalizedInputOptions,
+  pluginContextData: PluginContextData,
+): BindingPluginOptions['banner'] {
+  const hook = plugin.banner
+  if (!hook) {
+    return undefined
+  }
+
+  const [handler, _optionsIgnoredSofar] = normalizeHook(hook)
+
+  return async (ctx, chunk) => {
+    return handler.call(
+      new PluginContext(options, ctx, plugin, pluginContextData),
+      chunk,
+    )
+  }
+}
diff --git a/packages/rolldown/src/plugin/bindingify-plugin.ts b/packages/rolldown/src/plugin/bindingify-plugin.ts
@@ -16,6 +16,7 @@ import {
   bindingifyWriteBundle,
   bindingifyRenderError,
   bindingifyAugmentChunkHash,
+  bindingifyBanner,
 } from './bindingify-output-hooks'
 
 import type { Plugin } from './index'
@@ -73,5 +74,6 @@ export function bindingifyPlugin(
       outputOptions,
       pluginContextData,
     ),
+    banner: bindingifyBanner(plugin, options, pluginContextData),
   }
 }
diff --git a/packages/rolldown/src/plugin/index.ts b/packages/rolldown/src/plugin/index.ts
@@ -12,7 +12,7 @@ import type {
 import type { SourceMapInput } from '../types/sourcemap'
 import type { ModuleInfo } from '../types/module-info'
 import type { OutputBundle } from '../types/output-bundle'
-import type { PluginContext } from './plugin-context'
+import { PluginContext } from './plugin-context'
 import type { TransformPluginContext } from './transfrom-plugin-context'
 import type { NormalizedOutputOptions } from '../options/normalized-output-options'
 import type { LogLevel } from '../log/logging'
@@ -178,7 +178,7 @@ export type SequentialPluginHooks =
   | 'renderChunk'
   | 'transform'
 
-export type AddonHooks = '' /* 'banner' | 'footer' | 'intro' | 'outro' */
+export type AddonHooks = 'banner' /*| 'footer' | 'intro' | 'outro' */
 
 export type OutputPluginHooks =
   | 'augmentChunkHash'
@@ -213,7 +213,7 @@ export type AddonHookFunction = (
   chunk: RenderedChunk,
 ) => string | Promise<string>
 
-export type AddonHook = string | AddonHookFunction
+export type AddonHook = AddonHookFunction
 
 export interface OutputPlugin
   extends Partial<{ [K in OutputPluginHooks]: PluginHooks[K] }>,
diff --git a/packages/rolldown/tests/fixtures/plugin/banner/_config.ts b/packages/rolldown/tests/fixtures/plugin/banner/_config.ts
@@ -0,0 +1,21 @@
+import { expect, vi } from 'vitest'
+import path from 'node:path'
+import { defineTest } from '@tests'
+
+const entry = path.join(__dirname, './main.js')
+
+export default defineTest({
+  config: {
+    input: entry,
+    plugins: [
+      {
+        name: 'test-plugin',
+        banner: () => '/* Lorem ipsum */',
+      },
+    ],
+  },
+  afterTest: (output) => {
+    expect(output.output[0].code).toContain('/* Lorem ipsum */')
+    // expect(true).toBe(true)
+  },
+})
diff --git a/packages/rolldown/tests/fixtures/plugin/banner/main.js b/packages/rolldown/tests/fixtures/plugin/banner/main.js
@@ -0,0 +1 @@
+console.log()

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+pub mod hook_banner_args;`
`1`	`2`	`pub mod hook_build_end_args;`
`2`	`3`	`pub mod hook_load_args;`
`3`	`4`	`pub mod hook_load_output;`