More ReadRef fixes using VcCellMode::raw_cell API

crates/turbo-tasks/src/backend.rs

@@ -1,5 +1,4 @@
 use std::{
-    any::Any,
     borrow::Cow,
     fmt::{self, Debug, Display, Write},
     future::Future,
@@ -22,8 +21,9 @@ use crate::{
     raw_vc::CellId,
     registry,
     trait_helpers::{get_trait_method, has_trait, traits},
+    triomphe_utils::unchecked_sidecast_triomphe_arc,
     FunctionId, RawVc, ReadRef, SharedReference, TaskId, TaskIdProvider, TaskIdSet, TraitRef,
-    TraitTypeId, ValueTypeId, VcValueTrait, VcValueType,
+    TraitTypeId, ValueTypeId, VcRead, VcValueTrait, VcValueType,
 };
 
 pub enum TaskType {
@@ -347,11 +347,14 @@ impl Display for CellContent {
 }
 
 impl TypedCellContent {
-    pub fn cast<T: Any + VcValueType>(self) -> Result<ReadRef<T>> {
+    pub fn cast<T: VcValueType>(self) -> Result<ReadRef<T>> {
         let data = self.1 .0.ok_or_else(|| anyhow!("Cell is empty"))?;
         let data = data
-            .downcast()
+            .downcast::<<T::Read as VcRead<T>>::Repr>()
             .map_err(|_err| anyhow!("Unexpected type in cell"))?;
+        // SAFETY: `T` and `T::Read::Repr` must have equivalent memory representations,
+        // guaranteed by the unsafe implementation of `VcValueType`.
+        let data = unsafe { unchecked_sidecast_triomphe_arc(data) };
         Ok(ReadRef::new_arc(data))
     }
 
@@ -374,10 +377,6 @@ impl TypedCellContent {
         )
     }
 
-    pub fn try_cast<T: Any + VcValueType>(self) -> Option<ReadRef<T>> {
-        Some(ReadRef::new_arc(self.1 .0?.downcast().ok()?))
-    }
-
     pub fn into_untyped(self) -> CellContent {
         self.1
     }

crates/turbo-tasks/src/completion.rs

@@ -36,7 +36,7 @@ impl Completion {
         // only updates the cell when it is empty (Completion::cell opted-out of
         // that via `#[turbo_tasks::value(cell = "new")]`)
         let cell = turbo_tasks::macro_helpers::find_cell_by_type(*COMPLETION_VALUE_TYPE_ID);
-        cell.conditional_update_shared(|old| old.is_none().then_some(Completion));
+        cell.conditional_update(|old| old.is_none().then_some(Completion));
         let raw: RawVc = cell.into();
         raw.into()
     }

crates/turbo-tasks/src/manager.rs

@@ -1522,57 +1522,144 @@ pub struct CurrentCellRef {
     index: CellId,
 }
 
+type VcReadRepr<T> = <<T as VcValueType>::Read as VcRead<T>>::Repr;
+
 impl CurrentCellRef {
-    pub fn conditional_update_shared<
-        T: VcValueType + 'static,
-        F: FnOnce(Option<&T>) -> Option<T>,
-    >(
+    /// Updates the cell if the given `functor` returns a value.
+    pub fn conditional_update<T>(&self, functor: impl FnOnce(Option<&T>) -> Option<T>)
+    where
+        T: VcValueType,
+    {
+        self.conditional_update_with_shared_reference(|old_shared_reference| {
+            let old_ref = old_shared_reference
+                .and_then(|sr| sr.0.downcast_ref::<VcReadRepr<T>>())
+                .map(|content| <T::Read as VcRead<T>>::repr_to_value_ref(content));
+            let new_value = functor(old_ref)?;
+            Some(SharedReference::new(triomphe::Arc::new(
+                <T::Read as VcRead<T>>::value_to_repr(new_value),
+            )))
+        })
+    }
+
+    /// Updates the cell if the given `functor` returns a `SharedReference`.
+    pub fn conditional_update_with_shared_reference(
         &self,
-        functor: F,
+        functor: impl FnOnce(Option<&SharedReference>) -> Option<SharedReference>,
     ) {
         let tt = turbo_tasks();
-        let content = tt
-            .read_own_task_cell(self.current_task, self.index)
-            .ok()
-            .and_then(|v| v.try_cast::<T>());
-        let update =
-            functor(content.as_deref().map(|content| {
-                <<T as VcValueType>::Read as VcRead<T>>::target_to_value_ref(content)
-            }));
+        let cell_content = tt.read_own_task_cell(self.current_task, self.index).ok();
+        let update = functor(cell_content.as_ref().and_then(|cc| cc.1 .0.as_ref()));
         if let Some(update) = update {
-            tt.update_own_task_cell(
-                self.current_task,
-                self.index,
-                CellContent(Some(SharedReference::new(triomphe::Arc::new(update)))),
-            )
+            tt.update_own_task_cell(self.current_task, self.index, CellContent(Some(update)))
         }
     }
 
-    pub fn compare_and_update_shared<T: PartialEq + VcValueType + 'static>(&self, new_content: T) {
-        self.conditional_update_shared(|old_content| {
-            if let Some(old_content) = old_content {
-                if PartialEq::eq(&new_content, old_content) {
+    /// Replace the current cell's content with `new_value` if the current
+    /// content is not equal by value with the existing content.
+    ///
+    /// The comparison happens using the value itself, not the
+    /// [`VcRead::Target`] of that value.
+    ///
+    /// Take this example of a custom equality implementation on a transparent
+    /// wrapper type:
+    ///
+    /// ```
+    /// #[turbo_tasks::value(transparent, eq = "manual")]
+    /// struct Wrapper(Vec<u32>);
+    ///
+    /// impl PartialEq for Wrapper {
+    ///     fn eq(&self, other: Wrapper) {
+    ///         // Example: order doesn't matter for equality
+    ///         let (mut this, mut other) = (self.clone(), other.clone());
+    ///         this.sort_unstable();
+    ///         other.sort_unstable();
+    ///         this == other
+    ///     }
+    /// }
+    ///
+    /// impl Eq for Wrapper {}
+    /// ```
+    ///
+    /// Comparisons of `Vc<Wrapper>` used when updating the cell will use
+    /// `Wrapper`'s custom equality implementation, rather than the one
+    /// provided by the target (`Vec<u32>`) type.
+    ///
+    /// However, in most cases, the default derived implementation of
+    /// `PartialEq` is used which just forwards to the inner value's
+    /// `PartialEq`.
+    pub fn compare_and_update<T>(&self, new_value: T)
+    where
+        T: PartialEq + VcValueType,
+    {
+        self.conditional_update(|old_value| {
+            if let Some(old_value) = old_value {
+                if old_value == &new_value {
                     return None;
                 }
             }
-            Some(new_content)
+            Some(new_value)
         });
     }
 
-    pub fn update_shared<T: VcValueType + 'static>(&self, new_content: T) {
+    /// Replace the current cell's content with `new_shared_reference` if the
+    /// current content is not equal by value with the existing content.
+    ///
+    /// If you already have a `SharedReference`, this is a faster version of
+    /// [`CurrentCellRef::compare_and_update`].
+    ///
+    /// The [`SharedReference`] is expected to use the `<T::Read as
+    /// VcRead<T>>::Repr` type for its representation of the value.
+    pub fn compare_and_update_with_shared_reference<T>(&self, new_shared_reference: SharedReference)
+    where
+        T: VcValueType + PartialEq,
+    {
+        fn extract_sr_value<T: VcValueType>(sr: &SharedReference) -> &T {
+            <T::Read as VcRead<T>>::repr_to_value_ref(
+                sr.0.downcast_ref::<VcReadRepr<T>>()
+                    .expect("cannot update SharedReference of different type"),
+            )
+        }
+        self.conditional_update_with_shared_reference(|old_sr| {
+            if let Some(old_sr) = old_sr {
+                let old_value: &T = extract_sr_value(old_sr);
+                let new_value = extract_sr_value(&new_shared_reference);
+                if old_value == new_value {
+                    return None;
+                }
+            }
+            Some(new_shared_reference)
+        });
+    }
+
+    /// Unconditionally updates the content of the cell.
+    pub fn update<T>(&self, new_value: T)
+    where
+        T: VcValueType,
+    {
         let tt = turbo_tasks();
         tt.update_own_task_cell(
             self.current_task,
             self.index,
-            CellContent(Some(SharedReference::new(triomphe::Arc::new(new_content)))),
+            CellContent(Some(SharedReference::new(triomphe::Arc::new(
+                <T::Read as VcRead<T>>::value_to_repr(new_value),
+            )))),
         )
     }
 
-    pub fn update_shared_reference(&self, shared_ref: SharedReference) {
+    /// A faster version of [`Self::update`] if you already have a
+    /// [`SharedReference`].
+    ///
+    /// If the passed-in [`SharedReference`] is the same as the existing cell's
+    /// by identity, no update is performed.
+    ///
+    /// The [`SharedReference`] is expected to use the `<T::Read as
+    /// VcRead<T>>::Repr` type for its representation of the value.
+    pub fn update_with_shared_reference(&self, shared_ref: SharedReference) {
         let tt = turbo_tasks();
         let content = tt.read_own_task_cell(self.current_task, self.index).ok();
-        let update = if let Some(TypedCellContent(_, CellContent(shared_ref_exp))) = content {
-            shared_ref_exp.as_ref().ne(&Some(&shared_ref))
+        let update = if let Some(TypedCellContent(_, CellContent(Some(shared_ref_exp)))) = content {
+            // pointer equality (not value equality)
+            shared_ref_exp != shared_ref
         } else {
             true
         };

crates/turbo-tasks/src/read_ref.rs

@@ -10,9 +10,9 @@ use turbo_tasks_hash::DeterministicHash;
 
 use crate::{
     debug::{ValueDebugFormat, ValueDebugFormatString},
-    macro_helpers::find_cell_by_type,
     trace::{TraceRawVcs, TraceRawVcsContext},
     triomphe_utils::unchecked_sidecast_triomphe_arc,
+    vc::VcCellMode,
     SharedReference, Vc, VcRead, VcValueType,
 };
 
@@ -244,14 +244,15 @@ where
     /// reference.
     pub fn cell(read_ref: ReadRef<T>) -> Vc<T> {
         let type_id = T::get_value_type_id();
-        let local_cell = find_cell_by_type(type_id);
         // SAFETY: `T` and `T::Read::Repr` must have equivalent memory representations,
         // guaranteed by the unsafe implementation of `VcValueType`.
-        local_cell.update_shared_reference(SharedReference::new(unsafe {
+        let value = unsafe {
             unchecked_sidecast_triomphe_arc::<T, <T::Read as VcRead<T>>::Repr>(read_ref.0)
-        }));
+        };
         Vc {
-            node: local_cell.into(),
+            node: <T::CellMode as VcCellMode<T>>::raw_cell(
+                SharedReference::new(value).typed(type_id),
+            ),
             _t: PhantomData,
         }
     }

crates/turbo-tasks/src/trait_ref.rs

@@ -102,7 +102,7 @@ where
         // See Safety clause above.
         let TypedSharedReference(ty, shared_ref) = trait_ref.shared_reference;
         let local_cell = find_cell_by_type(ty);
-        local_cell.update_shared_reference(shared_ref);
+        local_cell.update_with_shared_reference(shared_ref);
         let raw_vc: RawVc = local_cell.into();
         raw_vc.into()
     }

crates/turbo-tasks/src/vc/cast.rs

@@ -1,8 +1,8 @@
-use std::{marker::PhantomData, mem::ManuallyDrop};
+use std::marker::PhantomData;
 
 use anyhow::Result;
 
-use crate::{backend::TypedCellContent, ReadRef, TraitRef, VcRead, VcValueTrait, VcValueType};
+use crate::{backend::TypedCellContent, ReadRef, TraitRef, VcValueTrait, VcValueType};
 
 /// Trait defined to share behavior between values and traits within
 /// [`ReadRawVcFuture`][crate::ReadRawVcFuture]. See [`VcValueTypeCast`] and
@@ -27,19 +27,7 @@ where
     type Output = ReadRef<T>;
 
     fn cast(content: TypedCellContent) -> Result<Self::Output> {
-        Ok(
-            // Safety: the `VcValueType` implementor must guarantee that both `T` and
-            // `Repr` are #[repr(transparent)].
-            unsafe {
-                // Downcast the cell content to the expected representation type, then
-                // transmute it to the expected type.
-                // See https://users.rust-lang.org/t/transmute-doesnt-work-on-generic-types/87272/9
-                std::mem::transmute_copy::<
-                    ManuallyDrop<ReadRef<<T::Read as VcRead<T>>::Repr>>,
-                    Self::Output,
-                >(&ManuallyDrop::new(content.cast()?))
-            },
-        )
+        content.cast()
     }
 }
 

crates/turbo-tasks/src/vc/cell_mode.rs

@@ -1,16 +1,25 @@
-use std::marker::PhantomData;
+use std::{any::type_name, marker::PhantomData};
 
 use super::{read::VcRead, traits::VcValueType};
-use crate::{manager::find_cell_by_type, Vc};
+use crate::{manager::find_cell_by_type, task::shared_reference::TypedSharedReference, RawVc, Vc};
 
-/// Trait that controls the behavior of `Vc::cell` on a value type basis.
+type VcReadTarget<T> = <<T as VcValueType>::Read as VcRead<T>>::Target;
+type VcReadRepr<T> = <<T as VcValueType>::Read as VcRead<T>>::Repr;
+
+/// Trait that controls the behavior of [`Vc::cell`] based on the value type's
+/// [`VcValueType::CellMode`].
 ///
 /// This trait must remain sealed within this crate.
 pub trait VcCellMode<T>
 where
     T: VcValueType,
 {
-    fn cell(value: <T::Read as VcRead<T>>::Target) -> Vc<T>;
+    /// Create a new cell.
+    fn cell(value: VcReadTarget<T>) -> Vc<T>;
+
+    /// Create a type-erased `RawVc` cell given a pre-existing type-erased
+    /// `SharedReference`. In some cases, we will re-use the shared value.
+    fn raw_cell(value: TypedSharedReference) -> RawVc;
 }
 
 /// Mode that always updates the cell's content.
@@ -22,14 +31,21 @@ impl<T> VcCellMode<T> for VcCellNewMode<T>
 where
     T: VcValueType,
 {
-    fn cell(inner: <T::Read as VcRead<T>>::Target) -> Vc<T> {
+    fn cell(inner: VcReadTarget<T>) -> Vc<T> {
         let cell = find_cell_by_type(T::get_value_type_id());
-        cell.update_shared(<T::Read as VcRead<T>>::target_to_repr(inner));
+        cell.update(<T::Read as VcRead<T>>::target_to_value(inner));
         Vc {
             node: cell.into(),
             _t: PhantomData,
         }
     }
+
+    fn raw_cell(content: TypedSharedReference) -> RawVc {
+        debug_assert_repr::<T>(&content);
+        let cell = find_cell_by_type(content.0);
+        cell.update_with_shared_reference(content.1);
+        cell.into()
+    }
 }
 
 /// Mode that compares the cell's content with the new value and only updates
@@ -40,15 +56,30 @@ pub struct VcCellSharedMode<T> {
 
 impl<T> VcCellMode<T> for VcCellSharedMode<T>
 where
-    T: VcValueType,
-    <T::Read as VcRead<T>>::Repr: PartialEq,
+    T: VcValueType + PartialEq,
 {
-    fn cell(inner: <T::Read as VcRead<T>>::Target) -> Vc<T> {
+    fn cell(inner: VcReadTarget<T>) -> Vc<T> {
         let cell = find_cell_by_type(T::get_value_type_id());
-        cell.compare_and_update_shared(<T::Read as VcRead<T>>::target_to_repr(inner));
+        cell.compare_and_update(<T::Read as VcRead<T>>::target_to_value(inner));
         Vc {
             node: cell.into(),
             _t: PhantomData,
         }
     }
+
+    fn raw_cell(content: TypedSharedReference) -> RawVc {
+        debug_assert_repr::<T>(&content);
+        let cell = find_cell_by_type(content.0);
+        cell.compare_and_update_with_shared_reference::<T>(content.1);
+        cell.into()
+    }
+}
+
+fn debug_assert_repr<T: VcValueType>(content: &TypedSharedReference) {
+    debug_assert!(
+        (*content.1 .0).is::<VcReadRepr<T>>(),
+        "SharedReference for type {} must use representation type {}",
+        type_name::<T>(),
+        type_name::<VcReadRepr<T>>(),
+    );
 }

crates/turbo-tasks/src/vc/mod.rs

@@ -16,10 +16,9 @@ use anyhow::Result;
 use auto_hash_map::AutoSet;
 use serde::{Deserialize, Serialize};
 
-use self::cell_mode::VcCellMode;
 pub use self::{
     cast::{VcCast, VcValueTraitCast, VcValueTypeCast},
-    cell_mode::{VcCellNewMode, VcCellSharedMode},
+    cell_mode::{VcCellMode, VcCellNewMode, VcCellSharedMode},
     default::ValueDefault,
     read::{ReadVcFuture, VcDefaultRead, VcRead, VcTransparentRead},
     resolved::{ResolvedValue, ResolvedVc},