feat(rust): support advanced_chunks option

hyf0 · hyf0 · commit 3e9b18a6a0ab · 2024-09-01T19:18:21.000Z
diff --git a/crates/rolldown/src/chunk_graph.rs b/crates/rolldown/src/chunk_graph.rs
@@ -1,4 +1,3 @@
-use arcstr::ArcStr;
 use oxc::index::{index_vec, IndexVec};
 use rolldown_common::{Chunk, ChunkIdx, ChunkTable, ModuleIdx, ModuleTable};
 use rustc_hash::FxHashMap;
@@ -10,7 +9,6 @@ pub struct ChunkGraph {
   /// Module to chunk that contains the module
   pub module_to_chunk: IndexVec<ModuleIdx, Option<ChunkIdx>>,
   pub entry_module_to_entry_chunk: FxHashMap<ModuleIdx, ChunkIdx>,
-  chunk_idx_by_name: FxHashMap<ArcStr, ChunkIdx>,
 }
 
 impl ChunkGraph {
@@ -20,7 +18,6 @@ impl ChunkGraph {
       module_to_chunk: index_vec![None; module_table.modules.len()],
       sorted_chunk_idx_vec: Vec::new(),
       entry_module_to_entry_chunk: FxHashMap::default(),
-      chunk_idx_by_name: FxHashMap::default(),
     }
   }
 
@@ -30,16 +27,7 @@ impl ChunkGraph {
   }
 
   pub fn add_chunk(&mut self, chunk: Chunk) -> ChunkIdx {
-    let idx = self.chunk_table.push(chunk);
-    let chunk = &self.chunk_table.chunks[idx];
-    if let Some(name) = &chunk.name {
-      debug_assert!(
-        !self.chunk_idx_by_name.contains_key(name),
-        "Should not have duplicate chunk name"
-      );
-      self.chunk_idx_by_name.insert(name.clone(), idx);
-    }
-    idx
+    self.chunk_table.push(chunk)
   }
 
   pub fn add_module_to_chunk(&mut self, module_idx: ModuleIdx, chunk_idx: ChunkIdx) {
diff --git a/crates/rolldown/src/stages/generate_stage/code_splitting.rs b/crates/rolldown/src/stages/generate_stage/code_splitting.rs
@@ -1,6 +1,7 @@
 use std::cmp::Ordering;
 
 use crate::chunk_graph::ChunkGraph;
+use arcstr::ArcStr;
 use itertools::Itertools;
 use oxc::index::IndexVec;
 use rolldown_common::{Chunk, ChunkIdx, ChunkKind, Module, ModuleIdx, OutputFormat};
@@ -12,7 +13,7 @@ use super::GenerateStage;
 
 impl<'a> GenerateStage<'a> {
   #[tracing::instrument(level = "debug", skip_all)]
-  pub fn generate_chunks(&mut self) -> ChunkGraph {
+  pub async fn generate_chunks(&mut self) -> anyhow::Result<ChunkGraph> {
     if matches!(self.options.format, OutputFormat::Iife) {
       let user_defined_entry_count =
         self.link_output.entries.iter().filter(|entry| entry.kind.is_user_defined()).count();
@@ -63,18 +64,47 @@ impl<'a> GenerateStage<'a> {
       );
     });
 
+    let mut module_to_assigned: IndexVec<ModuleIdx, bool> =
+      oxc::index::index_vec![false; self.link_output.module_table.modules.len()];
+
     // 1. Assign modules to corresponding chunks
     // 2. Create shared chunks to store modules that belong to multiple chunks.
     for normal_module in self.link_output.module_table.modules.iter().filter_map(Module::as_ecma) {
       if !normal_module.is_included {
         continue;
       }
 
+      if module_to_assigned[normal_module.idx] {
+        continue;
+      }
+      module_to_assigned[normal_module.idx] = true;
+
       let bits = &module_to_bits[normal_module.idx];
       debug_assert!(
         !bits.is_empty(),
         "Empty bits means the module is not reachable, so it should bail out with `is_included: false` {:?}", normal_module.stable_id
       );
+
+      let mut manual_chunk_by_name = FxHashMap::default();
+
+      if let Some(manual_chunks) = &self.options.advanced_chunks {
+        if let Some(matched) = (manual_chunks)(&normal_module.id).await? {
+          let chunk_name = ArcStr::from(matched.name);
+          if let Some(chunk_idx) = manual_chunk_by_name.get(&chunk_name).copied() {
+            let chunk: &mut Chunk = &mut chunk_graph.chunk_table[chunk_idx];
+            chunk.bits.union(bits);
+            chunk_graph.add_module_to_chunk(normal_module.idx, chunk_idx);
+          } else {
+            let chunk =
+              Chunk::new(Some(chunk_name.clone()), bits.clone(), vec![], ChunkKind::Common);
+            let chunk_id = chunk_graph.add_chunk(chunk);
+            chunk_graph.add_module_to_chunk(normal_module.idx, chunk_id);
+            manual_chunk_by_name.insert(chunk_name, chunk_id);
+          }
+          continue;
+        }
+      }
+
       if let Some(chunk_id) = bits_to_chunk.get(bits).copied() {
         chunk_graph.add_module_to_chunk(normal_module.idx, chunk_id);
       } else {
@@ -196,7 +226,7 @@ impl<'a> GenerateStage<'a> {
     chunk_graph.sorted_chunk_idx_vec = sorted_chunk_idx_vec;
     chunk_graph.entry_module_to_entry_chunk = entry_module_to_entry_chunk;
 
-    chunk_graph
+    Ok(chunk_graph)
   }
 
   fn determine_reachable_modules_for_entry(
diff --git a/crates/rolldown/src/stages/generate_stage/mod.rs b/crates/rolldown/src/stages/generate_stage/mod.rs
@@ -55,7 +55,7 @@ impl<'a> GenerateStage<'a> {
 
   #[tracing::instrument(level = "debug", skip_all)]
   pub async fn generate(&mut self) -> Result<BundleOutput> {
-    let mut chunk_graph = self.generate_chunks();
+    let mut chunk_graph = self.generate_chunks().await?;
 
     self.generate_chunk_name_and_preliminary_filenames(&mut chunk_graph).await?;
 
diff --git a/crates/rolldown/src/utils/normalize_options.rs b/crates/rolldown/src/utils/normalize_options.rs
@@ -120,6 +120,7 @@ pub fn normalize_options(mut raw_options: crate::BundlerOptions) -> NormalizeOpt
     extend: raw_options.extend.unwrap_or(false),
     external_live_bindings: raw_options.external_live_bindings.unwrap_or(true),
     inline_dynamic_imports: raw_options.inline_dynamic_imports.unwrap_or(false),
+    advanced_chunks: raw_options.advanced_chunks,
   };
 
   NormalizeOptionsReturn { options: normalized, resolve_options: raw_resolve }
diff --git a/crates/rolldown/tests/rolldown/function/advanced_chunks/basic/_config.json b/crates/rolldown/tests/rolldown/function/advanced_chunks/basic/_config.json
@@ -0,0 +1,20 @@
+{
+  "config": {
+    "input": [
+      {
+        "name": "a",
+        "import": "./a.js"
+      },
+      {
+        "name": "b",
+        "import": "./b.js"
+      }
+    ],
+    "advancedChunks": [
+      {
+        "filter": ".js",
+        "name": "common"
+      }
+    ]
+  }
+}
diff --git a/crates/rolldown/tests/rolldown/function/advanced_chunks/basic/a.js b/crates/rolldown/tests/rolldown/function/advanced_chunks/basic/a.js
@@ -0,0 +1 @@
+export const a = 'a'
diff --git a/crates/rolldown/tests/rolldown/function/advanced_chunks/basic/artifacts.snap b/crates/rolldown/tests/rolldown/function/advanced_chunks/basic/artifacts.snap
@@ -0,0 +1,39 @@
+---
+source: crates/rolldown_testing/src/integration_test.rs
+---
+# Assets
+
+## a.mjs
+
+```js
+import { a } from "./a~1.mjs";
+
+export { a };
+```
+## a~1.mjs
+
+```js
+
+//#region a.js
+const a = "a";
+
+//#endregion
+export { a };
+```
+## b.mjs
+
+```js
+import { b } from "./b~1.mjs";
+
+export { b };
+```
+## b~1.mjs
+
+```js
+
+//#region b.js
+const b = "a";
+
+//#endregion
+export { b };
+```
diff --git a/crates/rolldown/tests/rolldown/function/advanced_chunks/basic/b.js b/crates/rolldown/tests/rolldown/function/advanced_chunks/basic/b.js
@@ -0,0 +1 @@
+export const b = 'a'
diff --git a/crates/rolldown/tests/snapshots/integration_rolldown__filename_with_hash.snap b/crates/rolldown/tests/snapshots/integration_rolldown__filename_with_hash.snap
@@ -1454,6 +1454,13 @@ expression: "snapshot_outputs.join(\"\\n\")"
 # tests/rolldown/errors/unresolved_entry
 
 
+# tests/rolldown/function/advanced_chunks/basic
+
+- a-!~{000}~.mjs => a-SUCnoi0A.mjs
+- b-!~{001}~.mjs => b-SonL6YaA.mjs
+- a~1-!~{002}~.mjs => a~1-tSH6rdht.mjs
+- b~1-!~{004}~.mjs => b~1-xejoaD21.mjs
+
 # tests/rolldown/function/define/node_env
 
 - main-!~{000}~.mjs => main-zmFhpGWu.mjs
diff --git a/crates/rolldown_binding/src/utils/normalize_binding_options.rs b/crates/rolldown_binding/src/utils/normalize_binding_options.rs
@@ -181,6 +181,7 @@ pub fn normalize_binding_options(
       .map(|inner| inner.into_iter().map(normalize_binding_inject_import).collect()),
     external_live_bindings: output_options.external_live_bindings,
     inline_dynamic_imports: output_options.inline_dynamic_imports,
+    advanced_chunks: None,
   };
 
   #[cfg(not(target_family = "wasm"))]
diff --git a/crates/rolldown_common/src/inner_bundler_options/mod.rs b/crates/rolldown_common/src/inner_bundler_options/mod.rs
@@ -1,6 +1,7 @@
 use rolldown_utils::indexmap::FxIndexMap;
 use std::{collections::HashMap, fmt::Debug, path::PathBuf};
 use types::inject_import::InjectImport;
+use types::manual_chunks_option::ManualChunksOption;
 
 #[cfg(feature = "deserialize_bundler_options")]
 use schemars::JsonSchema;
@@ -13,8 +14,9 @@ use types::experimental_options::ExperimentalOptions;
 use self::types::treeshake::TreeshakeOptions;
 use self::types::{
   es_module_flag::EsModuleFlag, input_item::InputItem, is_external::IsExternal,
-  output_exports::OutputExports, output_format::OutputFormat, output_option::AddonOutputOption,
-  platform::Platform, resolve_options::ResolveOptions, source_map_type::SourceMapType,
+  manual_chunks_option::ManualChunksDescriptor, output_exports::OutputExports,
+  output_format::OutputFormat, output_option::AddonOutputOption, platform::Platform,
+  resolve_options::ResolveOptions, source_map_type::SourceMapType,
   sourcemap_path_transform::SourceMapPathTransform,
 };
 use crate::{ChunkFilenamesOutputOption, ModuleType, SourceMapIgnoreList};
@@ -129,6 +131,12 @@ pub struct BundlerOptions {
   pub inject: Option<Vec<InjectImport>>,
   pub external_live_bindings: Option<bool>,
   pub inline_dynamic_imports: Option<bool>,
+  #[cfg_attr(
+    feature = "deserialize_bundler_options",
+    serde(default, deserialize_with = "deserialize_advanced_chunks"),
+    schemars(with = "Option<Vec<ManualChunksDescriptor>>")
+  )]
+  pub advanced_chunks: Option<ManualChunksOption>,
 }
 
 #[cfg(feature = "deserialize_bundler_options")]
@@ -140,6 +148,17 @@ where
   Ok(deserialized.map(IsExternal::from_vec))
 }
 
+#[cfg(feature = "deserialize_bundler_options")]
+fn deserialize_advanced_chunks<'de, D>(
+  deserializer: D,
+) -> Result<Option<ManualChunksOption>, D::Error>
+where
+  D: Deserializer<'de>,
+{
+  let deserialized = Option::<Vec<ManualChunksDescriptor>>::deserialize(deserializer)?;
+  Ok(deserialized.map(ManualChunksOption::from_vec))
+}
+
 #[cfg(feature = "deserialize_bundler_options")]
 fn deserialize_addon<'de, D>(deserializer: D) -> Result<Option<AddonOutputOption>, D::Error>
 where
diff --git a/crates/rolldown_common/src/inner_bundler_options/types/manual_chunks_option.rs b/crates/rolldown_common/src/inner_bundler_options/types/manual_chunks_option.rs
@@ -0,0 +1,87 @@
+#[cfg(feature = "deserialize_bundler_options")]
+use schemars::JsonSchema;
+#[cfg(feature = "deserialize_bundler_options")]
+use serde::Deserialize;
+
+use std::{
+  fmt::Debug,
+  future::{self, Future},
+  ops::Deref,
+  pin::Pin,
+};
+
+type Inner = dyn Fn(
+    &str, // module id
+  )
+    -> Pin<Box<(dyn Future<Output = anyhow::Result<Option<ManualChunksMatch>>> + Send + 'static)>>
+  + Send
+  + Sync
+  + 'static;
+
+pub struct ManualChunksMatch {
+  // Chunk name
+  pub name: String,
+  pub include_dependencies_recursively: bool,
+}
+
+#[cfg_attr(
+  feature = "deserialize_bundler_options",
+  derive(Deserialize, JsonSchema),
+  serde(rename_all = "camelCase", deny_unknown_fields)
+)]
+pub struct ManualChunksDescriptor {
+  pub filter: String,
+  pub name: String,
+  pub include_dependencies_recursively: Option<bool>,
+  pub priority: Option<u32>,
+}
+
+pub struct ManualChunksOption(Box<Inner>);
+
+impl Debug for ManualChunksOption {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    f.debug_struct("ManualChunksOption").finish()
+  }
+}
+
+impl Deref for ManualChunksOption {
+  type Target = Inner;
+
+  fn deref(&self) -> &Self::Target {
+    &*self.0
+  }
+}
+
+impl ManualChunksOption {
+  pub fn from_vec(mut vec: Vec<ManualChunksDescriptor>) -> Self {
+    // Notice that we rely on stable sort here
+    vec.sort_by_key(|x| x.priority.unwrap_or(0));
+    Self::from_closure(move |module_id| {
+      Box::pin(future::ready(Ok(vec.iter().find_map(|desc| {
+        if module_id.contains(&desc.filter) {
+          Some(ManualChunksMatch {
+            name: desc.name.clone(),
+            include_dependencies_recursively: desc
+              .include_dependencies_recursively
+              .unwrap_or(false),
+          })
+        } else {
+          None
+        }
+      }))))
+    })
+  }
+
+  pub fn from_closure<F>(f: F) -> Self
+  where
+    F: Fn(
+        &str, // module id
+      ) -> Pin<
+        Box<(dyn Future<Output = anyhow::Result<Option<ManualChunksMatch>>> + Send + 'static)>,
+      > + Send
+      + Sync
+      + 'static,
+  {
+    Self(Box::new(f))
+  }
+}
diff --git a/crates/rolldown_common/src/inner_bundler_options/types/mod.rs b/crates/rolldown_common/src/inner_bundler_options/types/mod.rs
@@ -4,6 +4,7 @@ pub mod filename_template;
 pub mod inject_import;
 pub mod input_item;
 pub mod is_external;
+pub mod manual_chunks_option;
 pub mod module_type;
 pub mod normalized_bundler_options;
 pub mod output_exports;
diff --git a/crates/rolldown_common/src/inner_bundler_options/types/normalized_bundler_options.rs b/crates/rolldown_common/src/inner_bundler_options/types/normalized_bundler_options.rs
@@ -7,6 +7,7 @@ use oxc::minifier::InjectGlobalVariablesConfig;
 use rustc_hash::FxHashMap;
 
 use super::experimental_options::ExperimentalOptions;
+use super::manual_chunks_option::ManualChunksOption;
 use super::output_option::ChunkFilenamesOutputOption;
 use super::treeshake::TreeshakeOptions;
 use super::{
@@ -57,4 +58,5 @@ pub struct NormalizedBundlerOptions {
   pub oxc_inject_global_variables_config: InjectGlobalVariablesConfig,
   pub external_live_bindings: bool,
   pub inline_dynamic_imports: bool,
+  pub advanced_chunks: Option<ManualChunksOption>,
 }
diff --git a/crates/rolldown_testing/_config.schema.json b/crates/rolldown_testing/_config.schema.json
diff --git a/crates/rolldown_utils/src/bitset.rs b/crates/rolldown_utils/src/bitset.rs
