fix(rust): fix bench for react by removing a lot of unsafe (rolldown#522

)
cgb-vue3 · Mar 10, 2024 · 71da395 · 71da395
1 parent a00f7b7
commit 71da395
Show file tree

Hide file tree

Showing 10 changed files with 55 additions and 41 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,6 +8,10 @@ homepage   = "https://github.com/rolldown-rs/rolldown"
 license    = "MIT"
 repository = "https://github.com/rolldown-rs/rolldown"
 
+[profile.release-debug]
+debug    = true
+inherits = "release"
+
 [workspace.lints.rust]
 
 [workspace.lints.clippy]

diff --git a/crates/rolldown/examples/basic.rs b/crates/rolldown/examples/basic.rs
@@ -8,7 +8,10 @@ async fn main() {
   let root = PathBuf::from(&std::env::var("CARGO_MANIFEST_DIR").unwrap());
   let cwd = root.join("./examples").into_normalize();
   let mut bundler = Bundler::new(InputOptions {
-    input: vec![InputItem { name: Some("basic".to_string()), import: "react-dom".to_string() }],
+    input: vec![
+      InputItem { name: Some("react-dom".to_string()), import: "react-dom".to_string() },
+      InputItem { name: Some("react".to_string()), import: "react".to_string() },
+    ],
     cwd,
     ..Default::default()
   });

diff --git a/crates/rolldown/src/chunk/mod.rs b/crates/rolldown/src/chunk/mod.rs
@@ -27,7 +27,7 @@ use crate::{
   },
 };
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct CrossChunkImportItem {
   pub export_alias: Option<Specifier>,
   pub import_ref: SymbolRef,

diff --git a/crates/rolldown/src/plugin_driver/mod.rs b/crates/rolldown/src/plugin_driver/mod.rs
@@ -66,6 +66,7 @@ impl PluginDriver {
   }
 
   pub async fn build_end(&self, args: Option<&HookBuildEndArgs>) -> HookNoopReturn {
+    tracing::info!("PluginDriver::build_end");
     for plugin in &self.plugins {
       plugin.build_end(&mut PluginContext::new(), args).await?;
     }

diff --git a/crates/rolldown/src/stages/bundle_stage/compute_cross_chunk_links.rs b/crates/rolldown/src/stages/bundle_stage/compute_cross_chunk_links.rs
@@ -1,4 +1,4 @@
-use std::{borrow::Cow, ptr::addr_of};
+use std::{borrow::Cow, sync::Mutex};
 
 use crate::{
   OutputFormat,
@@ -28,8 +28,8 @@ impl<'a> BundleStage<'a> {
     chunk_meta_imports: &mut ChunkMetaImports,
     chunk_meta_imports_from_external_modules: &mut ChunkMetaImportsForExternalModules,
   ) {
-    let symbols = &self.link_output.symbols;
-
+    let symbols = &Mutex::new(&mut self.link_output.symbols);
+    tracing::info!("collect_potential_chunk_imports");
     chunk_graph
       .chunks
       .iter_enumerated()
@@ -64,9 +64,9 @@ impl<'a> BundleStage<'a> {
             if !stmt_info.is_included {
               return;
             }
-
+            let mut symbols = symbols.lock().expect("ignore poison error");
             stmt_info.declared_symbols.iter().for_each(|declared| {
-              let symbol = symbols.get(*declared);
+              let symbol = symbols.get_mut(*declared);
               debug_assert!(
                 symbol.chunk_id.unwrap_or(chunk_id) == chunk_id,
                 "Symbol: {:?}, {:?} in {:?} should only belong to one chunk",
@@ -75,14 +75,11 @@ impl<'a> BundleStage<'a> {
                 module.resource_id,
               );
 
-              // safety: No two threads are ever writing to the same location
-              unsafe {
-                (*addr_of!(*symbols).cast_mut()).get_mut(*declared).chunk_id = Some(chunk_id);
-              }
+              symbol.chunk_id = Some(chunk_id);
             });
 
             stmt_info.referenced_symbols.iter().for_each(|referenced| {
-              let mut canonical_ref = self.link_output.symbols.par_canonical_ref_for(*referenced);
+              let mut canonical_ref = symbols.par_canonical_ref_for(*referenced);
               if let Some(namespace_alias) = &symbols.get(canonical_ref).namespace_alias {
                 canonical_ref = namespace_alias.namespace_ref;
               }
@@ -100,9 +97,9 @@ impl<'a> BundleStage<'a> {
             chunk_meta_imports
               .insert(entry_linking_info.wrapper_ref.expect("cjs should be wrapped in esm output"));
           }
+          let symbols = symbols.lock().expect("ignore poison error");
           for export_ref in entry_linking_info.resolved_exports.values() {
-            let mut canonical_ref =
-              self.link_output.symbols.par_canonical_ref_for(export_ref.symbol_ref);
+            let mut canonical_ref = symbols.par_canonical_ref_for(export_ref.symbol_ref);
             let symbol = symbols.get(canonical_ref);
             if let Some(ns_alias) = &symbol.namespace_alias {
               canonical_ref = ns_alias.namespace_ref;
@@ -111,6 +108,7 @@ impl<'a> BundleStage<'a> {
           }
         }
       });
+    tracing::info!("collect_potential_chunk_imports end");
   }
 
   pub fn compute_cross_chunk_links(&mut self, chunk_graph: &mut ChunkGraph) {
@@ -120,12 +118,18 @@ impl<'a> BundleStage<'a> {
       index_vec![FxHashSet::<SymbolRef>::default(); chunk_graph.chunks.len()];
     let mut chunk_meta_imports_from_external_modules_vec: ChunkMetaImportsForExternalModules = index_vec![FxHashMap::<ExternalModuleId, Vec<NamedImport>>::default(); chunk_graph.chunks.len()];
 
+    let mut imports_from_other_chunks_vec: IndexVec<
+      ChunkId,
+      FxHashMap<ChunkId, Vec<CrossChunkImportItem>>,
+    > = index_vec![FxHashMap::<ChunkId, Vec<CrossChunkImportItem>>::default(); chunk_graph.chunks.len()];
+
     self.collect_potential_chunk_imports(
       chunk_graph,
       &mut chunk_meta_imports_vec,
       &mut chunk_meta_imports_from_external_modules_vec,
     );
 
+    tracing::info!("calculate cross chunk imports");
     // - Find out what imports are actually come from other chunks
     chunk_graph.chunks.iter_enumerated().for_each(|(chunk_id, chunk)| {
       let chunk_meta_imports = &chunk_meta_imports_vec[chunk_id];
@@ -142,10 +146,7 @@ impl<'a> BundleStage<'a> {
         });
         // Check if the import is from another chunk
         if chunk_id != importee_chunk_id {
-          let imports_from_other_chunks =
-          // Safety: No race condition here:
-          // - This loop is executed sequentially in a single thread
-            unsafe { &mut (*addr_of!(chunk.imports_from_other_chunks).cast_mut()) };
+          let imports_from_other_chunks = &mut imports_from_other_chunks_vec[chunk_id];
           imports_from_other_chunks
             .entry(importee_chunk_id)
             .or_default()
@@ -170,13 +171,13 @@ impl<'a> BundleStage<'a> {
               && importee_chunk.modules.first().copied() != Some(self.link_output.runtime.id())
           })
           .for_each(|(importee_chunk_id, _)| {
-            let imports_from_other_chunks =
-              unsafe { &mut (*addr_of!(chunk.imports_from_other_chunks).cast_mut()) };
+            let imports_from_other_chunks = &mut imports_from_other_chunks_vec[chunk_id];
             imports_from_other_chunks.entry(importee_chunk_id).or_default();
           });
       }
     });
 
+    tracing::info!("Generate cross-chunk exports");
     // Generate cross-chunk exports. These must be computed before cross-chunk
     // imports because of export alias renaming, which must consider all export
     // aliases simultaneously to avoid collisions.
@@ -197,27 +198,27 @@ impl<'a> BundleStage<'a> {
       }
     }
     for chunk_id in chunk_graph.chunks.indices() {
-      for (importee_chunk_id, import_items) in
-        &chunk_graph.chunks[chunk_id].imports_from_other_chunks
-      {
+      for (importee_chunk_id, import_items) in &mut imports_from_other_chunks_vec[chunk_id] {
         for item in import_items {
           if let Some(alias) =
             chunk_graph.chunks[*importee_chunk_id].exports_to_other_chunks.get(&item.import_ref)
           {
-            // safety: no other mutable reference to `item` exists
-            unsafe {
-              let item = (item as *const CrossChunkImportItem).cast_mut();
-              (*item).export_alias = Some(alias.clone().into());
-            }
+            item.export_alias = Some(alias.clone().into());
           }
         }
       }
     }
 
-    chunk_meta_imports_from_external_modules_vec.into_iter_enumerated().for_each(
-      |(chunk_id, imports_from_external_modules)| {
-        chunk_graph.chunks[chunk_id].imports_from_external_modules = imports_from_external_modules;
-      },
-    );
+    chunk_graph
+      .chunks
+      .iter_mut()
+      .zip(
+        imports_from_other_chunks_vec.into_iter().zip(chunk_meta_imports_from_external_modules_vec),
+      )
+      .par_bridge()
+      .for_each(|(chunk, (imports_from_other_chunks, imports_from_external_modules))| {
+        chunk.imports_from_other_chunks = imports_from_other_chunks;
+        chunk.imports_from_external_modules = imports_from_external_modules;
+      });
   }
 }
diff --git a/crates/rolldown/src/stages/bundle_stage/mod.rs b/crates/rolldown/src/stages/bundle_stage/mod.rs
@@ -40,12 +40,14 @@ impl<'a> BundleStage<'a> {
   #[tracing::instrument(skip_all)]
   pub async fn bundle(&mut self) -> BatchedResult<Vec<Output>> {
     use rayon::prelude::*;
-    tracing::debug!("Start bundle stage");
+    tracing::info!("Start bundle stage");
     let mut chunk_graph = self.generate_chunks();
 
     self.generate_chunk_filenames(&mut chunk_graph);
+    tracing::info!("generate_chunk_filenames");
 
     self.compute_cross_chunk_links(&mut chunk_graph);
+    tracing::info!("compute_cross_chunk_links");
 
     chunk_graph.chunks.iter_mut().par_bridge().for_each(|chunk| {
       chunk.de_conflict(self.link_output);
@@ -78,6 +80,7 @@ impl<'a> BundleStage<'a> {
           ast,
         );
       });
+    tracing::info!("finalizing modules");
 
     let chunks = chunk_graph.chunks.iter().map(|c| {
       let ret =
@@ -125,6 +128,8 @@ impl<'a> BundleStage<'a> {
       },
     )?;
 
+    tracing::info!("rendered chunks");
+
     Ok(assets)
   }
 

diff --git a/crates/rolldown/src/stages/link_stage/mod.rs b/crates/rolldown/src/stages/link_stage/mod.rs
@@ -111,6 +111,7 @@ impl<'a> LinkStage<'a> {
   }
 
   pub fn link(mut self) -> LinkStageOutput {
+    tracing::info!("Start link stage");
     self.sort_modules();
 
     self.determine_module_exports_kind();

diff --git a/crates/rolldown/src/stages/scan_stage.rs b/crates/rolldown/src/stages/scan_stage.rs
@@ -50,6 +50,7 @@ impl<Fs: FileSystem + Default + 'static> ScanStage<Fs> {
 
   #[tracing::instrument(skip_all)]
   pub async fn scan(&self) -> BatchedResult<ScanStageOutput> {
+    tracing::info!("Start scan stage");
     assert!(!self.input_options.input.is_empty(), "You must supply options.input to rolldown");
 
     let mut module_loader = ModuleLoader::new(

diff --git a/crates/rolldown/src/utils/render_chunks.rs b/crates/rolldown/src/utils/render_chunks.rs
@@ -15,11 +15,9 @@ pub async fn render_chunks<'a>(
 ) -> Result<Vec<(String, Option<SourceMap>, RenderedChunk)>, BatchedErrors> {
   // TODO support `render_chunk` hook return map
   let result = block_on_spawn_all(chunks.map(|(content, map, rendered_chunk)| async move {
+    tracing::info!("render_chunks");
     match plugin_driver
-      .render_chunk(RenderChunkArgs {
-        code: content,
-        chunk: unsafe { std::mem::transmute(&rendered_chunk) },
-      })
+      .render_chunk(RenderChunkArgs { code: content, chunk: &rendered_chunk })
       .await
     {
       Ok(value) => Ok((value, map, rendered_chunk)),

diff --git a/packages/bench/index.js b/packages/bench/index.js
@@ -34,11 +34,11 @@ const suites = [
     benchIteration: 3,
   },
   {
-    title: 'vue',
+    title: 'vue-stack',
     inputs: [path.join(dirname, 'vue-entry.js')],
   },
   {
-    title: 'react_and_react_dom',
+    title: 'react-stack',
     inputs: ['react', 'react-dom'],
   },
 ]
@@ -53,7 +53,7 @@ async function runRolldown(item) {
       // TODO
       // For now these are needed to align better w/ esbuild & Vite behavior
       // because internally we are still using the default behavior of oxc
-      // resovler. We should ship a more sensible resolver default that aligns
+      // resolver. We should ship a more sensible resolver default that aligns
       // with Vite's.
       conditionNames: ['import'],
       mainFields: ['module', 'browser', 'main'],