Implement resolution for local Vcs

crates/turbo-tasks/src/completion.rs

@@ -31,7 +31,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

@@ -1539,30 +1539,51 @@ pub(crate) async fn read_task_cell(
     }
 }
 
-#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+/// A reference to a task's cell with methods that allow updating the contents
+/// of the cell.
+///
+/// Mutations should not outside of the task that that owns this cell. Doing so
+/// is a logic error, and may lead to incorrect caching behavior.
+#[derive(Clone, Copy, Serialize, Deserialize)]
 pub struct CurrentCellRef {
     current_task: TaskId,
     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.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(
+                Some(self.index.type_id),
+                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.0.as_ref()));
         if let Some(update) = update {
+            debug_assert_eq!(
+                update.0,
+                Some(self.index.type_id),
+                "Cannot update cell with a SharedReference of a different type",
+            );
             tt.update_own_task_cell(
                 self.current_task,
                 self.index,
@@ -1574,18 +1595,54 @@ impl CurrentCellRef {
         }
     }
 
-    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.
+    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`].
+    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.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_content: T)
+    where
+        T: VcValueType,
+    {
         let tt = turbo_tasks();
         tt.update_own_task_cell(
             self.current_task,
@@ -1597,10 +1654,11 @@ impl CurrentCellRef {
         )
     }
 
-    pub fn update_shared_reference(&self, shared_ref: SharedReference) {
+    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(CellContent(Some(content))) = content {
+            // pointer equality (not value equality)
             content != shared_ref
         } else {
             true

crates/turbo-tasks/src/raw_vc.rs

@@ -173,26 +173,15 @@ impl RawVc {
 
     /// See [`crate::Vc::resolve`].
     pub(crate) async fn resolve(self) -> Result<RawVc> {
-        let tt = turbo_tasks();
-        let mut current = self;
-        let mut notified = false;
-        loop {
-            match current {
-                RawVc::TaskOutput(task) => {
-                    if !notified {
-                        tt.notify_scheduled_tasks();
-                        notified = true;
-                    }
-                    current = read_task_output(&*tt, task, false).await?;
-                }
-                RawVc::TaskCell(_, _) => return Ok(current),
-                RawVc::LocalCell(_, _) => todo!(),
-            }
-        }
+        self.resolve_inner(/* strongly_consistent */ false).await
     }
 
     /// See [`crate::Vc::resolve_strongly_consistent`].
     pub(crate) async fn resolve_strongly_consistent(self) -> Result<RawVc> {
+        self.resolve_inner(/* strongly_consistent */ true).await
+    }
+
+    pub(crate) async fn resolve_inner(self, strongly_consistent: bool) -> Result<RawVc> {
         let tt = turbo_tasks();
         let mut current = self;
         let mut notified = false;
@@ -203,10 +192,18 @@ impl RawVc {
                         tt.notify_scheduled_tasks();
                         notified = true;
                     }
-                    current = read_task_output(&*tt, task, true).await?;
+                    current = read_task_output(&*tt, task, strongly_consistent).await?;
                 }
                 RawVc::TaskCell(_, _) => return Ok(current),
-                RawVc::LocalCell(_, _) => todo!(),
+                RawVc::LocalCell(execution_id, local_cell_id) => {
+                    let shared_ref = read_local_cell(execution_id, local_cell_id);
+                    let value_type = get_value_type(
+                        shared_ref
+                            .0
+                            .expect("local cells always include ValueTypeId"),
+                    );
+                    return Ok((value_type.raw_cell)(shared_ref));
+                }
             }
         }
     }

crates/turbo-tasks/src/read_ref.rs

@@ -242,11 +242,9 @@ where
     /// Returns a new cell that points to the same value as the given
     /// reference.
     pub fn cell(read_ref: ReadRef<T>) -> Vc<T> {
-        let local_cell = find_cell_by_type(T::get_value_type_id());
-        local_cell.update_shared_reference(SharedReference::new(
-            Some(T::get_value_type_id()),
-            read_ref.0,
-        ));
+        let type_id = T::get_value_type_id();
+        let local_cell = find_cell_by_type(type_id);
+        local_cell.update_with_shared_reference(SharedReference::new(Some(type_id), read_ref.0));
         Vc {
             node: local_cell.into(),
             _t: PhantomData,

crates/turbo-tasks/src/task/shared_reference.rs

@@ -32,6 +32,10 @@ impl SharedReference {
             Err(data) => Err(Self(self.0, data)),
         }
     }
+
+    pub fn downcast_ref<T: Any>(&self) -> Option<&T> {
+        self.1.downcast_ref()
+    }
 }
 
 impl Hash for SharedReference {

crates/turbo-tasks/src/trait_ref.rs

@@ -102,7 +102,7 @@ where
         let SharedReference(ty, _) = trait_ref.shared_reference;
         let ty = ty.unwrap();
         let local_cell = find_cell_by_type(ty);
-        local_cell.update_shared_reference(trait_ref.shared_reference);
+        local_cell.update_with_shared_reference(trait_ref.shared_reference);
         let raw_vc: RawVc = local_cell.into();
         raw_vc.into()
     }

crates/turbo-tasks/src/value_type.rs

@@ -16,10 +16,13 @@ use crate::{
     id::{FunctionId, TraitTypeId},
     magic_any::{AnyDeserializeSeed, MagicAny, MagicAnyDeserializeSeed, MagicAnySerializeSeed},
     registry::{register_trait_type, register_value_type},
+    vc::VcCellMode,
+    RawVc, SharedReference, VcValueType,
 };
 
 type MagicSerializationFn = fn(&dyn MagicAny) -> &dyn erased_serde::Serialize;
 type AnySerializationFn = fn(&(dyn Any + Sync + Send)) -> &dyn erased_serde::Serialize;
+type RawCellFactoryFn = fn(SharedReference) -> RawVc;
 
 // TODO this type need some refactoring when multiple languages are added to
 // turbo-task In this case a trait_method might be of a different function type.
@@ -41,6 +44,17 @@ pub struct ValueType {
     /// Functors for serialization
     magic_serialization: Option<(MagicSerializationFn, MagicAnyDeserializeSeed)>,
     any_serialization: Option<(AnySerializationFn, AnyDeserializeSeed)>,
+
+    /// An implementation of
+    /// [`VcCellMode::raw_cell`][crate::vc::cell_mode::VcCellMode::raw_cell].
+    ///
+    /// Allows dynamically constructing a cell using the type id. Used inside of
+    /// [`RawVc`] where we have a type id, but not the concrete type `T` of
+    /// `Vc<T>`.
+    ///
+    /// Because we allow resolving `Vc<dyn Trait>`, it's otherwise not possible
+    /// for `RawVc` to know what the appropriate `VcCellMode` is.
+    pub(crate) raw_cell: RawCellFactoryFn,
 }
 
 impl Hash for ValueType {
@@ -88,19 +102,20 @@ pub fn any_as_serialize<T: Any + Serialize + Send + Sync + 'static>(
 
 impl ValueType {
     /// This is internally used by `#[turbo_tasks::value]`
-    pub fn new<T>() -> Self {
+    pub fn new<T: VcValueType>() -> Self {
         Self {
             name: std::any::type_name::<T>().to_string(),
             traits: AutoSet::new(),
             trait_methods: AutoMap::new(),
             magic_serialization: None,
             any_serialization: None,
+            raw_cell: <T::CellMode as VcCellMode<T>>::raw_cell,
         }
     }
 
     /// This is internally used by `#[turbo_tasks::value]`
     pub fn new_with_magic_serialization<
-        T: Debug + Eq + Hash + Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static,
+        T: VcValueType + Debug + Eq + Hash + Serialize + for<'de> Deserialize<'de>,
     >() -> Self {
         Self {
             name: std::any::type_name::<T>().to_string(),
@@ -111,19 +126,21 @@ impl ValueType {
                 MagicAnyDeserializeSeed::new::<T>(),
             )),
             any_serialization: Some((any_as_serialize::<T>, AnyDeserializeSeed::new::<T>())),
+            raw_cell: <T::CellMode as VcCellMode<T>>::raw_cell,
         }
     }
 
     /// This is internally used by `#[turbo_tasks::value]`
     pub fn new_with_any_serialization<
-        T: Any + Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static,
+        T: VcValueType + Any + Serialize + for<'de> Deserialize<'de>,
     >() -> Self {
         Self {
             name: std::any::type_name::<T>().to_string(),
             traits: AutoSet::new(),
             trait_methods: AutoMap::new(),
             magic_serialization: None,
             any_serialization: Some((any_as_serialize::<T>, AnyDeserializeSeed::new::<T>())),
+            raw_cell: <T::CellMode as VcCellMode<T>>::raw_cell,
         }
     }
 

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

@@ -1,16 +1,25 @@
 use std::marker::PhantomData;
 
 use super::{read::VcRead, traits::VcValueType};
-use crate::{manager::find_cell_by_type, Vc};
+use crate::{manager::find_cell_by_type, RawVc, SharedReference, 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: SharedReference) -> RawVc;
 }
 
 /// Mode that always updates the cell's content.
@@ -22,14 +31,20 @@ 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_repr(inner));
         Vc {
             node: cell.into(),
             _t: PhantomData,
         }
     }
+
+    fn raw_cell(value: SharedReference) -> RawVc {
+        let cell = find_cell_by_type(T::get_value_type_id());
+        cell.update_with_shared_reference(value);
+        cell.into()
+    }
 }
 
 /// Mode that compares the cell's content with the new value and only updates
@@ -43,12 +58,18 @@ where
     T: VcValueType,
     <T::Read as VcRead<T>>::Repr: 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_repr(inner));
         Vc {
             node: cell.into(),
             _t: PhantomData,
         }
     }
+
+    fn raw_cell(value: SharedReference) -> RawVc {
+        let cell = find_cell_by_type(T::get_value_type_id());
+        cell.compare_and_update_with_shared_reference::<VcReadRepr<T>>(value);
+        cell.into()
+    }
 }

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},