refactor(compilation): process dependencies should happen after all d…

…ependencies get factorized

crates/rspack_core/src/compiler/compilation.rs

@@ -41,11 +41,12 @@ use crate::{
   CleanTask, CleanTaskResult, CodeGenerationResult, CodeGenerationResults, CompilationLogger,
   CompilationLogging, CompilerOptions, ContentHashArgs, ContextDependency, DependencyId,
   DependencyParents, DependencyType, Entry, EntryData, EntryOptions, Entrypoint, ErrorSpan,
-  FactorizeQueue, FactorizeTask, FactorizeTaskResult, Filename, Logger, Module, ModuleFactory,
-  ModuleGraph, ModuleIdentifier, ModuleProfile, PathData, ProcessAssetsArgs,
-  ProcessDependenciesQueue, ProcessDependenciesResult, ProcessDependenciesTask, RenderManifestArgs,
-  Resolve, ResolverFactory, RuntimeGlobals, RuntimeModule, RuntimeRequirementsInTreeArgs,
-  RuntimeSpec, SharedPluginDriver, SourceType, Stats, TaskResult, WorkerTask,
+  FactorizeQueue, FactorizeTask, FactorizeTaskResult, Filename, Logger, Module,
+  ModuleCreationCallback, ModuleFactory, ModuleGraph, ModuleIdentifier, ModuleProfile, PathData,
+  ProcessAssetsArgs, ProcessDependenciesQueue, ProcessDependenciesResult, ProcessDependenciesTask,
+  RenderManifestArgs, Resolve, ResolverFactory, RuntimeGlobals, RuntimeModule,
+  RuntimeRequirementsInTreeArgs, RuntimeSpec, SharedPluginDriver, SourceType, Stats, TaskResult,
+  WorkerTask,
 };
 use crate::{tree_shaking::visitor::OptimizeAnalyzeResult, Context};
 
@@ -525,6 +526,7 @@ impl Compilation {
           parent_module
             .and_then(|m| m.as_normal_module())
             .and_then(|module| module.name_for_condition()),
+          None,
         );
       });
 
@@ -638,7 +640,10 @@ impl Compilation {
           .module_by_identifier(original_module_identifier)
           .expect("Module expected");
 
+        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
+        let mut remaining = sorted_dependencies.len();
         for dependencies in sorted_dependencies.into_values() {
+          let tx = tx.clone();
           self.handle_module_creation(
             &mut factorize_queue,
             Some(module.identifier()),
@@ -650,16 +655,33 @@ impl Compilation {
             module
               .as_normal_module()
               .and_then(|module| module.name_for_condition()),
+            Some(Box::new(move |_| {
+              tx.send(())
+                .expect("Failed to send callback to process_dependencies");
+            })),
           );
         }
+        drop(tx);
 
-        result_tx
-          .send(Ok(TaskResult::ProcessDependencies(Box::new(
+        let tx = result_tx.clone();
+
+        tokio::spawn(async move {
+          loop {
+            if remaining == 0 {
+              break;
+            }
+
+            rx.recv().await;
+            remaining -= 1;
+          }
+
+          tx.send(Ok(TaskResult::ProcessDependencies(Box::new(
             ProcessDependenciesResult {
               module_identifier: task.original_module_identifier,
             },
           ))))
           .expect("Failed to send process dependencies result");
+        });
       }
       process_deps_time.end(start);
 
@@ -677,6 +699,8 @@ impl Compilation {
                 current_profile,
                 exports_info_related,
                 from_cache,
+                callback,
+                ..
               } = task_result;
 
               if let Some(counter) = &mut factorize_cache_counter {
@@ -727,6 +751,7 @@ impl Compilation {
                 dependencies,
                 is_entry,
                 current_profile,
+                callback,
               });
             }
             Ok(TaskResult::Add(box task_result)) => match task_result {
@@ -1047,6 +1072,7 @@ impl Compilation {
     resolve_options: Option<Box<Resolve>>,
     lazy_visit_modules: std::collections::HashSet<String>,
     issuer: Option<Box<str>>,
+    callback: Option<ModuleCreationCallback>,
   ) {
     let current_profile = self.options.profile.then(Box::<ModuleProfile>::default);
     let dependency = dependencies[0].get_dependency(&self.module_graph).clone();
@@ -1066,6 +1092,7 @@ impl Compilation {
       plugin_driver: self.plugin_driver.clone(),
       cache: self.cache.clone(),
       current_profile,
+      callback,
     });
   }
 

crates/rspack_core/src/compiler/queue.rs

@@ -8,7 +8,7 @@ use crate::{
   ModuleGraph, ModuleGraphModule, ModuleIdentifier, ModuleProfile, Resolve, ResolverFactory,
   SharedPluginDriver, WorkerQueue,
 };
-use crate::{DependencyId, ExportInfo, ExportsInfo, UsageState};
+use crate::{BoxModule, DependencyId, ExportInfo, ExportsInfo, UsageState};
 
 #[derive(Debug)]
 pub enum TaskResult {
@@ -39,6 +39,7 @@ pub struct FactorizeTask {
   pub plugin_driver: SharedPluginDriver,
   pub cache: Arc<Cache>,
   pub current_profile: Option<Box<ModuleProfile>>,
+  pub callback: Option<ModuleCreationCallback>,
 }
 
 /// a struct temporarily used creating ExportsInfo
@@ -48,7 +49,7 @@ pub struct ExportsInfoRelated {
   pub other_exports_info: ExportInfo,
   pub side_effects_info: ExportInfo,
 }
-#[derive(Debug)]
+
 pub struct FactorizeTaskResult {
   pub original_module_identifier: Option<ModuleIdentifier>,
   pub factory_result: ModuleFactoryResult,
@@ -59,6 +60,26 @@ pub struct FactorizeTaskResult {
   pub current_profile: Option<Box<ModuleProfile>>,
   pub exports_info_related: ExportsInfoRelated,
   pub from_cache: bool,
+  pub callback: Option<ModuleCreationCallback>,
+}
+
+impl std::fmt::Debug for FactorizeTaskResult {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    f.debug_struct("FactorizeTaskResult")
+      .field(
+        "original_module_identifier",
+        &self.original_module_identifier,
+      )
+      .field("factory_result", &self.factory_result)
+      .field("module_graph_module", &self.module_graph_module)
+      .field("dependencies", &self.dependencies)
+      .field("diagnostics", &self.diagnostics)
+      .field("is_entry", &self.is_entry)
+      .field("current_profile", &self.current_profile)
+      .field("exports_info_related", &self.exports_info_related)
+      .field("from_cache", &self.from_cache)
+      .finish()
+  }
 }
 
 #[async_trait::async_trait]
@@ -121,6 +142,7 @@ impl WorkerTask for FactorizeTask {
         other_exports_info,
         side_effects_info: side_effects_only_info,
       },
+      callback: self.callback,
     })))
   }
 }
@@ -134,6 +156,7 @@ pub struct AddTask {
   pub dependencies: Vec<DependencyId>,
   pub is_entry: bool,
   pub current_profile: Option<Box<ModuleProfile>>,
+  pub callback: Option<ModuleCreationCallback>,
 }
 
 #[derive(Debug)]
@@ -166,6 +189,10 @@ impl AddTask {
         module_identifier,
       )?;
 
+      if let Some(callback) = self.callback {
+        callback(&self.module);
+      }
+
       return Ok(TaskResult::Add(Box::new(AddTaskResult::ModuleReused {
         module: self.module,
       })));
@@ -192,6 +219,10 @@ impl AddTask {
       current_profile.mark_integration_end();
     }
 
+    if let Some(callback) = self.callback {
+      callback(&self.module);
+    }
+
     Ok(TaskResult::Add(Box::new(AddTaskResult::ModuleAdded {
       module: self.module,
       current_profile: self.current_profile,
@@ -366,3 +397,5 @@ impl CleanTask {
 }
 
 pub type CleanQueue = WorkerQueue<CleanTask>;
+
+pub type ModuleCreationCallback = Box<dyn FnOnce(&BoxModule) + Send>;