feat: support return sourcemap in render chunk hook (#776)

crates/rolldown/src/chunk/mod.rs

@@ -4,6 +4,8 @@ pub mod render_chunk;
 mod render_chunk_exports;
 mod render_chunk_imports;
 
+use std::sync::Arc;
+
 use index_vec::IndexVec;
 use path_slash::PathBufExt;
 use rolldown_common::{ChunkId, FileNameTemplate};
@@ -48,7 +50,7 @@ pub struct Chunk {
 
 pub struct ChunkRenderReturn {
   pub code: String,
-  pub map: Option<SourceMap>,
+  pub map: Option<Arc<SourceMap>>,
   pub rendered_chunk: RenderedChunk,
 }
 
@@ -105,7 +107,7 @@ impl Chunk {
           // TODO(underfin): refactor the relative path
           m.resource_id.expect_file().relative_path(file_dir).to_slash_lossy().as_ref(),
           options,
-          file_dir,
+          file_dir.to_string_lossy().as_ref(),
         )
       })
       .collect::<Vec<_>>()
@@ -143,8 +145,12 @@ impl Chunk {
       concat_source.add_source(Box::new(RawSource::new(footer_txt)));
     }
 
-    let (content, map) = concat_source.content_and_sourcemap();
+    let (content, mut map) = concat_source.content_and_sourcemap();
+
+    if let Some(x) = map.as_mut() {
+      x.set_file(&rendered_chunk.file_name);
+    }
 
-    Ok(ChunkRenderReturn { code: content, map, rendered_chunk })
+    Ok(ChunkRenderReturn { code: content, map: map.map(Arc::new), rendered_chunk })
   }
 }

crates/rolldown/src/stages/bundle_stage/mod.rs

@@ -93,9 +93,8 @@ impl<'a> BundleStage<'a> {
 
     render_chunks(self.plugin_driver, chunks).await?.into_iter().try_for_each(
       |chunk| -> Result<(), BuildError> {
-        let ChunkRenderReturn { mut map, rendered_chunk, mut code } = chunk;
-        if let Some(map) = map.as_mut() {
-          map.set_file(rendered_chunk.file_name.as_str());
+        let ChunkRenderReturn { map, rendered_chunk, mut code } = chunk;
+        if let Some(map) = map.as_ref() {
           match self.options.sourcemap {
             SourceMapType::File => {
               let map_file_name = format!("{}.map", rendered_chunk.file_name);

crates/rolldown/src/utils/render_chunks.rs

@@ -1,5 +1,6 @@
 use rolldown_common::IntoBatchedResult;
 use rolldown_plugin::{HookRenderChunkArgs, SharedPluginDriver};
+use rolldown_sourcemap::collapse_sourcemaps;
 use rolldown_utils::block_on_spawn_all;
 
 use crate::{chunk::ChunkRenderReturn, error::BatchedErrors};
@@ -8,16 +9,25 @@ pub async fn render_chunks<'a>(
   plugin_driver: &SharedPluginDriver,
   chunks: Vec<ChunkRenderReturn>,
 ) -> Result<Vec<ChunkRenderReturn>, BatchedErrors> {
-  // TODO support `render_chunk` hook return map
   let result = block_on_spawn_all(chunks.into_iter().map(|chunk| async move {
     tracing::info!("render_chunks");
+    let mut sourcemap_chain = vec![];
+    if let Some(sourcemap) = chunk.map {
+      sourcemap_chain.push(sourcemap);
+    }
+
     match plugin_driver
-      .render_chunk(HookRenderChunkArgs { code: chunk.code, chunk: &chunk.rendered_chunk })
+      .render_chunk(
+        HookRenderChunkArgs { code: chunk.code, chunk: &chunk.rendered_chunk },
+        &mut sourcemap_chain,
+      )
       .await
     {
-      Ok(code) => {
-        Ok(ChunkRenderReturn { code, map: chunk.map, rendered_chunk: chunk.rendered_chunk })
-      }
+      Ok(code) => Ok(ChunkRenderReturn {
+        code,
+        map: collapse_sourcemaps(sourcemap_chain, None),
+        rendered_chunk: chunk.rendered_chunk,
+      }),
       Err(e) => Err(e),
     }
   }));

crates/rolldown/src/utils/render_normal_module.rs

@@ -1,5 +1,4 @@
-use std::{path::Path, sync::Arc};
-// cSpell:disable
+use std::sync::Arc;
 
 use rolldown_common::{NormalModule, RenderedModule};
 use rolldown_oxc_utils::{OxcCompiler, OxcProgram};
@@ -12,7 +11,7 @@ pub fn render_normal_module<'a>(
   ast: &OxcProgram,
   source_name: &str,
   options: &SharedOptions,
-  file_dir: impl AsRef<Path>,
+  file_dir: &str,
 ) -> Option<ModuleRenderOutput<'a>> {
   if ast.program().body.is_empty() {
     None
@@ -36,7 +35,7 @@ pub fn render_normal_module<'a>(
         if let Some(sourcemap) = render_output.source_map {
           sourcemap_chain.push(Arc::new(sourcemap));
         }
-        collapse_sourcemaps(sourcemap_chain, file_dir)
+        collapse_sourcemaps(sourcemap_chain, Some(file_dir))
       },
     })
   }

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

@@ -107,7 +107,12 @@ impl Plugin for JsPlugin {
     args: &rolldown_plugin::HookRenderChunkArgs,
   ) -> rolldown_plugin::HookRenderChunkReturn {
     if let Some(cb) = &self.render_chunk {
-      Ok(cb.await_call((args.code.to_string(), args.chunk.clone().into())).await?.map(Into::into))
+      Ok(
+        cb.await_call((args.code.to_string(), args.chunk.clone().into()))
+          .await?
+          .map(TryInto::try_into)
+          .transpose()?,
+      )
     } else {
       Ok(None)
     }

crates/rolldown_binding/src/options/plugin/types/binding_hook_render_chunk_output.rs

@@ -1,4 +1,5 @@
 use derivative::Derivative;
+use rolldown_error::BuildError;
 use serde::Deserialize;
 
 #[napi_derive::napi(object)]
@@ -7,10 +8,22 @@ use serde::Deserialize;
 #[derivative(Debug)]
 pub struct BindingHookRenderChunkOutput {
   pub code: String,
+  pub map: Option<String>,
 }
 
-impl From<BindingHookRenderChunkOutput> for rolldown_plugin::HookRenderChunkOutput {
-  fn from(value: BindingHookRenderChunkOutput) -> Self {
-    Self { code: value.code }
+impl TryFrom<BindingHookRenderChunkOutput> for rolldown_plugin::HookRenderChunkOutput {
+  type Error = BuildError;
+
+  fn try_from(value: BindingHookRenderChunkOutput) -> Result<Self, Self::Error> {
+    Ok(rolldown_plugin::HookRenderChunkOutput {
+      code: value.code,
+      map: value
+        .map
+        .map(|content| {
+          rolldown_sourcemap::SourceMap::from_json_string(&content)
+            .map_err(BuildError::sourcemap_error)
+        })
+        .transpose()?,
+    })
   }
 }

crates/rolldown_common/src/types/output_chunk.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use rolldown_sourcemap::SourceMap;
 use rustc_hash::FxHashMap;
 
@@ -17,6 +19,6 @@ pub struct OutputChunk {
   pub modules: FxHashMap<String, RenderedModule>,
   // OutputChunk
   pub code: String,
-  pub map: Option<SourceMap>,
+  pub map: Option<Arc<SourceMap>>,
   pub sourcemap_file_name: Option<String>,
 }

crates/rolldown_plugin/src/plugin_driver/build_hooks.rs

@@ -64,10 +64,14 @@ impl PluginDriver {
   pub async fn render_chunk(
     &self,
     mut args: HookRenderChunkArgs<'_>,
+    sourcemap_chain: &mut Vec<Arc<SourceMap>>,
   ) -> Result<String, BuildError> {
     for (plugin, ctx) in &self.plugins {
       if let Some(r) = plugin.render_chunk(ctx, &args).await? {
         args.code = r.code;
+        if let Some(map) = r.map {
+          sourcemap_chain.push(Arc::new(map));
+        }
       }
     }
     Ok(args.code)

crates/rolldown_plugin/src/types/hook_render_chunk_output.rs

@@ -1,4 +1,7 @@
+use rolldown_sourcemap::SourceMap;
+
 #[derive(Debug)]
 pub struct HookRenderChunkOutput {
   pub code: String,
+  pub map: Option<SourceMap>,
 }

crates/rolldown_sourcemap/src/lib.rs

@@ -1,5 +1,5 @@
 use path_slash::PathBufExt;
-use std::{path::Path, sync::Arc};
+use std::sync::Arc;
 use sugar_path::{AsPath, SugarPath};
 
 // cSpell:disable
@@ -11,7 +11,7 @@ mod concat_sourcemap;
 
 pub fn collapse_sourcemaps(
   mut sourcemap_chain: Vec<Arc<SourceMap>>,
-  file_dir: impl AsRef<Path>,
+  file_dir: Option<&str>,
 ) -> Option<Arc<SourceMap>> {
   let last_map = sourcemap_chain.pop()?;
   // If there is only one sourcemap, return it as result.
@@ -43,9 +43,13 @@ pub fn collapse_sourcemaps(
       let name_id = original_token.get_name().map(|name| sourcemap_builder.add_name(name));
 
       let source_id = original_token.get_source_and_content().map(|(source, source_content)| {
-        let relative_path = source.as_path().relative(&file_dir);
-        sourcemap_builder
-          .add_source_and_content(relative_path.to_slash_lossy().as_ref(), source_content)
+        if let Some(file_dir) = file_dir.as_ref() {
+          let relative_path = source.as_path().relative(file_dir);
+          sourcemap_builder
+            .add_source_and_content(relative_path.to_slash_lossy().as_ref(), source_content)
+        } else {
+          sourcemap_builder.add_source_and_content(source, source_content)
+        }
       });
 
       sourcemap_builder.add_token(
@@ -59,7 +63,14 @@ pub fn collapse_sourcemaps(
     }
   }
 
-  Some(Arc::new(sourcemap_builder.into_sourcemap()))
+  // TODO(underfin) using sourcemap_builder.set_file
+  let mut map = sourcemap_builder.into_sourcemap();
+
+  if let Some(file) = sourcemap_chain.first().and_then(|x| x.get_file()) {
+    map.set_file(file);
+  }
+
+  Some(Arc::new(map))
 }
 
 #[cfg(test)]
@@ -99,7 +110,7 @@ mod tests {
     ];
 
     let result = {
-      let map = super::collapse_sourcemaps(sourcemaps, "/project/dist").unwrap();
+      let map = super::collapse_sourcemaps(sourcemaps, Some("/project/dist")).unwrap();
       map.to_json_string().unwrap()
     };
 

cspell.json

@@ -112,6 +112,7 @@
     "vite",
     "vitepress",
     "Yinan",
-    "Yunfei"
+    "Yunfei",
+    "sourcemaps"
   ]
 }

packages/rolldown/src/binding.d.ts

@@ -49,6 +49,7 @@ export interface BindingHookLoadOutput {
 
 export interface BindingHookRenderChunkOutput {
   code: string
+  map?: string
 }
 
 export interface BindingHookResolveIdExtraOptions {

packages/rolldown/src/plugin/bindingify-build-hooks.ts

@@ -164,8 +164,22 @@ export function bindingifyRenderChunk(
       return
     }
 
+    if (typeof ret === 'string') {
+      return { code: ret }
+    }
+
+    if (!ret.map) {
+      return { code: ret.code }
+    }
+
+    let map =
+      typeof ret.map === 'object'
+        ? ret.map
+        : (JSON.parse(ret.map) as SourceMapInputObject)
+
     return {
-      code: ret,
+      code: ret.code,
+      map: JSON.stringify(map),
     }
   }
 }

packages/rolldown/src/plugin/index.ts

@@ -79,7 +79,14 @@ export interface Plugin {
       this: null,
       code: string,
       chunk: RenderedChunk,
-    ) => MaybePromise<NullValue | string>
+    ) => MaybePromise<
+      | NullValue
+      | string
+      | {
+          code: string
+          map?: string | null | SourceMapInput
+        }
+    >
   >
 
   buildEnd?: Hook<(this: null, err?: string) => MaybePromise<NullValue>>

packages/rollup-tests/src/failed-tests.json

@@ -721,7 +721,6 @@
   "rollup@sourcemaps@names: names are recovered (https://github.com/rollup/rollup/issues/101)@generates es",
   "rollup@sourcemaps@reified-namespace: generates correct sourcemap with reified namespace (#668)@generates es",
   "rollup@sourcemaps@render-chunk-babili: generates valid sourcemap when source could not be determined@generates es",
-  "rollup@sourcemaps@render-chunk: preserves sourcemap chains when transforming@generates es",
   "rollup@sourcemaps@single-length-segments: handles single-length sourcemap segments@generates es",
   "rollup@sourcemaps@sourcemap-base-url-without-trailing-slash: add a trailing slash automatically if it is missing@generates es",
   "rollup@sourcemaps@sourcemap-base-url: adds a sourcemap base url@generates es",

packages/rollup-tests/src/status.json

@@ -1,6 +1,6 @@
 {
   "failed": 0,
-  "skipFailed": 676,
+  "skipFailed": 675,
   "skipped": 0,
-  "passed": 243
+  "passed": 244
 }

packages/rollup-tests/src/status.md

@@ -1,6 +1,6 @@
 |  | number |
 |----| ---- |
 | failed | 0|
-| skipFailed | 676|
+| skipFailed | 675|
 | skipped | 0|
-| passed | 243|
+| passed | 244|