feat: Adding runtime param when executing codegen (#4263)

* chore: 🤖 init * chore: 🤖 add runtime * chore: 🤖 update runtime * chore: 🤖 refactor * chore: 🤖 pass test * chore: 🤖 add runtime param * chore: 🤖 fix compile issue * chore: 🤖 fix test case * chore: 🤖 recover binding.d.ts * chore: 🤖 remove unnecessary comment
web-infra-dev · Oct 10, 2023 · 8a2990b · 8a2990b
1 parent 05815ca
commit 8a2990b
Show file tree

Hide file tree

Showing 19 changed files with 200 additions and 75 deletions.
diff --git a/crates/rspack_core/src/cache/occasion/code_generate.rs b/crates/rspack_core/src/cache/occasion/code_generate.rs
@@ -2,8 +2,9 @@ use rspack_error::Result;
 use rspack_identifier::Identifier;
 
 use crate::{cache::storage, BoxModule, CodeGenerationResult, Compilation, NormalModuleSource};
+use crate::{RuntimeSpec, RuntimeSpecSet};
 
-type Storage = dyn storage::Storage<CodeGenerationResult>;
+type Storage = dyn storage::Storage<Vec<(CodeGenerationResult, RuntimeSpec)>>;
 
 #[derive(Debug)]
 pub struct CodeGenerateOccasion {
@@ -19,16 +20,17 @@ impl CodeGenerateOccasion {
   pub fn use_cache<'a, G>(
     &self,
     module: &'a BoxModule,
+    runtimes: RuntimeSpecSet,
     compilation: &Compilation,
     generator: G,
-  ) -> Result<(CodeGenerationResult, bool)>
+  ) -> Result<(Vec<(CodeGenerationResult, RuntimeSpec)>, bool)>
   where
-    G: Fn(&'a BoxModule) -> Result<CodeGenerationResult>,
+    G: Fn(&'a BoxModule, RuntimeSpecSet) -> Result<Vec<(CodeGenerationResult, RuntimeSpec)>>,
   {
     let storage = match &self.storage {
       Some(s) => s,
       // no cache return directly
-      None => return Ok((generator(module)?, false)),
+      None => return Ok((generator(module, runtimes)?, false)),
     };
 
     let mut cache_id = None;
@@ -55,7 +57,7 @@ impl CodeGenerateOccasion {
     }
 
     // run generator and save to cache
-    let data = generator(module)?;
+    let data = generator(module, runtimes)?;
     if let Some(id) = cache_id {
       storage.set(id, data.clone());
     }

diff --git a/crates/rspack_core/src/code_generation_results.rs b/crates/rspack_core/src/code_generation_results.rs
@@ -1,17 +1,20 @@
 use std::collections::hash_map::Entry;
 use std::hash::Hash;
 use std::ops::{Deref, DerefMut};
+use std::sync::atomic::AtomicU32;
 
 use anymap::CloneAny;
+use once_cell::sync::Lazy;
 use rspack_error::{internal_error, Result};
 use rspack_hash::{HashDigest, HashFunction, HashSalt, RspackHash, RspackHashDigest};
 use rspack_identifier::IdentifierMap;
 use rspack_sources::BoxSource;
 use rustc_hash::FxHashMap as HashMap;
+use serde::Serialize;
 
 use crate::{
-  AssetInfo, ChunkInitFragments, ModuleIdentifier, RuntimeGlobals, RuntimeSpec, RuntimeSpecMap,
-  SourceType,
+  AssetInfo, ChunkInitFragments, ModuleIdentifier, RuntimeGlobals, RuntimeMode, RuntimeSpec,
+  RuntimeSpecMap, SourceType,
 };
 
 #[derive(Clone, Debug)]
@@ -86,6 +89,7 @@ pub struct CodeGenerationResult {
   pub chunk_init_fragments: ChunkInitFragments,
   pub runtime_requirements: RuntimeGlobals,
   pub hash: Option<RspackHashDigest>,
+  pub id: CodeGenResultId,
 }
 
 impl CodeGenerationResult {
@@ -133,14 +137,48 @@ impl CodeGenerationResult {
   }
 }
 
+#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize)]
+pub struct CodeGenResultId(u32);
+
+impl Default for CodeGenResultId {
+  fn default() -> Self {
+    Self(CODE_GEN_RESULT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed))
+  }
+}
+
+pub static CODE_GEN_RESULT_ID: Lazy<AtomicU32> = Lazy::new(|| AtomicU32::new(0));
+
 #[derive(Debug, Default)]
 pub struct CodeGenerationResults {
-  // TODO: This should be a map of ModuleIdentifier to CodeGenerationResult
-  pub module_generation_result_map: IdentifierMap<CodeGenerationResult>,
-  map: IdentifierMap<RuntimeSpecMap<ModuleIdentifier>>,
+  pub module_generation_result_map: HashMap<CodeGenResultId, CodeGenerationResult>,
+  map: IdentifierMap<RuntimeSpecMap<CodeGenResultId>>,
 }
 
 impl CodeGenerationResults {
+  pub fn get_one(&self, module_identifier: &ModuleIdentifier) -> Option<&CodeGenerationResult> {
+    self
+      .map
+      .get(module_identifier)
+      .and_then(|spec| match spec.mode {
+        RuntimeMode::Empty => None,
+        RuntimeMode::SingleEntry => spec
+          .single_value
+          .and_then(|result_id| self.module_generation_result_map.get(&result_id)),
+        RuntimeMode::Map => spec
+          .map
+          .values()
+          .next()
+          .and_then(|result_id| self.module_generation_result_map.get(result_id)),
+      })
+  }
+
+  pub fn clear_entry(
+    &mut self,
+    module_identifier: &ModuleIdentifier,
+  ) -> Option<(ModuleIdentifier, RuntimeSpecMap<CodeGenResultId>)> {
+    self.map.remove_entry(module_identifier)
+  }
+
   pub fn get(
     &self,
     module_identifier: &ModuleIdentifier,
@@ -195,7 +233,7 @@ impl CodeGenerationResults {
     &mut self,
     module_identifier: ModuleIdentifier,
     runtime: RuntimeSpec,
-    result: ModuleIdentifier,
+    result: CodeGenResultId,
   ) {
     match self.map.entry(module_identifier) {
       Entry::Occupied(mut record) => {

diff --git a/crates/rspack_core/src/compiler/compilation.rs b/crates/rspack_core/src/compiler/compilation.rs
@@ -960,55 +960,85 @@ impl Compilation {
       codegen_cache_counter: &mut Option<CacheCount>,
       filter_op: impl Fn(&(&ModuleIdentifier, &Box<dyn Module>)) -> bool + Sync + Send,
     ) -> Result<()> {
+      // If the runtime optimization is not opt out, a module codegen should be executed for each runtime.
+      // Else, share same codegen result for all runtimes.
+      let used_exports_optimization = compilation.options.is_new_tree_shaking()
+        && compilation.options.optimization.used_exports.is_true();
       let results = compilation
         .module_graph
         .modules()
         .par_iter()
         .filter(filter_op)
-        .filter(|(module_identifier, _)| {
+        .filter_map(|(module_identifier, module)| {
           let runtimes = compilation
             .chunk_graph
-            .get_module_runtimes(**module_identifier, &compilation.chunk_by_ukey);
-          !runtimes.is_empty()
-        })
-        .map(|(module_identifier, module)| {
-          compilation
+            .get_module_runtimes(*module_identifier, &compilation.chunk_by_ukey);
+          if runtimes.is_empty() {
+            return None;
+          }
+
+          let res = compilation
             .cache
             .code_generate_occasion
-            .use_cache(module, compilation, |module| {
-              module.code_generation(compilation)
+            .use_cache(module, runtimes, compilation, |module, runtimes| {
+              let take_length = if used_exports_optimization {
+                runtimes.len()
+              } else {
+                // Only codegen once
+                1
+              };
+              let mut codegen_list = vec![];
+              for runtime in runtimes.into_values().take(take_length) {
+                codegen_list.push((
+                  module.code_generation(compilation, Some(&runtime))?,
+                  runtime,
+                ));
+              }
+              Ok(codegen_list)
             })
-            .map(|(result, from_cache)| (*module_identifier, result, from_cache))
+            .map(|(result, from_cache)| (*module_identifier, result, from_cache));
+          Some(res)
         })
-        .collect::<Result<Vec<(ModuleIdentifier, CodeGenerationResult, bool)>>>()?;
-
+        .collect::<Result<
+          Vec<(
+            ModuleIdentifier,
+            Vec<(CodeGenerationResult, RuntimeSpec)>,
+            bool,
+          )>,
+        >>()?;
       results
         .into_iter()
-        .for_each(|(module_identifier, result, from_cache)| {
-          if let Some(counter) = codegen_cache_counter {
-            if from_cache {
-              counter.hit();
-            } else {
-              counter.miss();
+        .for_each(|(module_identifier, item, from_cache)| {
+          item.into_iter().for_each(|(result, runtime)| {
+            if let Some(counter) = codegen_cache_counter {
+              if from_cache {
+                counter.hit();
+              } else {
+                counter.miss();
+              }
             }
-          }
-          compilation.code_generated_modules.insert(module_identifier);
+            compilation.code_generated_modules.insert(module_identifier);
 
-          let runtimes = compilation
-            .chunk_graph
-            .get_module_runtimes(module_identifier, &compilation.chunk_by_ukey);
-
-          compilation
-            .code_generation_results
-            .module_generation_result_map
-            .insert(module_identifier, result);
-          for runtime in runtimes.values() {
-            compilation.code_generation_results.add(
-              module_identifier,
-              runtime.clone(),
-              module_identifier,
-            );
-          }
+            let runtimes = compilation
+              .chunk_graph
+              .get_module_runtimes(module_identifier, &compilation.chunk_by_ukey);
+            let result_id = result.id;
+            compilation
+              .code_generation_results
+              .module_generation_result_map
+              .insert(result.id, result);
+            if used_exports_optimization {
+              compilation
+                .code_generation_results
+                .add(module_identifier, runtime, result_id);
+            } else {
+              for runtime in runtimes.into_values() {
+                compilation
+                  .code_generation_results
+                  .add(module_identifier, runtime, result_id);
+              }
+            }
+          })
         });
       Ok(())
     }

diff --git a/crates/rspack_core/src/compiler/hmr.rs b/crates/rspack_core/src/compiler/hmr.rs
@@ -169,7 +169,9 @@ pub fn collect_changed_modules(
     .par_iter()
     .filter_map(|(identifier, cgm)| {
       let cid = cgm.id.as_deref();
-      if let Some(code_generation_result) = compilation.code_generation_results.module_generation_result_map.get(identifier) && let Some(module_hash) = &code_generation_result.hash && let Some(cid) = cid {
+      // TODO: Determine how to calc module hash if module related to multiple runtime code
+      // gen  
+      if let Some(code_generation_result) = compilation.code_generation_results.get_one(identifier) && let Some(module_hash) = &code_generation_result.hash && let Some(cid) = cid {
         Some((
             *identifier,
             (

diff --git a/crates/rspack_core/src/context_module.rs b/crates/rspack_core/src/context_module.rs
@@ -20,7 +20,7 @@ use crate::{
   BuildMeta, BuildResult, ChunkGraph, CodeGenerationResult, Compilation, ContextElementDependency,
   DependencyCategory, DependencyId, DependencyType, ExportsType, FakeNamespaceObjectMode,
   LibIdentOptions, Module, ModuleType, Resolve, ResolveInnerOptions,
-  ResolveOptionsWithDependencyType, ResolverFactory, RuntimeGlobals, SourceType,
+  ResolveOptionsWithDependencyType, ResolverFactory, RuntimeGlobals, RuntimeSpec, SourceType,
 };
 
 #[derive(Debug, Clone)]
@@ -482,7 +482,11 @@ impl Module for ContextModule {
     self.resolve_dependencies(build_context)
   }
 
-  fn code_generation(&self, compilation: &Compilation) -> Result<CodeGenerationResult> {
+  fn code_generation(
+    &self,
+    compilation: &Compilation,
+    _runtime: Option<&RuntimeSpec>,
+  ) -> Result<CodeGenerationResult> {
     let mut code_generation_result = CodeGenerationResult::default();
     code_generation_result
       .runtime_requirements

diff --git a/crates/rspack_core/src/dependency/dependency_template.rs b/crates/rspack_core/src/dependency/dependency_template.rs
@@ -3,13 +3,14 @@ use std::fmt::Debug;
 use dyn_clone::{clone_trait_object, DynClone};
 use rspack_sources::{BoxSource, ReplaceSource};
 
-use crate::{Compilation, Module, ModuleInitFragments, RuntimeGlobals};
+use crate::{Compilation, Module, ModuleInitFragments, RuntimeGlobals, RuntimeSpec};
 
 pub struct TemplateContext<'a, 'b> {
   pub compilation: &'a Compilation,
   pub module: &'a dyn Module,
   pub runtime_requirements: &'a mut RuntimeGlobals,
   pub init_fragments: &'a mut ModuleInitFragments<'b>,
+  pub runtime: Option<&'a RuntimeSpec>,
 }
 
 pub type TemplateReplaceSource = ReplaceSource<BoxSource>;

diff --git a/crates/rspack_core/src/external_module.rs b/crates/rspack_core/src/external_module.rs
@@ -14,7 +14,7 @@ use crate::{
   to_identifier, BuildContext, BuildInfo, BuildMetaExportsType, BuildResult, ChunkInitFragments,
   ChunkUkey, CodeGenerationDataUrl, CodeGenerationResult, Compilation, Context, ExternalType,
   InitFragmentExt, InitFragmentKey, InitFragmentStage, LibIdentOptions, Module, ModuleType,
-  NormalInitFragment, RuntimeGlobals, SourceType,
+  NormalInitFragment, RuntimeGlobals, RuntimeSpec, SourceType,
 };
 
 static EXTERNAL_MODULE_JS_SOURCE_TYPES: &[SourceType] = &[SourceType::JavaScript];
@@ -316,7 +316,11 @@ impl Module for ExternalModule {
     Ok(build_result.with_empty_diagnostic())
   }
 
-  fn code_generation(&self, compilation: &Compilation) -> Result<CodeGenerationResult> {
+  fn code_generation(
+    &self,
+    compilation: &Compilation,
+    _runtime: Option<&RuntimeSpec>,
+  ) -> Result<CodeGenerationResult> {
     let mut cgr = CodeGenerationResult::default();
     let (request, external_type) = self.get_request_and_external_type();
     match self.external_type.as_str() {

diff --git a/crates/rspack_core/src/missing_module.rs b/crates/rspack_core/src/missing_module.rs
@@ -6,7 +6,9 @@ use rspack_identifier::{Identifiable, Identifier};
 use rspack_sources::{RawSource, Source, SourceExt};
 use serde_json::json;
 
-use crate::{CodeGenerationResult, Compilation, Module, ModuleIdentifier, ModuleType, SourceType};
+use crate::{
+  CodeGenerationResult, Compilation, Module, ModuleIdentifier, ModuleType, RuntimeSpec, SourceType,
+};
 
 #[derive(Debug, Eq)]
 pub struct MissingModule {
@@ -52,7 +54,11 @@ impl Module for MissingModule {
     160.0
   }
 
-  fn code_generation(&self, compilation: &Compilation) -> Result<CodeGenerationResult> {
+  fn code_generation(
+    &self,
+    compilation: &Compilation,
+    _runtime: Option<&RuntimeSpec>,
+  ) -> Result<CodeGenerationResult> {
     let mut code_gen = CodeGenerationResult::default().with_javascript(
       RawSource::from(format!(
         "throw new Error({});\n",

diff --git a/crates/rspack_core/src/module.rs b/crates/rspack_core/src/module.rs
@@ -16,8 +16,8 @@ use crate::tree_shaking::visitor::OptimizeAnalyzeResult;
 use crate::{
   BoxDependency, ChunkUkey, CodeGenerationResult, Compilation, CompilerContext, CompilerOptions,
   ConnectionState, Context, ContextModule, DependencyId, DependencyTemplate, ExternalModule,
-  ModuleDependency, ModuleGraph, ModuleType, NormalModule, RawModule, Resolve, SharedPluginDriver,
-  SourceType,
+  ModuleDependency, ModuleGraph, ModuleType, NormalModule, RawModule, Resolve, RuntimeSpec,
+  SharedPluginDriver, SourceType,
 };
 
 pub struct BuildContext<'a> {
@@ -186,7 +186,11 @@ pub trait Module: Debug + Send + Sync + AsAny + DynHash + DynEq + Identifiable {
   ///
   /// Code generation will often iterate through every `source_types` given by the module
   /// to provide multiple code generation results for different `source_type`s.
-  fn code_generation(&self, _compilation: &Compilation) -> Result<CodeGenerationResult>;
+  fn code_generation(
+    &self,
+    _compilation: &Compilation,
+    _runtime: Option<&RuntimeSpec>,
+  ) -> Result<CodeGenerationResult>;
 
   /// Name matched against bundle-splitting conditions.
   fn name_for_condition(&self) -> Option<Box<str>> {
@@ -339,7 +343,7 @@ mod test {
   use super::Module;
   use crate::{
     BuildContext, BuildResult, CodeGenerationResult, Compilation, Context, ModuleExt, ModuleType,
-    SourceType,
+    RuntimeSpec, SourceType,
   };
 
   #[derive(Debug, Eq)]
@@ -409,7 +413,11 @@ mod test {
           unreachable!()
         }
 
-        fn code_generation(&self, _compilation: &Compilation) -> Result<CodeGenerationResult> {
+        fn code_generation(
+          &self,
+          _compilation: &Compilation,
+          _runtime: Option<&RuntimeSpec>,
+        ) -> Result<CodeGenerationResult> {
           unreachable!()
         }
       }