feat: concatenateWrappedModule

crates/rolldown/src/ecmascript/format/esm.rs

@@ -3,7 +3,7 @@ use itertools::Itertools;
 use rolldown_common::{AddonRenderContext, ExportsKind, Specifier};
 use rolldown_sourcemap::SourceJoiner;
 use rolldown_utils::{concat_string, ecmascript::to_module_import_export_name};
-use rustc_hash::FxHashSet;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use crate::{
   ecmascript::ecma_generator::{RenderedModuleSource, RenderedModuleSources},
@@ -67,15 +67,7 @@ pub fn render_esm<'code>(
   }
 
   // chunk content
-  module_sources.iter().for_each(
-    |RenderedModuleSource { sources: module_render_output, .. }| {
-      if let Some(emitted_sources) = module_render_output {
-        for source in emitted_sources.as_ref() {
-          source_joiner.append_source(source);
-        }
-      }
-    },
-  );
+  render_chunk_content(ctx, module_sources, &mut source_joiner);
 
   if let Some(source) = render_wrapped_entry_chunk(ctx, None) {
     source_joiner.append_source(source);
@@ -96,6 +88,106 @@ pub fn render_esm<'code>(
   source_joiner
 }
 
+fn render_chunk_content<'code>(
+  ctx: &GenerateContext<'_>,
+  module_sources: &'code [RenderedModuleSource],
+  source_joiner: &mut SourceJoiner<'code>,
+) {
+  // If there is no concatenate_wrapping_modules, just concate all modules by exec order.
+  if ctx.chunk.module_groups.is_empty() {
+    module_sources.iter().for_each(
+      |RenderedModuleSource { sources: module_render_output, .. }| {
+        if let Some(emitted_sources) = module_render_output {
+          for source in emitted_sources.as_ref() {
+            source_joiner.append_source(source);
+          }
+        }
+      },
+    );
+    return;
+  }
+  let module_idx_to_source_idx = module_sources.iter().enumerate().fold(
+    FxHashMap::default(),
+    |mut acc, (idx, module_source)| {
+      acc.insert(module_source.module_idx, idx);
+      acc
+    },
+  );
+
+  for group in &ctx.chunk.module_groups {
+    // If the group is not belong to any concatenated module, we just render it as a single module.
+    if group.modules.len() == 1 {
+      let source =
+        module_sources.get(module_idx_to_source_idx[&group.entry]).expect("should have source");
+      if let Some(emitted_sources) = source.sources.as_ref() {
+        for source in emitted_sources.as_ref() {
+          source_joiner.append_source(source);
+        }
+      }
+      continue;
+    }
+    let (hoisted_fns, hoisted_vars) = group
+      .modules
+      .iter()
+      .filter_map(|idx| {
+        let render_concatenated_module =
+          ctx.chunk.module_idx_to_render_concatenated_module.get(idx)?;
+        let hoisted_fns = render_concatenated_module.hoisted_functions_or_module_ns_decl.join("");
+        let hoisted_vars = render_concatenated_module.hoisted_vars.join(", ");
+        Some((hoisted_fns, hoisted_vars))
+      })
+      .fold(
+        (String::new(), String::new()),
+        |(mut acc_hoisted_fns, mut acc_hoisted_vars), (hoisted_fn, hoisted_vars)| {
+          acc_hoisted_fns += &hoisted_fn;
+          if !hoisted_vars.is_empty() && !acc_hoisted_vars.is_empty() {
+            acc_hoisted_vars += ", ";
+          }
+          acc_hoisted_vars += &hoisted_vars;
+          (acc_hoisted_fns, acc_hoisted_vars)
+        },
+      );
+
+    let entry_module_stable_id = ctx.link_output.module_table[group.entry].stable_id();
+    // render var init_entry = __esm("", () => {
+    let rendered_esm_runtime_expr = ctx.chunk.module_idx_to_render_concatenated_module
+      [&group.entry]
+      .rendered_esm_runtime_expr
+      .as_ref()
+      .unwrap()
+      .trim_end_matches([';', '\n']);
+    let wrap_name = ctx.chunk.module_idx_to_render_concatenated_module[&group.entry]
+      .wrap_ref_name
+      .as_ref()
+      .unwrap();
+
+    source_joiner.append_source(hoisted_fns);
+    if !hoisted_vars.is_empty() {
+      source_joiner.append_source(concat_string!("var ", hoisted_vars, ";"));
+    }
+    source_joiner.append_source(concat_string!(
+      "var ",
+      wrap_name,
+      " = ",
+      rendered_esm_runtime_expr,
+      "({\"",
+      entry_module_stable_id,
+      "\": () => {"
+    ));
+    // we render each module in the group by exec order.
+    group.modules.iter().for_each(|module_idx| {
+      if let Some(rendered) =
+        module_sources.get(module_idx_to_source_idx[module_idx]).and_then(|m| m.sources.as_ref())
+      {
+        for source in rendered.iter() {
+          source_joiner.append_source(source);
+        }
+      }
+    });
+    source_joiner.append_source("}});\n");
+  }
+}
+
 fn render_esm_chunk_imports(ctx: &GenerateContext<'_>) -> Option<String> {
   let mut s = String::new();
   ctx.chunk.imports_from_other_chunks.iter().for_each(|(exporter_id, items)| {

crates/rolldown/src/module_finalizers/finalizer_context.rs

@@ -1,6 +1,6 @@
 use rolldown_common::{
   ChunkIdx, ConstExportMeta, ImportRecordIdx, IndexModules, ModuleIdx, NormalModule,
-  RuntimeModuleBrief, SharedFileEmitter, SymbolRef, SymbolRefDb,
+  RenderedConcatenatedModuleParts, RuntimeModuleBrief, SharedFileEmitter, SymbolRef, SymbolRefDb,
 };
 
 use oxc::span::CompactStr;
@@ -32,4 +32,5 @@ pub struct ScopeHoistingFinalizerContext<'me> {
   pub needs_hosted_top_level_binding: bool,
   pub module_namespace_included: bool,
   pub transferred_import_record: FxIndexMap<ImportRecordIdx, String>,
+  pub rendered_concatenated_wrapped_module_parts: RenderedConcatenatedModuleParts,
 }

crates/rolldown/src/module_finalizers/impl_visit_mut.rs

@@ -1,19 +1,22 @@
+use itertools::Itertools;
 use oxc::allocator::FromIn;
-use oxc::span::Atom;
+use oxc::span::{Atom, CompactStr};
 use oxc::{
   allocator::{self, IntoIn, TakeIn},
   ast::{
     NONE,
-    ast::{self, BindingPatternKind, Expression, SimpleAssignmentTarget},
+    ast::{self, BindingPatternKind, Expression, SimpleAssignmentTarget, Statement},
     match_member_expression,
   },
   ast_visit::{VisitMut, walk_mut},
   semantic::ScopeFlags,
   span::{SPAN, Span},
 };
 use rolldown_common::{
-  ExportsKind, ModuleNamespaceIncludedReason, SymbolRef, ThisExprReplaceKind, WrapKind,
+  ConcatenateWrappedModuleKind, ExportsKind, ModuleNamespaceIncludedReason, SymbolRef,
+  ThisExprReplaceKind, WrapKind,
 };
+use rolldown_ecmascript::ToSourceString;
 use rolldown_ecmascript_utils::{ExpressionExt, JsxExt};
 
 use crate::hmr::utils::HmrAstBuilder;
@@ -154,14 +157,10 @@ impl<'ast> VisitMut<'ast> for ScopeHoistingFinalizer<'_, 'ast> {
         ));
       }
       Some(WrapKind::Esm) => {
-        use ast::Statement;
-        let wrap_ref_name = self.canonical_name_for(self.ctx.linking_info.wrapper_ref.unwrap());
-        let esm_ref = if self.ctx.options.profiler_names {
-          self.canonical_ref_for_runtime("__esm")
-        } else {
-          self.canonical_ref_for_runtime("__esmMin")
-        };
-        let esm_ref_expr = self.finalized_expr_for_symbol_ref(esm_ref, false, false);
+        let is_concatenated_wrapped_module = !matches!(
+          self.ctx.linking_info.concatenated_wrapped_module_kind,
+          ConcatenateWrappedModuleKind::None
+        );
         let old_body = program.body.take_in(self.alloc);
 
         let mut fn_stmts = allocator::Vec::new_in(self.alloc);
@@ -184,9 +183,25 @@ impl<'ast> VisitMut<'ast> for ScopeHoistingFinalizer<'_, 'ast> {
               stmts_inside_closure.push(stmt);
             }
           });
-        program.body.extend(declaration_of_module_namespace_object);
-        program.body.extend(fn_stmts);
-        if !self.top_level_var_bindings.is_empty() {
+
+        if is_concatenated_wrapped_module {
+          self.ctx.rendered_concatenated_wrapped_module_parts.hoisted_functions_or_module_ns_decl =
+            declaration_of_module_namespace_object
+              .iter()
+              .chain(fn_stmts.iter())
+              .map(rolldown_ecmascript::ToSourceString::to_source_string)
+              .collect_vec();
+          self.ctx.rendered_concatenated_wrapped_module_parts.hoisted_vars = self
+            .top_level_var_bindings
+            .iter()
+            .map(|var_name| CompactStr::new(var_name))
+            .collect_vec();
+        } else {
+          program.body.extend(declaration_of_module_namespace_object);
+          program.body.extend(fn_stmts);
+        }
+
+        if !is_concatenated_wrapped_module && !self.top_level_var_bindings.is_empty() {
           let builder = self.builder();
           let decorations = self.top_level_var_bindings.iter().map(|var_name| {
             builder.variable_declarator(
@@ -210,6 +225,36 @@ impl<'ast> VisitMut<'ast> for ScopeHoistingFinalizer<'_, 'ast> {
             false,
           )));
         }
+
+        // The wrapping would happen during the chunk codegen phase
+        if matches!(
+          self.ctx.linking_info.concatenated_wrapped_module_kind,
+          ConcatenateWrappedModuleKind::Inner
+        ) {
+          program.body.extend(stmts_inside_closure);
+          return;
+        }
+
+        let esm_ref = if self.ctx.options.profiler_names {
+          self.canonical_ref_for_runtime("__esm")
+        } else {
+          self.canonical_ref_for_runtime("__esmMin")
+        };
+        let esm_ref_expr = self.finalized_expr_for_symbol_ref(esm_ref, false, false);
+        let wrap_ref_name = self.canonical_name_for(self.ctx.linking_info.wrapper_ref.unwrap());
+
+        if matches!(
+          self.ctx.linking_info.concatenated_wrapped_module_kind,
+          ConcatenateWrappedModuleKind::Root
+        ) {
+          self.ctx.rendered_concatenated_wrapped_module_parts.rendered_esm_runtime_expr =
+            Some(self.builder().expression_statement(SPAN, esm_ref_expr).to_source_string());
+          self.ctx.rendered_concatenated_wrapped_module_parts.wrap_ref_name =
+            Some(wrap_ref_name.clone());
+          program.body.extend(stmts_inside_closure);
+          return;
+        }
+
         program.body.push(self.snippet.esm_wrapper_stmt(
           wrap_ref_name,
           esm_ref_expr,

crates/rolldown/src/module_finalizers/mod.rs

@@ -12,9 +12,9 @@ use oxc::{
   span::{Atom, GetSpan, GetSpanMut, SPAN},
 };
 use rolldown_common::{
-  AstScopes, ExportsKind, ImportRecordIdx, ImportRecordMeta, MemberExprRefResolution, Module,
-  ModuleIdx, ModuleNamespaceIncludedReason, ModuleType, OutputFormat, Platform, SymbolRef,
-  WrapKind,
+  AstScopes, ConcatenateWrappedModuleKind, ExportsKind, ImportRecordIdx, ImportRecordMeta,
+  MemberExprRefResolution, Module, ModuleIdx, ModuleNamespaceIncludedReason, ModuleType,
+  OutputFormat, Platform, SymbolRef, WrapKind,
 };
 use rolldown_ecmascript::ToSourceString;
 use rolldown_ecmascript_utils::{
@@ -105,6 +105,7 @@ impl<'me, 'ast> ScopeHoistingFinalizer<'me, 'ast> {
         // import React from './node_modules/react/index.js';
         // ```
         if rec.meta.contains(ImportRecordMeta::SafelyMergeCjsNs)
+          // TODO: merge cjs namepspace in module group level
           && self.ctx.linking_infos[self.ctx.module.idx].wrap_kind.is_none()
         {
           let chunk_idx = self.ctx.chunk_id;
@@ -154,7 +155,11 @@ impl<'me, 'ast> ScopeHoistingFinalizer<'me, 'ast> {
       // Replace the import statement with `init_foo()` if `ImportDeclaration` is not a plain import
       // or the importee have side effects.
       WrapKind::Esm => {
-        if self.generated_init_esm_importee_ids.contains(&importee.idx) {
+        if matches!(
+          importee_linking_info.concatenated_wrapped_module_kind,
+          ConcatenateWrappedModuleKind::Inner
+        ) || self.generated_init_esm_importee_ids.contains(&importee.idx)
+        {
           return true;
         }
         self.generated_init_esm_importee_ids.insert(importee.idx);

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

@@ -11,8 +11,9 @@ use rolldown_std_utils::OptionExt;
 use rustc_hash::FxHashMap;
 
 use rolldown_common::{
-  ChunkIdx, ChunkKind, CssAssetNameReplacer, ImportMetaRolldownAssetReplacer, ImportRecordIdx,
-  Module, ModuleIdx, PreliminaryFilename, PrependRenderedImport, RollupPreRenderedAsset,
+  ChunkIdx, ChunkKind, ConcatenateWrappedModuleKind, CssAssetNameReplacer,
+  ImportMetaRolldownAssetReplacer, ImportRecordIdx, Module, ModuleIdx, PreliminaryFilename,
+  PrependRenderedImport, RenderedConcatenatedModuleParts, RollupPreRenderedAsset,
 };
 use rolldown_plugin::SharedPluginDriver;
 use rolldown_std_utils::{PathBufExt, PathExt, representative_file_name_for_preserve_modules};
@@ -64,6 +65,7 @@ impl<'a> GenerateStage<'a> {
   }
 
   #[tracing::instrument(level = "debug", skip_all)]
+  #[allow(clippy::too_many_lines)]
   pub async fn generate(&mut self) -> BuildResult<BundleOutput> {
     self.plugin_driver.render_start(self.options).await?;
 
@@ -160,9 +162,19 @@ impl<'a> GenerateStage<'a> {
               idxs.into_iter().map(|idx| (idx, String::new())).collect::<FxIndexMap<_, _>>()
             })
             .unwrap_or_default(),
+          rendered_concatenated_wrapped_module_parts: RenderedConcatenatedModuleParts::default(),
         };
         let ctx = finalize_normal_module(ctx, ast, ast_scope);
-        (!ctx.transferred_import_record.is_empty()).then_some((idx, ctx.transferred_import_record))
+        (!ctx.transferred_import_record.is_empty()
+          || !matches!(
+            ctx.linking_info.concatenated_wrapped_module_kind,
+            ConcatenateWrappedModuleKind::None
+          ))
+        .then_some((
+          idx,
+          ctx.transferred_import_record,
+          ctx.rendered_concatenated_wrapped_module_parts,
+        ))
       })
       .collect::<Vec<_>>();
 
@@ -173,13 +185,28 @@ impl<'a> GenerateStage<'a> {
   fn apply_transfer_parts_mutation(
     &mut self,
     chunk_graph: &mut ChunkGraph,
-    transfer_parts_rendered_maps: Vec<(ModuleIdx, FxIndexMap<ImportRecordIdx, String>)>,
+    transfer_parts_rendered_maps: Vec<(
+      ModuleIdx,
+      FxIndexMap<ImportRecordIdx, String>,
+      RenderedConcatenatedModuleParts,
+    )>,
   ) {
     let mut normalized_transfer_parts_rendered_maps = FxHashMap::default();
-    for (idx, transferred_import_record) in transfer_parts_rendered_maps {
+    for (idx, transferred_import_record, rendered_concatenated_module_parts) in
+      transfer_parts_rendered_maps
+    {
       for (rec_idx, rendered_string) in transferred_import_record {
         normalized_transfer_parts_rendered_maps.insert((idx, rec_idx), rendered_string);
       }
+      let chunk_idx = chunk_graph.module_to_chunk[idx].expect("should have chunk idx");
+      let chunk = &mut chunk_graph.chunk_table[chunk_idx];
+      chunk
+        .module_idx_to_render_concatenated_module
+        .insert(idx, rendered_concatenated_module_parts);
+    }
+
+    if normalized_transfer_parts_rendered_maps.is_empty() {
+      return;
     }
     for chunk in chunk_graph.chunk_table.iter_mut() {
       for (module_idx, recs) in &chunk.insert_map {