[Merged by Bors] - remove potential ub in render_resource_wrapper

crates/bevy_render/src/render_resource/resource_macros.rs

@@ -1,90 +1,90 @@
 // structs containing wgpu types take a long time to compile. this is particularly bad for generic
 // structs containing wgpu structs. we avoid that in debug builds (and for cargo check and rust analyzer)
-// by boxing and type-erasing with the `render_resource_wrapper` macro.
+// by type-erasing with the `render_resource_wrapper` macro. The resulting type behaves like Arc<$wgpu_type>,
+// but avoids explicitly storing an Arc<$wgpu_type> member.
 // analysis from https://github.com/bevyengine/bevy/pull/5950#issuecomment-1243473071 indicates this is
 // due to `evaluate_obligations`. we should check if this can be removed after a fix lands for
 // https://github.com/rust-lang/rust/issues/99188 (and after other `evaluate_obligations`-related changes).
 #[cfg(debug_assertions)]
 #[macro_export]
 macro_rules! render_resource_wrapper {
     ($wrapper_type:ident, $wgpu_type:ty) => {
-        #[derive(Clone, Debug)]
-        pub struct $wrapper_type(Option<std::sync::Arc<Box<()>>>);
+        #[derive(Debug)]
+        // SAFETY: when self.0 is Some(ptr), it comes from `into_raw` of an Arc<$wgpu_type> with a strong ref.
+        pub struct $wrapper_type(Option<*const ()>);
 
         impl $wrapper_type {
             pub fn new(value: $wgpu_type) -> Self {
-                unsafe {
-                    Self(Some(std::sync::Arc::new(std::mem::transmute(Box::new(
-                        value,
-                    )))))
-                }
+                let arc = std::sync::Arc::new(value);
+                let value_ptr = std::sync::Arc::into_raw(arc);
+                let unit_ptr = value_ptr.cast();
+                Self(Some(unit_ptr))
             }
 
             pub fn try_unwrap(mut self) -> Option<$wgpu_type> {
-                let inner = self.0.take();
-                if let Some(inner) = inner {
-                    match std::sync::Arc::try_unwrap(inner) {
-                        Ok(untyped_box) => {
-                            let typed_box = unsafe {
-                                std::mem::transmute::<Box<()>, Box<$wgpu_type>>(untyped_box)
-                            };
-                            Some(*typed_box)
-                        }
-                        Err(inner) => {
-                            let _ = unsafe {
-                                std::mem::transmute::<
-                                    std::sync::Arc<Box<()>>,
-                                    std::sync::Arc<Box<$wgpu_type>>,
-                                >(inner)
-                            };
-                            None
-                        }
-                    }
-                } else {
-                    None
-                }
+                self.0
+                    .take()
+                    .and_then(|unit_ptr| {
+                        let value_ptr = unit_ptr.cast::<$wgpu_type>();
+                        // SAFETY: pointer refers to a valid Arc, and was created from Arc::into_raw.
+                        // we take the pointer here, and this reconstructed arc is dropped/decremented within this scope.
+                        let arc = unsafe { std::sync::Arc::from_raw(value_ptr) };
+                        std::sync::Arc::try_unwrap(arc).ok()
+                    })
             }
         }
 
         impl std::ops::Deref for $wrapper_type {
             type Target = $wgpu_type;
 
             fn deref(&self) -> &Self::Target {
-                let untyped_box = self
-                    .0
-                    .as_ref()
-                    .expect("render_resource_wrapper inner value has already been taken (via drop or try_unwrap")
-                    .as_ref();
-
-                let typed_box =
-                    unsafe { std::mem::transmute::<&Box<()>, &Box<$wgpu_type>>(untyped_box) };
-                typed_box.as_ref()
+                let unit_ptr = self.0.as_ref().expect("render_resource_wrapper inner value has already been taken (via drop or try_unwrap");
+                let value_ptr = unit_ptr.cast::<$wgpu_type>();
+                // SAFETY: the arc lives for 'self, so the ref lives for 'self
+                let value_ref = unsafe { value_ptr.as_ref() };
+                value_ref.unwrap()
             }
         }
 
         impl Drop for $wrapper_type {
             fn drop(&mut self) {
-                let inner = self.0.take();
-                if let Some(inner) = inner {
-                    let _ = unsafe {
-                        std::mem::transmute::<
-                            std::sync::Arc<Box<()>>,
-                            std::sync::Arc<Box<$wgpu_type>>,
-                        >(inner)
-                    };
-                }
+                self.0
+                .map(|unit_ptr| {
+                    let value_ptr = unit_ptr.cast::<$wgpu_type>();
+                    // SAFETY: pointer refers to a valid Arc, and was created from Arc::into_raw.
+                    // we take the pointer here, and this reconstructed arc is dropped/decremented within this scope.
+                    unsafe { std::sync::Arc::from_raw(value_ptr) };
+                });
             }
         }
 
-        // Arc<Box<()>> and Arc<()> will be Sync and Send even when $wgpu_type is not Sync or Send.
+        // SAFETY: We manually implement Send and Sync, which is valid for Arc<T> when T: Send + Sync.
         // We ensure correctness by checking that $wgpu_type does implement Send and Sync.
         // If in future there is a case where a wrapper is required for a non-send/sync type
-        // we can implement a macro variant that also does `impl !Send for $wrapper_type {}` and
-        // `impl !Sync for $wrapper_type {}`
+        // we can implement a macro variant that omits these manual Send + Sync impls
+        unsafe impl Send for $wrapper_type {}
+        unsafe impl Sync for $wrapper_type {}
         const _: () = {
             trait AssertSendSyncBound: Send + Sync {}
             impl AssertSendSyncBound for $wgpu_type {}
         };
+
+
+        impl Clone for $wrapper_type {
+            fn clone(&self) -> Self {
+                let inner = self.0.map(|unit_ptr| {
+                    let value_ptr = unit_ptr.cast::<$wgpu_type>();
+                    // SAFETY: pointer refers to a valid Arc, and was created from Arc::into_raw.
+                    let arc = unsafe { std::sync::Arc::from_raw(value_ptr.cast::<$wgpu_type>()) };
+                    let cloned = std::sync::Arc::clone(&arc);
+                    // we have not taken the pointer, so we must forget the Arc to avoid decrementing the ref counter.
+                    std::mem::forget(arc);
+                    std::sync::Arc::into_raw(cloned).cast()
+                });
+
+                Self(inner)
+            }
+        }
     };
 }