feat: introduce ChunkGraph (#90)

.vscode/settings.json

@@ -24,6 +24,13 @@
     "*.snap": "markdown",
     "*.snap.new": "markdown"
   },
-  "cSpell.words": ["Finalizer", "importee", "rolldown", "rollup", "rustc"],
+  "cSpell.words": [
+    "Finalizer",
+    "hasher",
+    "importee",
+    "rolldown",
+    "rollup",
+    "rustc"
+  ],
   "git.ignoreLimitWarning": true
 }

crates/rolldown/src/bundler/bundle/bundle.rs

@@ -1,7 +1,9 @@
+use std::hash::BuildHasherDefault;
+
 use super::asset::Asset;
 use crate::bundler::{
   bitset::BitSet,
-  chunk::{chunk::Chunk, ChunkId, ChunksVec},
+  chunk::{chunk::Chunk, chunk_graph::ChunkGraph, ChunkId, ChunksVec},
   graph::graph::Graph,
   module::module::Module,
   options::{
@@ -24,115 +26,122 @@ impl<'a> Bundle<'a> {
     Self { graph, output_options }
   }
 
-  pub fn mark_modules_entry_bit(
+  fn determine_reachable_modules_for_entry(
     &self,
     module_id: ModuleId,
-    index: usize,
-    modules_entry_bit: &mut IndexVec<ModuleId, BitSet>,
+    entry_index: u32,
+    module_to_bits: &mut IndexVec<ModuleId, BitSet>,
   ) {
-    if modules_entry_bit[module_id].has_bit(index.try_into().unwrap()) {
+    if module_to_bits[module_id].has_bit(entry_index) {
       return;
     }
-    modules_entry_bit[module_id].set_bit(index.try_into().unwrap());
-    if let Module::Normal(m) = &self.graph.modules[module_id] {
-      m.import_records.iter().for_each(|i| {
-        // because dynamic import is already as entry, so here ignore it
-        if i.kind != ImportKind::DynamicImport {
-          self.mark_modules_entry_bit(i.resolved_module, index, modules_entry_bit);
-        }
-      });
-    }
+    module_to_bits[module_id].set_bit(entry_index);
+    let Module::Normal(module) = &self.graph.modules[module_id] else { return };
+    module.import_records.iter().for_each(|rec| {
+      // Module imported dynamically will be considered as a entry,
+      // so we don't need to include it in this chunk
+      if rec.kind != ImportKind::DynamicImport {
+        self.determine_reachable_modules_for_entry(
+          rec.resolved_module,
+          entry_index,
+          module_to_bits,
+        );
+      }
+    });
   }
 
-  pub fn generate_chunks(&self) -> (ChunksVec, IndexVec<ModuleId, Option<ChunkId>>) {
+  fn generate_chunks(&self) -> ChunkGraph {
+    let entries_len: u32 = self.graph.entries.len().try_into().unwrap();
+
     let mut module_to_bits = index_vec::index_vec![
-      BitSet::new(self.graph.entries.len().try_into().unwrap());
+      BitSet::new(entries_len);
       self.graph.modules.len()
     ];
-
-    let mut chunks = FxHashMap::default();
-    chunks.shrink_to(self.graph.entries.len());
-
-    for (i, (name, module_id)) in self.graph.entries.iter().enumerate() {
-      let count: u32 = u32::try_from(i).unwrap();
-      let mut entry_bits = BitSet::new(self.graph.entries.len().try_into().unwrap());
-      entry_bits.set_bit(count);
-      let c = Chunk::new(name.clone(), Some(*module_id), entry_bits.clone(), vec![]);
-      chunks.insert(entry_bits, c);
+    let mut bits_to_chunk =
+      FxHashMap::with_capacity_and_hasher(self.graph.entries.len(), BuildHasherDefault::default());
+    let mut chunks = ChunksVec::with_capacity(self.graph.entries.len());
+
+    // Create chunk for each static and dynamic entry
+    for (entry_index, (name, module_id)) in self.graph.entries.iter().enumerate() {
+      let count: u32 = u32::try_from(entry_index).unwrap();
+      let mut bits = BitSet::new(entries_len);
+      bits.set_bit(count);
+      let chunk = chunks.push(Chunk::new(name.clone(), Some(*module_id), bits.clone(), vec![]));
+      bits_to_chunk.insert(bits, chunk);
     }
 
+    // Determine which modules belong to which chunk. A module could belong to multiple chunks.
     self.graph.entries.iter().enumerate().for_each(|(i, (_, entry))| {
-      self.mark_modules_entry_bit(*entry, i, &mut module_to_bits);
+      self.determine_reachable_modules_for_entry(
+        *entry,
+        i.try_into().unwrap(),
+        &mut module_to_bits,
+      );
     });
 
-    self
-      .graph
-      .modules
-      .iter()
-      .enumerate()
-      // TODO avoid generate runtime module
-      .skip_while(|(module_id, _)| module_id.eq(&self.graph.runtime.id)) // TODO avoid generate runtime module
-      .for_each(|(_, module)| {
-        let bits = &module_to_bits[module.id()];
-        if let Some(chunk) = chunks.get_mut(bits) {
-          chunk.modules.push(module.id());
-        } else {
-          // TODO share chunk name
-          let len = chunks.len();
-          chunks.insert(
-            bits.clone(),
-            Chunk::new(Some(len.to_string()), None, bits.clone(), vec![module.id()]),
-          );
-        }
-      });
-
-    let chunks = chunks
-      .into_values()
-      .map(|mut chunk| {
-        chunk.modules.sort_by_key(|id| self.graph.modules[*id].exec_order());
-        chunk
-      })
-      .collect::<ChunksVec>();
-
     let mut module_to_chunk: IndexVec<ModuleId, Option<ChunkId>> = index_vec::index_vec![
       None;
       self.graph.modules.len()
     ];
 
-    // perf: this process could be done with computing chunks together
-    for (i, chunk) in chunks.iter_enumerated() {
-      for module_id in &chunk.modules {
-        module_to_chunk[*module_id] = Some(i);
+    // 1. Assign modules to corresponding chunks
+    // 2. Create shared chunks to store modules that belong to multiple chunks.
+    for module in &self.graph.modules {
+      if module.id() == self.graph.runtime.id {
+        // TODO: render runtime module
+        continue;
+      }
+      let bits = &module_to_bits[module.id()];
+      if let Some(chunk_id) = bits_to_chunk.get(bits).copied() {
+        chunks[chunk_id].modules.push(module.id());
+        module_to_chunk[module.id()] = Some(chunk_id);
+      } else {
+        let len = bits_to_chunk.len();
+        // FIXME: https://github.com/rolldown-rs/rolldown/issues/49
+        let chunk = Chunk::new(Some(len.to_string()), None, bits.clone(), vec![module.id()]);
+        let chunk_id = chunks.push(chunk);
+        module_to_chunk[module.id()] = Some(chunk_id);
+        bits_to_chunk.insert(bits.clone(), chunk_id);
       }
     }
 
-    (chunks, module_to_chunk)
+    // Sort modules in each chunk by execution order
+    chunks.iter_mut().for_each(|chunk| {
+      chunk.modules.sort_by_key(|module_id| self.graph.modules[*module_id].exec_order());
+    });
+
+    ChunkGraph { chunks, module_to_chunk }
   }
 
   pub fn generate(
     &mut self,
     _input_options: &'a NormalizedInputOptions,
   ) -> anyhow::Result<Vec<Asset>> {
     use rayon::prelude::*;
-    let (mut chunks, module_to_chunk) = self.generate_chunks();
+    let mut chunk_graph = self.generate_chunks();
 
-    chunks.iter_mut().par_bridge().for_each(|chunk| chunk.render_file_name(self.output_options));
+    chunk_graph
+      .chunks
+      .iter_mut()
+      .par_bridge()
+      .for_each(|chunk| chunk.render_file_name(self.output_options));
 
-    chunks.iter_mut().par_bridge().for_each(|chunk| {
+    chunk_graph.chunks.iter_mut().par_bridge().for_each(|chunk| {
       chunk.de_conflict(self.graph);
     });
 
-    chunks.iter_mut().for_each(|chunk| {
+    chunk_graph.chunks.iter_mut().for_each(|chunk| {
       if chunk.entry_module.is_some() {
         chunk.initialize_exports(&self.graph.linker_modules, &self.graph.symbols);
       }
     });
 
-    let assets = chunks
+    let assets = chunk_graph
+      .chunks
       .iter()
       .enumerate()
       .map(|(_chunk_id, c)| {
-        let content = c.render(self.graph, &module_to_chunk, &chunks).unwrap();
+        let content = c.render(self.graph, &chunk_graph).unwrap();
 
         Asset { file_name: c.file_name.clone().unwrap(), content }
       })

crates/rolldown/src/bundler/chunk/chunk.rs

@@ -1,4 +1,3 @@
-use index_vec::IndexVec;
 use oxc::span::Atom;
 use rolldown_common::{ModuleId, ResolvedExport, SymbolRef};
 use rustc_hash::FxHashMap;
@@ -17,7 +16,7 @@ use crate::bundler::{
   },
 };
 
-use super::{ChunkId, ChunksVec};
+use super::{chunk_graph::ChunkGraph, ChunkId};
 
 #[derive(Debug, Default)]
 pub struct Chunk {
@@ -94,12 +93,7 @@ impl Chunk {
   }
 
   #[allow(clippy::unnecessary_wraps)]
-  pub fn render(
-    &self,
-    graph: &Graph,
-    module_to_chunk: &IndexVec<ModuleId, Option<ChunkId>>,
-    chunks: &ChunksVec,
-  ) -> anyhow::Result<String> {
+  pub fn render(&self, graph: &Graph, chunk_graph: &ChunkGraph) -> anyhow::Result<String> {
     use rayon::prelude::*;
     let mut joiner = Joiner::with_options(JoinerOptions { separator: Some("\n".to_string()) });
     self
@@ -108,12 +102,7 @@ impl Chunk {
       .copied()
       .map(|id| &graph.modules[id])
       .filter_map(|m| {
-        m.render(ModuleRenderContext {
-          canonical_names: &self.canonical_names,
-          graph,
-          module_to_chunk,
-          chunks,
-        })
+        m.render(ModuleRenderContext { canonical_names: &self.canonical_names, graph, chunk_graph })
       })
       .collect::<Vec<_>>()
       .into_iter()

crates/rolldown/src/bundler/chunk/chunk_graph.rs

@@ -0,0 +1,10 @@
+use index_vec::IndexVec;
+use rolldown_common::ModuleId;
+
+use super::{ChunkId, ChunksVec};
+
+#[derive(Debug)]
+pub struct ChunkGraph {
+  pub chunks: ChunksVec,
+  pub module_to_chunk: IndexVec<ModuleId, Option<ChunkId>>,
+}

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

@@ -1,5 +1,6 @@
 #[allow(clippy::module_inception)]
 pub mod chunk;
+pub mod chunk_graph;
 mod de_conflict;
 
 use index_vec::IndexVec;

crates/rolldown/src/bundler/module/module.rs

@@ -5,7 +5,7 @@ use rustc_hash::FxHashMap;
 use string_wizard::MagicString;
 
 use crate::bundler::{
-  chunk::{ChunkId, ChunksVec},
+  chunk::chunk_graph::ChunkGraph,
   graph::{graph::Graph, linker::LinkerModule},
 };
 
@@ -78,6 +78,5 @@ impl Module {
 pub struct ModuleRenderContext<'a> {
   pub canonical_names: &'a FxHashMap<SymbolRef, Atom>,
   pub graph: &'a Graph,
-  pub module_to_chunk: &'a IndexVec<ModuleId, Option<ChunkId>>,
-  pub chunks: &'a ChunksVec,
+  pub chunk_graph: &'a ChunkGraph,
 }

crates/rolldown/src/bundler/module/normal_module.rs

@@ -78,8 +78,7 @@ impl NormalModule {
       ctx.graph,
       ctx.canonical_names,
       &mut source,
-      ctx.module_to_chunk,
-      ctx.chunks,
+      ctx.chunk_graph,
       self,
       self_linker_module,
     );

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

@@ -2,7 +2,6 @@ pub mod commonjs_source_render;
 pub mod esm_source_render;
 pub mod esm_wrap_source_render;
 pub mod scanner;
-use index_vec::IndexVec;
 use oxc::{
   semantic::ReferenceId,
   span::{Atom, GetSpan, Span},
@@ -12,7 +11,7 @@ use rustc_hash::FxHashMap;
 use string_wizard::{MagicString, UpdateOptions};
 
 use super::{
-  chunk::{chunk::Chunk, ChunkId},
+  chunk::chunk_graph::ChunkGraph,
   graph::{
     graph::Graph,
     linker::LinkerModule,
@@ -25,8 +24,7 @@ pub struct RendererContext<'ast> {
   graph: &'ast Graph,
   final_names: &'ast FxHashMap<SymbolRef, Atom>,
   source: &'ast mut MagicString<'static>,
-  module_to_chunk: &'ast IndexVec<ModuleId, Option<ChunkId>>,
-  chunks: &'ast IndexVec<ChunkId, Chunk>,
+  chunk_graph: &'ast ChunkGraph,
   module: &'ast NormalModule,
   linker_module: &'ast LinkerModule,
   wrap_symbol_name: Option<&'ast Atom>,
@@ -42,8 +40,7 @@ impl<'ast> RendererContext<'ast> {
     graph: &'ast Graph,
     final_names: &'ast FxHashMap<SymbolRef, Atom>,
     source: &'ast mut MagicString<'static>,
-    module_to_chunk: &'ast IndexVec<ModuleId, Option<ChunkId>>,
-    chunks: &'ast IndexVec<ChunkId, Chunk>,
+    chunk_graph: &'ast ChunkGraph,
     module: &'ast NormalModule,
     linker_module: &'ast LinkerModule,
   ) -> Self {
@@ -61,8 +58,7 @@ impl<'ast> RendererContext<'ast> {
       graph,
       final_names,
       source,
-      module_to_chunk,
-      chunks,
+      chunk_graph,
       module,
       linker_module,
       wrap_symbol_name,
@@ -233,8 +229,8 @@ impl<'ast> RendererContext<'ast> {
     if let oxc::ast::ast::Expression::StringLiteral(str) = &expr.source {
       let rec = &self.module.import_records[self.module.imports.get(&expr.span).copied().unwrap()];
 
-      if let Some(chunk_id) = self.module_to_chunk[rec.resolved_module] {
-        let chunk = &self.chunks[chunk_id];
+      if let Some(chunk_id) = self.chunk_graph.module_to_chunk[rec.resolved_module] {
+        let chunk = &self.chunk_graph.chunks[chunk_id];
         self.overwrite(
           str.span.start,
           str.span.end,