From 5ca29c11e0b27f226727d0ce363f165192048c76 Mon Sep 17 00:00:00 2001 From: Benjamin Woodruff Date: Mon, 22 Jul 2024 16:31:41 -0700 Subject: [PATCH] Implement resolution for local Vcs --- crates/turbo-tasks-memory/tests/local_cell.rs | 55 ++++++++++++++++++- crates/turbo-tasks/src/backend.rs | 45 ++++++++++++++- crates/turbo-tasks/src/manager.rs | 21 +++++-- crates/turbo-tasks/src/persisted_graph.rs | 4 +- crates/turbo-tasks/src/raw_vc.rs | 31 ++++------- crates/turbo-tasks/src/read_ref.rs | 2 +- .../turbo-tasks/src/task/shared_reference.rs | 21 +++++-- crates/turbo-tasks/src/value_type.rs | 24 +++++++- crates/turbo-tasks/src/vc/mod.rs | 7 +-- 9 files changed, 169 insertions(+), 41 deletions(-) diff --git a/crates/turbo-tasks-memory/tests/local_cell.rs b/crates/turbo-tasks-memory/tests/local_cell.rs index e2f20b5501fc6..9274c789c6404 100644 --- a/crates/turbo-tasks-memory/tests/local_cell.rs +++ b/crates/turbo-tasks-memory/tests/local_cell.rs @@ -12,7 +12,7 @@ struct Wrapper(u32); struct TransparentWrapper(u32); #[tokio::test] -async fn store_and_read() { +async fn test_store_and_read() { run(®ISTRATION, async { let a: Vc = Vc::local_cell(42); assert_eq!(*a.await.unwrap(), 42); @@ -27,7 +27,7 @@ async fn store_and_read() { } #[tokio::test] -async fn store_and_read_generic() { +async fn test_store_and_read_generic() { run(®ISTRATION, async { // `Vc>>` is stored as `Vc>>` and requires special // transmute handling @@ -43,3 +43,54 @@ async fn store_and_read_generic() { }) .await } + +#[turbo_tasks::function] +async fn returns_resolved_local_vc() -> Vc { + Vc::::local_cell(42).resolve().await.unwrap() +} + +#[tokio::test] +async fn test_return_resolved() { + run(®ISTRATION, async { + assert_eq!(*returns_resolved_local_vc().await.unwrap(), 42); + }) + .await +} + +#[turbo_tasks::value(eq = "manual")] +#[derive(Default)] +struct Untracked { + #[turbo_tasks(debug_ignore, trace_ignore)] + #[serde(skip)] + cell: Vc, +} + +impl PartialEq for Untracked { + fn eq(&self, other: &Self) -> bool { + std::ptr::eq(self as *const _, other as *const _) + } +} + +impl Eq for Untracked {} + +#[turbo_tasks::function] +async fn get_untracked_local_cell() -> Vc { + Untracked { + cell: Vc::local_cell(42), + } + .cell() +} + +#[tokio::test] +#[should_panic(expected = "Local Vcs must only be accessed within their own task")] +async fn test_panics_on_local_cell_escape() { + run(®ISTRATION, async { + get_untracked_local_cell() + .await + .unwrap() + .cell + .await + .unwrap(); + }) + .await +} diff --git a/crates/turbo-tasks/src/backend.rs b/crates/turbo-tasks/src/backend.rs index 8377535dff840..a9ce58a13d14a 100644 --- a/crates/turbo-tasks/src/backend.rs +++ b/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 for TypedCellContent { + fn from(value: TypedSharedReference) -> Self { + TypedCellContent(value.0, CellContent(Some(value.1))) + } +} + +impl TryFrom for TypedSharedReference { + type Error = TypedCellContent; + + fn try_from(content: TypedCellContent) -> Result { + 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 for CellContent { + fn from(value: SharedReference) -> Self { + CellContent(Some(value)) + } +} + +impl From> for CellContent { + fn from(value: Option) -> Self { + CellContent(value) + } +} + +impl TryFrom for SharedReference { + type Error = CellContent; + + fn try_from(content: CellContent) -> Result { + if let CellContent(Some(shared_reference)) = content { + Ok(shared_reference) + } else { + Err(content) + } + } +} + pub type TaskCollectiblesMap = AutoMap, 1>; pub trait Backend: Sync + Send { diff --git a/crates/turbo-tasks/src/manager.rs b/crates/turbo-tasks/src/manager.rs index 73a551b20b308..4abaec6b33998 100644 --- a/crates/turbo-tasks/src/manager.rs +++ b/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 the task finishes executing. - local_cells: Vec, + local_cells: Vec, } 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, diff --git a/crates/turbo-tasks/src/persisted_graph.rs b/crates/turbo-tasks/src/persisted_graph.rs index e317dc903714b..14609f8b6071d 100644 --- a/crates/turbo-tasks/src/persisted_graph.rs +++ b/crates/turbo-tasks/src/persisted_graph.rs @@ -41,7 +41,7 @@ struct SerializableTaskCell(Option>); impl From 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, }); diff --git a/crates/turbo-tasks/src/raw_vc.rs b/crates/turbo-tasks/src/raw_vc.rs index bf8185c2fd917..9dc748bb282f4 100644 --- a/crates/turbo-tasks/src/raw_vc.rs +++ b/crates/turbo-tasks/src/raw_vc.rs @@ -164,26 +164,15 @@ impl RawVc { /// See [`crate::Vc::resolve`]. pub(crate) async fn resolve(self) -> Result { - 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 { + self.resolve_inner(/* strongly_consistent */ true).await + } + + pub(crate) async fn resolve_inner(self, strongly_consistent: bool) -> Result { 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 diff --git a/crates/turbo-tasks/src/read_ref.rs b/crates/turbo-tasks/src/read_ref.rs index b5525a29b4dc0..97f55bc2d2d74 100644 --- a/crates/turbo-tasks/src/read_ref.rs +++ b/crates/turbo-tasks/src/read_ref.rs @@ -251,7 +251,7 @@ where }; Vc { node: >::raw_cell( - SharedReference::new(value).typed(type_id), + SharedReference::new(value).into_typed(type_id), ), _t: PhantomData, } diff --git a/crates/turbo-tasks/src/task/shared_reference.rs b/crates/turbo-tasks/src/task/shared_reference.rs index 0194bfdfc597d..bb1e94dda1d03 100644 --- a/crates/turbo-tasks/src/task/shared_reference.rs +++ b/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(&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 } } diff --git a/crates/turbo-tasks/src/value_type.rs b/crates/turbo-tasks/src/value_type.rs index d9d6e0246d8c2..067df22ca9b20 100644 --- a/crates/turbo-tasks/src/value_type.rs +++ b/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`. + /// + /// Because we allow resolving `Vc`, 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( impl ValueType { /// This is internally used by `#[turbo_tasks::value]` - pub fn new() -> Self { + pub fn new() -> Self { Self { name: std::any::type_name::().to_string(), traits: AutoSet::new(), trait_methods: AutoMap::new(), magic_serialization: None, any_serialization: None, + raw_cell: >::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::().to_string(), @@ -111,12 +127,13 @@ impl ValueType { MagicAnyDeserializeSeed::new::(), )), any_serialization: Some((any_as_serialize::, AnyDeserializeSeed::new::())), + raw_cell: >::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::().to_string(), @@ -124,6 +141,7 @@ impl ValueType { trait_methods: AutoMap::new(), magic_serialization: None, any_serialization: Some((any_as_serialize::, AnyDeserializeSeed::new::())), + raw_cell: >::raw_cell, } } diff --git a/crates/turbo-tasks/src/vc/mod.rs b/crates/turbo-tasks/src/vc/mod.rs index 54ce5d040d386..bb4beaf9f349f 100644 --- a/crates/turbo-tasks/src/vc/mod.rs +++ b/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),