perf: more parallel real content hash plugin

Author: LingyuCoder

crates/rspack_plugin_real_content_hash/src/lib.rs

@@ -131,49 +131,99 @@ async fn inner_impl(compilation: &mut Compilation) -> Result<()> {
     })
     .collect();
 
-  let ordered_hashes = OrderedHashesBuilder::new(&hash_to_asset_names, &assets_data).build();
+  let (ordered_hashes, mut hash_dependencies) =
+    OrderedHashesBuilder::new(&hash_to_asset_names, &assets_data).build();
+  let mut ordered_hashes_iter = ordered_hashes.iter();
+
   logger.time_end(start);
 
   let start = logger.time("old hash to new hash");
   let mut hash_to_new_hash = HashMap::default();
 
   let hooks = RealContentHashPlugin::get_compilation_hooks(compilation.id());
-  for old_hash in &ordered_hashes {
-    if let Some(asset_names) = hash_to_asset_names.get_mut(old_hash.as_str()) {
-      asset_names.sort();
-      let mut asset_contents: Vec<_> = asset_names
-        .par_iter()
-        .filter_map(|name| assets_data.get(name))
-        .map(|data| {
-          data
-            .compute_new_source(
-              data.own_hashes.contains(old_hash),
-              &hash_to_new_hash,
-              &hash_regexp,
-            )
-            .clone()
-        })
-        .collect();
-      asset_contents.dedup();
-      let updated_hash = hooks
-        .update_hash
-        .call(compilation, &asset_contents, old_hash)
-        .await?;
-
-      let new_hash = if let Some(new_hash) = updated_hash {
-        new_hash
-      } else {
-        let mut hasher = RspackHash::from(&compilation.options.output);
-        for asset_content in asset_contents {
-          hasher.write(&asset_content.buffer());
-        }
-        let new_hash = hasher.digest(&compilation.options.output.hash_digest);
-        let new_hash = new_hash.rendered(old_hash.len()).to_string();
-        new_hash
+
+  let mut computed_hashes = HashSet::default();
+  let mut top_task = ordered_hashes_iter.next();
+
+  loop {
+    let Some(top) = top_task else {
+      break;
+    };
+    let mut batch = vec![top];
+    top_task = None;
+
+    for hash in ordered_hashes_iter.by_ref() {
+      let Some(dependencies) = hash_dependencies.remove(hash.as_str()) else {
+        top_task = Some(hash);
+        break;
       };
-      hash_to_new_hash.insert(old_hash, new_hash);
+      if dependencies.iter().all(|dep| computed_hashes.contains(dep)) {
+        batch.push(hash);
+      } else {
+        top_task = Some(hash);
+        break;
+      }
+    }
+
+    let batch_source_tasks = batch
+      .iter()
+      .filter_map(|hash| {
+        let assets_names = hash_to_asset_names.get(hash.as_str())?;
+        let tasks = assets_names
+          .iter()
+          .filter_map(|name| {
+            let data = assets_data.get(name)?;
+            Some((hash.as_str(), *name, data))
+          })
+          .collect::<Vec<_>>();
+        Some(tasks)
+      })
+      .flatten()
+      .collect::<Vec<_>>();
+
+    let mut batch_sources = batch_source_tasks
+      .into_par_iter()
+      .map(|(hash, name, data)| {
+        let new_source = data.compute_new_source(
+          data.own_hashes.contains(hash),
+          &hash_to_new_hash,
+          &hash_regexp,
+        );
+        ((hash, name), new_source)
+      })
+      .collect::<HashMap<_, _>>();
+
+    for old_hash in batch.iter() {
+      if let Some(asset_names) = hash_to_asset_names.get_mut(old_hash.as_str()) {
+        asset_names.sort();
+        let mut asset_contents = asset_names
+          .iter()
+          .filter_map(|name| batch_sources.remove(&(old_hash.as_str(), name)))
+          .collect::<Vec<_>>();
+        asset_contents.dedup();
+        let updated_hash = hooks
+          .update_hash
+          .call(compilation, &asset_contents, old_hash)
+          .await?;
+
+        let new_hash = if let Some(new_hash) = updated_hash {
+          new_hash
+        } else {
+          let mut hasher = RspackHash::from(&compilation.options.output);
+          for asset_content in asset_contents {
+            hasher.write(&asset_content.buffer());
+          }
+          let new_hash = hasher.digest(&compilation.options.output.hash_digest);
+          let new_hash = new_hash.rendered(old_hash.len()).to_string();
+          new_hash
+        };
+        hash_to_new_hash.insert(old_hash, new_hash);
+      }
     }
+
+    computed_hashes.extend(batch);
   }
+
   logger.time_end(start);
 
   let start = logger.time("collect hash updates");
@@ -193,7 +243,7 @@ async fn inner_impl(compilation: &mut Compilation) -> Result<()> {
         })
         .into_owned();
       let new_name = (name != new_name).then_some(new_name);
-      Some((name.to_owned(), new_source.clone(), new_name))
+      Some((name.to_owned(), new_source, new_name))
     })
     .collect();
   logger.time_end(start);
@@ -279,7 +329,7 @@ impl AssetData {
     without_own: bool,
     hash_to_new_hash: &HashMap<&str, String>,
     hash_regexp: &Regex,
-  ) -> &BoxSource {
+  ) -> BoxSource {
     (if without_own {
       &self.new_source_without_own
     } else {
@@ -309,6 +359,7 @@ impl AssetData {
       }
       self.old_source.clone()
     })
+    .clone()
   }
 }
 
@@ -328,12 +379,29 @@ impl<'a> OrderedHashesBuilder<'a> {
     }
   }
 
-  pub fn build(&self) -> IndexSet<String> {
+  pub fn build(&self) -> (IndexSet<String>, HashMap<String, HashSet<String>>) {
     let mut ordered_hashes = IndexSet::default();
+    let mut hash_dependencies = HashMap::default();
     for hash in self.hash_to_asset_names.keys() {
-      self.add_to_ordered_hashes(hash, &mut ordered_hashes, &mut HashSet::default());
+      self.add_to_ordered_hashes(
+        hash,
+        &mut ordered_hashes,
+        &mut HashSet::default(),
+        &mut hash_dependencies,
+      );
     }
-    ordered_hashes
+    (
+      ordered_hashes,
+      hash_dependencies
+        .into_iter()
+        .map(|(k, v)| {
+          (
+            k.to_string(),
+            v.into_iter().map(|s| s.to_string()).collect(),
+          )
+        })
+        .collect(),
+    )
   }
 }
 
@@ -364,8 +432,12 @@ impl OrderedHashesBuilder<'_> {
     hash: &'b str,
     ordered_hashes: &mut IndexSet<String>,
     stack: &mut HashSet<&'b str>,
+    hash_dependencies: &mut HashMap<&'b str, HashSet<&'b str>>,
   ) {
-    let deps = self.get_hash_dependencies(hash);
+    let deps = hash_dependencies
+      .entry(hash)
+      .or_insert_with(|| self.get_hash_dependencies(hash))
+      .clone();
     stack.insert(hash);
     for dep in deps {
       if ordered_hashes.contains(dep) {
@@ -376,7 +448,7 @@ impl OrderedHashesBuilder<'_> {
         // so there shouldn't have circular hash dependency between chunks
         panic!("RealContentHashPlugin: circular hash dependency");
       }
-      self.add_to_ordered_hashes(dep, ordered_hashes, stack);
+      self.add_to_ordered_hashes(dep, ordered_hashes, stack, hash_dependencies);
     }
     ordered_hashes.insert(hash.to_string());
     stack.remove(hash);