fix(rust/rolldown): prevent segment fault in generate_bundle by rem…

…oving unsafe code (#1905)

<!-- Thank you for contributing! -->

### Description

Another fix compared to #1904
for the segment fault problem.

<!-- Please insert your description here and provide especially info
about the "what" this PR is solving -->

crates/rolldown_binding/src/lib.rs

@@ -24,3 +24,4 @@ pub mod types;
 pub mod utils;
 mod worker_manager;
 pub use oxc_transform_napi;
+mod type_aliases;

crates/rolldown_binding/src/options/plugin/js_plugin.rs

@@ -3,7 +3,13 @@ use crate::types::{
   js_callback::MaybeAsyncJsCallbackExt,
 };
 use rolldown_plugin::{Plugin, __inner::SharedPluginable, typedmap::TypedMapKey};
-use std::{borrow::Cow, ops::Deref, sync::Arc};
+use rolldown_utils::unique_arc::UniqueArc;
+use std::{
+  borrow::Cow,
+  mem,
+  ops::Deref,
+  sync::{Arc, Mutex},
+};
 
 use super::{
   binding_transform_context::BindingTransformPluginContext,
@@ -295,12 +301,10 @@ impl Plugin for JsPlugin {
     is_write: bool,
   ) -> rolldown_plugin::HookNoopReturn {
     if let Some(cb) = &self.generate_bundle {
-      cb.await_call((
-        Arc::clone(ctx).into(),
-        BindingOutputs::new(unsafe { std::mem::transmute(bundle) }),
-        is_write,
-      ))
-      .await?;
+      let old_bundle = UniqueArc::new(Mutex::new(mem::take(bundle)));
+      cb.await_call((Arc::clone(ctx).into(), BindingOutputs::new(old_bundle.weak_ref()), is_write))
+        .await?;
+      *bundle = old_bundle.into_inner().into_inner()?;
     }
     Ok(())
   }
@@ -311,11 +315,9 @@ impl Plugin for JsPlugin {
     bundle: &mut Vec<rolldown_common::Output>,
   ) -> rolldown_plugin::HookNoopReturn {
     if let Some(cb) = &self.write_bundle {
-      cb.await_call((
-        Arc::clone(ctx).into(),
-        BindingOutputs::new(unsafe { std::mem::transmute(bundle) }),
-      ))
-      .await?;
+      let old_bundle = UniqueArc::new(Mutex::new(mem::take(bundle)));
+      cb.await_call((Arc::clone(ctx).into(), BindingOutputs::new(old_bundle.weak_ref()))).await?;
+      *bundle = old_bundle.into_inner().into_inner()?;
     }
     Ok(())
   }

crates/rolldown_binding/src/type_aliases.rs

@@ -0,0 +1,6 @@
+use std::sync::Mutex;
+
+use rolldown_utils::unique_arc::{UniqueArc, WeakRef};
+
+pub type UniqueArcMutex<T> = UniqueArc<Mutex<T>>;
+pub type WeakRefMutex<T> = WeakRef<Mutex<T>>;

crates/rolldown_binding/src/types/binding_outputs.rs

@@ -1,28 +1,32 @@
 use napi_derive::napi;
 
+use crate::type_aliases::{UniqueArcMutex, WeakRefMutex};
+
 use super::{binding_output_asset::BindingOutputAsset, binding_output_chunk::BindingOutputChunk};
 
 /// The `BindingOutputs` owner `Vec<Output>` the mutable reference, it avoid `Clone` at call `writeBundle/generateBundle` hook, and make it mutable.
 #[napi]
 pub struct BindingOutputs {
-  inner: &'static mut Vec<rolldown_common::Output>,
+  inner: WeakRefMutex<Vec<rolldown_common::Output>>,
 }
 
 #[napi]
 impl BindingOutputs {
-  pub fn new(inner: &'static mut Vec<rolldown_common::Output>) -> Self {
+  pub fn new(inner: WeakRefMutex<Vec<rolldown_common::Output>>) -> Self {
     Self { inner }
   }
 
   #[napi(getter)]
   pub fn chunks(&mut self) -> Vec<BindingOutputChunk> {
     let mut chunks: Vec<BindingOutputChunk> = vec![];
-
-    self.inner.iter_mut().for_each(|o| match o {
-      rolldown_common::Output::Chunk(chunk) => {
-        chunks.push(BindingOutputChunk::new(unsafe { std::mem::transmute(chunk.as_mut()) }));
-      }
-      rolldown_common::Output::Asset(_) => {}
+    self.inner.with_inner(|inner| {
+      let mut inner = inner.lock().expect("PoisonError raised");
+      inner.iter_mut().for_each(|o| match o {
+        rolldown_common::Output::Chunk(chunk) => {
+          chunks.push(BindingOutputChunk::new(unsafe { std::mem::transmute(chunk.as_mut()) }));
+        }
+        rolldown_common::Output::Asset(_) => {}
+      });
     });
 
     chunks
@@ -32,45 +36,53 @@ impl BindingOutputs {
   pub fn assets(&mut self) -> Vec<BindingOutputAsset> {
     let mut assets: Vec<BindingOutputAsset> = vec![];
 
-    self.inner.iter_mut().for_each(|o| match o {
-      rolldown_common::Output::Asset(asset) => {
-        assets.push(BindingOutputAsset::new(unsafe { std::mem::transmute(asset.as_mut()) }));
-      }
-      rolldown_common::Output::Chunk(_) => {}
+    self.inner.with_inner(|inner| {
+      let mut inner = inner.lock().expect("PoisonError raised");
+      inner.iter_mut().for_each(|o| match o {
+        rolldown_common::Output::Asset(asset) => {
+          assets.push(BindingOutputAsset::new(unsafe { std::mem::transmute(asset.as_mut()) }));
+        }
+        rolldown_common::Output::Chunk(_) => {}
+      });
     });
     assets
   }
 
   #[napi]
   pub fn delete(&mut self, file_name: String) {
-    if let Some(index) = self.inner.iter().position(|o| o.filename() == file_name) {
-      self.inner.remove(index);
-    }
+    self.inner.with_inner(|inner| {
+      let mut inner = inner.lock().expect("PoisonError raised");
+      if let Some(index) = inner.iter().position(|o| o.filename() == file_name) {
+        inner.remove(index);
+      }
+    });
   }
 }
 
 /// The `FinalBindingOutputs` is used at `write()` or `generate()`, it is similar to `BindingOutputs`, if using `BindingOutputs` has unexpected behavior.
 /// TODO find a way to export it gracefully.
 #[napi]
 pub struct FinalBindingOutputs {
-  inner: Vec<rolldown_common::Output>,
+  inner: UniqueArcMutex<Vec<rolldown_common::Output>>,
 }
 
 #[napi]
 impl FinalBindingOutputs {
   pub fn new(inner: Vec<rolldown_common::Output>) -> Self {
-    Self { inner }
+    Self { inner: UniqueArcMutex::new(inner.into()) }
   }
 
   #[napi(getter)]
   pub fn chunks(&mut self) -> Vec<BindingOutputChunk> {
     let mut chunks: Vec<BindingOutputChunk> = vec![];
-
-    self.inner.iter_mut().for_each(|o| match o {
-      rolldown_common::Output::Chunk(chunk) => {
-        chunks.push(BindingOutputChunk::new(unsafe { std::mem::transmute(chunk.as_mut()) }));
-      }
-      rolldown_common::Output::Asset(_) => {}
+    self.inner.weak_ref().with_inner(|inner| {
+      let mut inner = inner.lock().expect("PoisonError raised");
+      inner.iter_mut().for_each(|o| match o {
+        rolldown_common::Output::Chunk(chunk) => {
+          chunks.push(BindingOutputChunk::new(unsafe { std::mem::transmute(chunk.as_mut()) }));
+        }
+        rolldown_common::Output::Asset(_) => {}
+      });
     });
 
     chunks
@@ -80,11 +92,14 @@ impl FinalBindingOutputs {
   pub fn assets(&mut self) -> Vec<BindingOutputAsset> {
     let mut assets: Vec<BindingOutputAsset> = vec![];
 
-    self.inner.iter_mut().for_each(|o| match o {
-      rolldown_common::Output::Asset(asset) => {
-        assets.push(BindingOutputAsset::new(unsafe { std::mem::transmute(asset.as_mut()) }));
-      }
-      rolldown_common::Output::Chunk(_) => {}
+    self.inner.weak_ref().with_inner(|inner| {
+      let mut inner = inner.lock().expect("PoisonError raised");
+      inner.iter_mut().for_each(|o| match o {
+        rolldown_common::Output::Asset(asset) => {
+          assets.push(BindingOutputAsset::new(unsafe { std::mem::transmute(asset.as_mut()) }));
+        }
+        rolldown_common::Output::Chunk(_) => {}
+      });
     });
     assets
   }

crates/rolldown_utils/src/lib.rs

@@ -17,3 +17,4 @@ pub mod rustc_hash;
 pub mod sanitize_file_name;
 pub mod xxhash;
 pub use bitset::BitSet;
+pub mod unique_arc;

crates/rolldown_utils/src/unique_arc.rs

@@ -0,0 +1,57 @@
+use std::sync::{Arc, Weak};
+
+use crate::debug::pretty_type_name;
+
+#[derive(Debug)]
+pub struct UniqueArc<T: ?Sized>(Arc<T>);
+
+impl<T> UniqueArc<T> {
+  pub fn new(value: T) -> Self {
+    Self(Arc::new(value))
+  }
+
+  /// # Panics
+  /// - Don't call this method while calling the `WeakRef#with_inner` method.
+  pub fn into_inner(self) -> T {
+    Arc::into_inner(self.0).unwrap_or_else(|| {
+      let t_name = pretty_type_name::<T>();
+      panic!("UniqueArc<{t_name}> has multiple references")
+    })
+  }
+
+  pub fn weak_ref(&self) -> WeakRef<T> {
+    WeakRef(Arc::downgrade(&self.0))
+  }
+}
+
+#[derive(Debug)]
+pub struct WeakRef<T: ?Sized>(Weak<T>);
+
+impl<T> WeakRef<T> {
+  /// This API pattern is intended to prevent users from getting `Arc<T>` and storing it.
+  /// This will cause `UniqueArc#into_inner` panic if `Arc<T>` got stored.
+  pub fn try_with_inner<F, R>(&self, f: F) -> Option<R>
+  where
+    F: FnOnce(&T) -> R,
+  {
+    self.0.upgrade().map(|arc| f(&*arc))
+  }
+
+  pub fn with_inner<F, R>(&self, f: F) -> R
+  where
+    F: FnOnce(&T) -> R,
+  {
+    self.try_with_inner(f).unwrap_or_else(|| {
+      let t_name = pretty_type_name::<T>();
+      panic!(
+        "UniqueArc<{t_name}> is already dropped. You can't access it by this WeakRef<{t_name}>",
+      )
+    })
+  }
+}
+
+impl<T> Clone for WeakRef<T> {
+  fn clone(&self) -> Self {
+    Self(Weak::clone(&self.0))
+  }
+}