Implement resolution for local Vcs

crates/turbo-tasks/src/backend.rs

@@ -20,6 +20,7 @@ use crate::{
     manager::TurboTasksBackendApi,
     raw_vc::CellId,
     registry,
+    task::shared_reference::TypedSharedReference,
     trait_helpers::{get_trait_method, has_trait, traits},
     triomphe_utils::unchecked_sidecast_triomphe_arc,
     FunctionId, RawVc, ReadRef, SharedReference, TaskId, TaskIdProvider, TaskIdSet, TraitRef,
@@ -370,7 +371,7 @@ impl TypedCellContent {
             .1
              .0
             .ok_or_else(|| anyhow!("Cell is empty"))?
-            .typed(self.0);
+            .into_typed(self.0);
         Ok(
             // Safety: It is a TypedSharedReference
             TraitRef::new(shared_reference),
@@ -382,12 +383,54 @@ impl TypedCellContent {
     }
 }
 
+impl From<TypedSharedReference> for TypedCellContent {
+    fn from(value: TypedSharedReference) -> Self {
+        TypedCellContent(value.0, CellContent(Some(value.1)))
+    }
+}
+
+impl TryFrom<TypedCellContent> for TypedSharedReference {
+    type Error = TypedCellContent;
+
+    fn try_from(content: TypedCellContent) -> Result<Self, TypedCellContent> {
+        if let TypedCellContent(type_id, CellContent(Some(shared_reference))) = content {
+            Ok(TypedSharedReference(type_id, shared_reference))
+        } else {
+            Err(content)
+        }
+    }
+}
+
 impl CellContent {
     pub fn into_typed(self, type_id: ValueTypeId) -> TypedCellContent {
         TypedCellContent(type_id, self)
     }
 }
 
+impl From<SharedReference> for CellContent {
+    fn from(value: SharedReference) -> Self {
+        CellContent(Some(value))
+    }
+}
+
+impl From<Option<SharedReference>> for CellContent {
+    fn from(value: Option<SharedReference>) -> Self {
+        CellContent(value)
+    }
+}
+
+impl TryFrom<CellContent> for SharedReference {
+    type Error = CellContent;
+
+    fn try_from(content: CellContent) -> Result<Self, CellContent> {
+        if let CellContent(Some(shared_reference)) = content {
+            Ok(shared_reference)
+        } else {
+            Err(content)
+        }
+    }
+}
+
 pub type TaskCollectiblesMap = AutoMap<RawVc, i32, BuildHasherDefault<FxHasher>, 1>;
 
 pub trait Backend: Sync + Send {

crates/turbo-tasks/src/manager.rs

@@ -35,6 +35,7 @@ use crate::{
     magic_any::MagicAny,
     raw_vc::{CellId, RawVc},
     registry,
+    task::shared_reference::TypedSharedReference,
     trace::TraceRawVcs,
     trait_helpers::get_trait_method,
     util::StaticOrArc,
@@ -274,7 +275,7 @@ struct CurrentTaskState {
 
     /// Cells for locally allocated Vcs (`RawVc::LocalCell`). This is freed
     /// (along with `CurrentTaskState`) when
-    local_cells: Vec<TypedCellContent>,
+    local_cells: Vec<TypedSharedReference>,
 }
 
 impl CurrentTaskState {
@@ -1542,7 +1543,12 @@ 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,
@@ -1715,7 +1721,7 @@ pub fn find_cell_by_type(ty: ValueTypeId) -> CurrentCellRef {
     })
 }
 
-pub(crate) fn create_local_cell(value: TypedCellContent) -> (ExecutionId, LocalCellId) {
+pub(crate) fn create_local_cell(value: TypedSharedReference) -> (ExecutionId, LocalCellId) {
     CURRENT_TASK_STATE.with(|cell| {
         let CurrentTaskState {
             execution_id,
@@ -1738,11 +1744,18 @@ pub(crate) fn create_local_cell(value: TypedCellContent) -> (ExecutionId, LocalC
     })
 }
 
+/// Returns the contents of the given local cell. Panics if a local cell is
+/// attempted to be accessed outside of its task.
+///
+/// Returns [`TypedSharedReference`] instead of [`TypedCellContent`] because
+/// local cells are always filled. The returned value can be cheaply converted
+/// with `.into()`.
+///
 /// Panics if the ExecutionId does not match the expected value.
 pub(crate) fn read_local_cell(
     execution_id: ExecutionId,
     local_cell_id: LocalCellId,
-) -> TypedCellContent {
+) -> TypedSharedReference {
     CURRENT_TASK_STATE.with(|cell| {
         let CurrentTaskState {
             execution_id: expected_execution_id,

crates/turbo-tasks/src/persisted_graph.rs

@@ -41,7 +41,7 @@ struct SerializableTaskCell(Option<Option<TypedSharedReference>>);
 impl From<SerializableTaskCell> for TaskCell {
     fn from(val: SerializableTaskCell) -> Self {
         match val.0 {
-            Some(d) => TaskCell::Content(CellContent(d.map(|d| d.untyped().1))),
+            Some(d) => TaskCell::Content(d.map(TypedSharedReference::into_untyped).into()),
             None => TaskCell::NeedComputation,
         }
     }
@@ -56,7 +56,7 @@ impl Serialize for TaskCells {
         for (cell_id, cell) in &self.0 {
             let task_cell = SerializableTaskCell(match cell {
                 TaskCell::Content(CellContent(opt)) => {
-                    Some(opt.as_ref().map(|d| d.typed(cell_id.type_id)))
+                    Some(opt.clone().map(|d| d.into_typed(cell_id.type_id)))
                 }
                 TaskCell::NeedComputation => None,
             });

crates/turbo-tasks/src/raw_vc.rs

@@ -164,26 +164,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;
@@ -194,10 +183,14 @@ 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_reference = read_local_cell(execution_id, local_cell_id);
+                    let value_type = get_value_type(shared_reference.0);
+                    return Ok((value_type.raw_cell)(shared_reference));
+                }
             }
         }
     }
@@ -355,7 +348,7 @@ impl Future for ReadRawVcFuture {
                     }
                 }
                 RawVc::LocalCell(execution_id, local_cell_id) => {
-                    return Poll::Ready(Ok(read_local_cell(execution_id, local_cell_id)));
+                    return Poll::Ready(Ok(read_local_cell(execution_id, local_cell_id).into()));
                 }
             };
             // SAFETY: listener is from previous pinned this

crates/turbo-tasks/src/read_ref.rs

@@ -251,7 +251,7 @@ where
         };
         Vc {
             node: <T::CellMode as VcCellMode<T>>::raw_cell(
-                SharedReference::new(value).typed(type_id),
+                SharedReference::new(value).into_typed(type_id),
             ),
             _t: PhantomData,
         }

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

@@ -2,6 +2,7 @@ use std::{
     any::Any,
     fmt::{Debug, Display},
     hash::Hash,
+    ops::Deref,
 };
 
 use anyhow::Result;
@@ -36,14 +37,26 @@ impl SharedReference {
         }
     }
 
-    pub(crate) fn typed(&self, type_id: ValueTypeId) -> TypedSharedReference {
-        TypedSharedReference(type_id, self.clone())
+    pub fn downcast_ref<T: Any>(&self) -> Option<&T> {
+        self.0.downcast_ref()
+    }
+
+    pub fn into_typed(self, type_id: ValueTypeId) -> TypedSharedReference {
+        TypedSharedReference(type_id, self)
     }
 }
 
 impl TypedSharedReference {
-    pub(crate) fn untyped(&self) -> (ValueTypeId, SharedReference) {
-        (self.0, self.1.clone())
+    pub fn into_untyped(self) -> SharedReference {
+        self.1
+    }
+}
+
+impl Deref for TypedSharedReference {
+    type Target = SharedReference;
+
+    fn deref(&self) -> &Self::Target {
+        &self.1
     }
 }
 

crates/turbo-tasks/src/value_type.rs

@@ -16,10 +16,14 @@ use crate::{
     id::{FunctionId, TraitTypeId},
     magic_any::{AnyDeserializeSeed, MagicAny, MagicAnyDeserializeSeed, MagicAnySerializeSeed},
     registry::{register_trait_type, register_value_type},
+    task::shared_reference::TypedSharedReference,
+    vc::VcCellMode,
+    RawVc, VcValueType,
 };
 
 type MagicSerializationFn = fn(&dyn MagicAny) -> &dyn erased_serde::Serialize;
 type AnySerializationFn = fn(&(dyn Any + Sync + Send)) -> &dyn erased_serde::Serialize;
+type RawCellFactoryFn = fn(TypedSharedReference) -> 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 +45,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 +103,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 +127,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/mod.rs

@@ -25,7 +25,6 @@ pub use self::{
     traits::{Dynamic, TypedForInput, Upcast, VcValueTrait, VcValueType},
 };
 use crate::{
-    backend::CellContent,
     debug::{ValueDebug, ValueDebugFormat, ValueDebugFormatString},
     manager::create_local_cell,
     registry,
@@ -283,10 +282,8 @@ where
         // cells aren't stored across executions, so there can be no concept of
         // "updating" the cell across multiple executions.
         let (execution_id, local_cell_id) = create_local_cell(
-            CellContent(Some(SharedReference::new(triomphe::Arc::new(
-                T::Read::target_to_repr(inner),
-            ))))
-            .into_typed(T::get_value_type_id()),
+            SharedReference::new(triomphe::Arc::new(T::Read::target_to_repr(inner)))
+                .into_typed(T::get_value_type_id()),
         );
         Vc {
             node: RawVc::LocalCell(execution_id, local_cell_id),