perf: introduce ModuleGraphCache and cache the result of get_mode and determine_export_assignments (#10584)

* Remove useless parameters

* Add ModuleGraphCacheArtifact

* Cache FlagDependencyExportsPlugin and seal

* Cache determine_export_assignments

* Reduce DetermineExportAssignmentsKey size

* Reduce GetModeCacheKey size

* Unfreeze in rebuild_module

* Freeze and unfreeze in finish

* Remove todo

* Add comments

* Split modules

* Update DetermineExportAssignmentsKey

* Use fxhashmap

* Fix rebase

* Fix clippy

* Remove module name

* Mutex -> RwLock

* Update memory orderings

Author: CPunisher

crates/rspack_core/src/artifacts/mod.rs

@@ -7,12 +7,14 @@ mod cgm_hash_artifact;
 mod cgm_runtime_requirement_artifact;
 mod chunk_hashes_artifact;
 mod code_generation_results;
+mod module_graph_cache_artifact;
 mod side_effects_do_optimize_artifact;
 
 pub use cgm_hash_artifact::*;
 pub use cgm_runtime_requirement_artifact::*;
 pub use chunk_hashes_artifact::*;
 pub use code_generation_results::*;
+pub use module_graph_cache_artifact::*;
 pub use side_effects_do_optimize_artifact::*;
 
 pub type AsyncModulesArtifact = IdentifierSet;

crates/rspack_core/src/artifacts/module_graph_cache_artifact.rs

@@ -0,0 +1,206 @@
+use std::sync::{
+  atomic::{AtomicBool, Ordering},
+  Arc, RwLock,
+};
+
+pub use determine_export_assignments::DetermineExportAssignmentsKey;
+use determine_export_assignments::*;
+use get_mode::*;
+use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
+use swc_core::atoms::Atom;
+
+use crate::{DependencyId, ExportInfo, RuntimeKey};
+pub type ModuleGraphCacheArtifact = Arc<ModuleGraphCacheArtifactInner>;
+
+/// This is a rust port of `ModuleGraph.cached` and `ModuleGraph.dependencyCacheProvide` in webpack.
+/// We use this to cache the result of functions with high computational overhead.
+#[derive(Debug, Default)]
+pub struct ModuleGraphCacheArtifactInner {
+  /// Webpack enables module graph caches by creating new cache maps and disable them by setting them to undefined.
+  /// But in rust I think it's better to use a bool flag to avoid memory reallocation.
+  freezed: AtomicBool,
+  get_mode_cache: GetModeCache,
+  determine_export_assignments_cache: DetermineExportAssignmentsCache,
+}
+
+impl ModuleGraphCacheArtifactInner {
+  pub fn freeze(&self) {
+    self.get_mode_cache.freeze();
+    self.determine_export_assignments_cache.freeze();
+    self.freezed.store(true, Ordering::Release);
+  }
+
+  pub fn unfreeze(&self) {
+    self.freezed.store(false, Ordering::Release);
+  }
+
+  pub fn cached_get_mode<F: FnOnce() -> ExportMode>(
+    &self,
+    key: GetModeCacheKey,
+    f: F,
+  ) -> ExportMode {
+    if !self.freezed.load(Ordering::Acquire) {
+      return f();
+    }
+
+    match self.get_mode_cache.get(&key) {
+      Some(value) => value,
+      None => {
+        let value = f();
+        self.get_mode_cache.set(key, value.clone());
+        value
+      }
+    }
+  }
+
+  pub fn cached_determine_export_assignments<F: FnOnce() -> DetermineExportAssignmentsValue>(
+    &self,
+    key: DetermineExportAssignmentsKey,
+    f: F,
+  ) -> DetermineExportAssignmentsValue {
+    if !self.freezed.load(Ordering::Acquire) {
+      return f();
+    }
+
+    match self.determine_export_assignments_cache.get(&key) {
+      Some(value) => value,
+      None => {
+        let value = f();
+        self
+          .determine_export_assignments_cache
+          .set(key, value.clone());
+        value
+      }
+    }
+  }
+}
+
+pub(super) mod get_mode {
+  use super::*;
+
+  pub type GetModeCacheKey = (DependencyId, Option<RuntimeKey>);
+
+  #[derive(Debug, Default)]
+  pub struct GetModeCache {
+    cache: RwLock<HashMap<GetModeCacheKey, ExportMode>>,
+  }
+
+  impl GetModeCache {
+    pub fn freeze(&self) {
+      self.cache.write().expect("should get lock").clear();
+    }
+
+    pub fn get(&self, key: &GetModeCacheKey) -> Option<ExportMode> {
+      let inner = self.cache.read().expect("should get lock");
+      inner.get(key).cloned()
+    }
+
+    pub fn set(&self, key: GetModeCacheKey, value: ExportMode) {
+      self
+        .cache
+        .write()
+        .expect("should get lock")
+        .insert(key, value);
+    }
+  }
+}
+
+pub(super) mod determine_export_assignments {
+  use super::*;
+  use crate::ModuleIdentifier;
+
+  /// Webpack cache the result of `determineExportAssignments` with the keys of dependencies arraris of `allStarExports.dependencies` and `otherStarExports` + `this`(DependencyId).
+  /// See: https://github.com/webpack/webpack/blob/19ca74127f7668aaf60d59f4af8fcaee7924541a/lib/dependencies/HarmonyExportImportedSpecifierDependency.js#L645
+  ///
+  /// However, we can simplify the cache key since dependencies under the same parent module share `allStarExports` and copy their own `otherStarExports`.
+  #[derive(Debug, PartialEq, Eq, Hash)]
+  pub enum DetermineExportAssignmentsKey {
+    All(ModuleIdentifier),
+    Other(DependencyId),
+  }
+  pub type DetermineExportAssignmentsValue = (Vec<Atom>, Vec<usize>);
+
+  #[derive(Debug, Default)]
+  pub struct DetermineExportAssignmentsCache {
+    cache: RwLock<HashMap<DetermineExportAssignmentsKey, DetermineExportAssignmentsValue>>,
+  }
+
+  impl DetermineExportAssignmentsCache {
+    pub fn freeze(&self) {
+      self.cache.write().expect("should get lock").clear();
+    }
+
+    pub fn get(
+      &self,
+      key: &DetermineExportAssignmentsKey,
+    ) -> Option<DetermineExportAssignmentsValue> {
+      let inner = self.cache.read().expect("should get lock");
+      inner.get(key).cloned()
+    }
+
+    pub fn set(&self, key: DetermineExportAssignmentsKey, value: DetermineExportAssignmentsValue) {
+      self
+        .cache
+        .write()
+        .expect("should get lock")
+        .insert(key, value);
+    }
+  }
+}
+
+#[derive(Debug, Clone)]
+pub struct NormalReexportItem {
+  pub name: Atom,
+  pub ids: Vec<Atom>,
+  pub hidden: bool,
+  pub checked: bool,
+  pub export_info: ExportInfo,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum ExportModeType {
+  Missing,
+  Unused,
+  EmptyStar,
+  ReexportDynamicDefault,
+  ReexportNamedDefault,
+  ReexportNamespaceObject,
+  ReexportFakeNamespaceObject,
+  ReexportUndefined,
+  NormalReexport,
+  DynamicReexport,
+}
+
+#[derive(Debug, Clone)]
+pub struct ExportMode {
+  /// corresponding to `type` field in webpack's `EpxortMode`
+  pub ty: ExportModeType,
+  pub items: Option<Vec<NormalReexportItem>>,
+  pub name: Option<Atom>,
+  pub fake_type: u8,
+  pub partial_namespace_export_info: Option<ExportInfo>,
+  pub ignored: Option<HashSet<Atom>>,
+  pub hidden: Option<HashSet<Atom>>,
+}
+
+impl ExportMode {
+  pub fn new(ty: ExportModeType) -> Self {
+    Self {
+      ty,
+      items: None,
+      name: None,
+      fake_type: 0,
+      partial_namespace_export_info: None,
+      ignored: None,
+      hidden: None,
+    }
+  }
+}
+
+#[derive(Debug, Default)]
+pub struct StarReexportsInfo {
+  pub exports: Option<HashSet<Atom>>,
+  pub checked: Option<HashSet<Atom>>,
+  pub ignored_exports: HashSet<Atom>,
+  pub hidden: Option<HashSet<Atom>>,
+}

crates/rspack_core/src/build_chunk_graph/code_splitter.rs

@@ -24,7 +24,8 @@ use crate::{
   merge_runtime, AsyncDependenciesBlockIdentifier, ChunkGroup, ChunkGroupKind, ChunkGroupOptions,
   ChunkGroupUkey, ChunkLoading, ChunkUkey, Compilation, ConnectionState, DependenciesBlock,
   DependencyId, DependencyLocation, EntryDependency, EntryRuntime, GroupOptions, Logger,
-  ModuleDependency, ModuleGraph, ModuleIdentifier, RuntimeSpec, SyntheticDependencyLocation,
+  ModuleDependency, ModuleGraph, ModuleGraphCacheArtifact, ModuleIdentifier, RuntimeSpec,
+  SyntheticDependencyLocation,
 };
 
 type IndexMap<K, V, H = FxHasher> = RawIndexMap<K, V, BuildHasherDefault<H>>;
@@ -310,21 +311,22 @@ fn get_active_state_of_connections(
   connections: &[DependencyId],
   runtime: Option<&RuntimeSpec>,
   module_graph: &ModuleGraph,
+  module_graph_cache: &ModuleGraphCacheArtifact,
 ) -> ConnectionState {
   let mut iter = connections.iter();
   let id = iter.next().expect("should have connection");
   let mut merged = module_graph
     .connection_by_dependency_id(id)
     .expect("should have connection")
-    .active_state(module_graph, runtime);
+    .active_state(module_graph, runtime, module_graph_cache);
   if merged.is_true() {
     return merged;
   }
   for c in iter {
     let c = module_graph
       .connection_by_dependency_id(c)
       .expect("should have connection");
-    merged = merged + c.active_state(module_graph, runtime);
+    merged = merged + c.active_state(module_graph, runtime, module_graph_cache);
     if merged.is_true() {
       return merged;
     }
@@ -1648,8 +1650,12 @@ Or do you want to use the entrypoints '{name}' and '{runtime}' independently on
       let modules = map
         .get_mut(&block_id)
         .expect("should have modules in block_modules_runtime_map");
-      let active_state =
-        get_active_state_of_connections(&connections, runtime, &compilation.get_module_graph());
+      let active_state = get_active_state_of_connections(
+        &connections,
+        runtime,
+        &compilation.get_module_graph(),
+        &compilation.module_graph_cache_artifact,
+      );
       modules.push((module_identifier, active_state, connections));
     }
   }
@@ -1771,6 +1777,7 @@ Or do you want to use the entrypoints '{name}' and '{runtime}' independently on
             connections,
             Some(&cgi.runtime),
             &compilation.get_module_graph(),
+            &compilation.module_graph_cache_artifact,
           );
           if active_state.is_false() {
             continue;

crates/rspack_core/src/build_chunk_graph/new_code_splitter.rs

@@ -18,8 +18,8 @@ use crate::{
   merge_runtime, AsyncDependenciesBlockIdentifier, Chunk, ChunkGroup, ChunkGroupKind,
   ChunkGroupOptions, ChunkGroupUkey, ChunkLoading, ChunkUkey, Compilation, DependenciesBlock,
   DependencyLocation, EntryData, EntryDependency, EntryOptions, EntryRuntime, GroupOptions,
-  ModuleDependency, ModuleGraph, ModuleGraphConnection, ModuleIdentifier, RuntimeSpec,
-  SyntheticDependencyLocation,
+  ModuleDependency, ModuleGraph, ModuleGraphCacheArtifact, ModuleGraphConnection, ModuleIdentifier,
+  RuntimeSpec, SyntheticDependencyLocation,
 };
 
 type ModuleDeps = HashMap<
@@ -197,6 +197,7 @@ impl CreateChunkRoot {
 
   fn create(&self, splitter: &CodeSplitter, compilation: &Compilation) -> Vec<ChunkDesc> {
     let module_graph = compilation.get_module_graph();
+    let module_graph_cache = &compilation.module_graph_cache_artifact;
 
     match self {
       CreateChunkRoot::Entry(entry, data, runtime) => {
@@ -237,7 +238,7 @@ impl CreateChunkRoot {
           )
           .filter_map(|dep_id| module_graph.module_identifier_by_dependency_id(dep_id))
         {
-          splitter.fill_chunk_modules(*m, runtime, &module_graph, &mut ctx);
+          splitter.fill_chunk_modules(*m, runtime, &module_graph, module_graph_cache, &mut ctx);
         }
 
         vec![ChunkDesc::Entry(Box::new(EntryChunkDesc {
@@ -291,7 +292,7 @@ impl CreateChunkRoot {
             continue;
           };
 
-          splitter.fill_chunk_modules(*m, runtime, &module_graph, &mut ctx);
+          splitter.fill_chunk_modules(*m, runtime, &module_graph, module_graph_cache, &mut ctx);
         }
 
         if let Some(group_option) = block.get_group_options()
@@ -451,6 +452,7 @@ impl CodeSplitter {
     let mut stack = vec![];
     let mut visited = HashSet::default();
     let module_graph = compilation.get_module_graph();
+    let module_graph_cache = &compilation.module_graph_cache_artifact;
     let global_chunk_loading = &compilation.options.output.chunk_loading;
     let mut roots = HashMap::<AsyncDependenciesBlockIdentifier, CreateChunkRoot>::default();
     let mut named_roots = HashMap::<String, CreateChunkRoot>::default();
@@ -501,7 +503,8 @@ impl CodeSplitter {
         continue;
       }
 
-      let guard = self.outgoings_modules(&module, runtime.as_ref(), &module_graph);
+      let guard =
+        self.outgoings_modules(&module, runtime.as_ref(), &module_graph, module_graph_cache);
       let (modules, blocks) = guard.value();
       let blocks = blocks.clone();
       for m in modules {
@@ -749,6 +752,7 @@ impl CodeSplitter {
     module: &ModuleIdentifier,
     runtime: &RuntimeSpec,
     module_graph: &ModuleGraph,
+    module_graph_cache: &ModuleGraphCacheArtifact,
   ) -> Ref<ModuleIdentifier, (Vec<ModuleIdentifier>, Vec<AsyncDependenciesBlockIdentifier>)> {
     let module_map = self.module_deps.get(runtime).expect("should have value");
 
@@ -783,14 +787,14 @@ impl CodeSplitter {
 
     'outer: for (m, conns) in outgoings.iter() {
       for conn in conns {
-        let conn_state = conn.active_state(module_graph, Some(runtime));
+        let conn_state = conn.active_state(module_graph, Some(runtime), module_graph_cache);
         match conn_state {
           crate::ConnectionState::Active(true) => {
             modules.insert(*m);
             continue 'outer;
           }
           crate::ConnectionState::TransitiveOnly => {
-            let transitive = self.outgoings_modules(m, runtime, module_graph);
+            let transitive = self.outgoings_modules(m, runtime, module_graph, module_graph_cache);
             let (extra_modules, extra_blocks) = transitive.value();
             modules.extend(extra_modules.iter().copied());
             blocks.extend(extra_blocks.iter().copied());
@@ -811,6 +815,7 @@ impl CodeSplitter {
     target_module: ModuleIdentifier,
     runtime: &RuntimeSpec,
     module_graph: &ModuleGraph,
+    module_graph_cache: &ModuleGraphCacheArtifact,
     ctx: &mut FillCtx,
   ) {
     enum Task {
@@ -855,7 +860,8 @@ impl CodeSplitter {
               value
             });
 
-          let guard = self.outgoings_modules(&target_module, runtime, module_graph);
+          let guard =
+            self.outgoings_modules(&target_module, runtime, module_graph, module_graph_cache);
           let (_, (outgoing_modules, blocks)) = guard.pair();
           let mut outgoing_modules = outgoing_modules.clone();
 
@@ -940,6 +946,7 @@ impl CodeSplitter {
 
     let mut visited = HashSet::default();
     let module_graph = compilation.get_module_graph();
+    let module_graph_cache = &compilation.module_graph_cache_artifact;
 
     queue.reverse();
 
@@ -983,7 +990,7 @@ impl CodeSplitter {
             if !self.module_deps.contains_key(runtime) {
               self.module_deps.insert(runtime.clone(), Default::default());
             }
-            let guard = self.outgoings_modules(&m, runtime, &module_graph);
+            let guard = self.outgoings_modules(&m, runtime, &module_graph, module_graph_cache);
             let (modules, blocks) = guard.value();
 
             for m in modules.iter().rev() {