Chunking Refactoring

crates/turbo-tasks/src/keyed_cell.rs

@@ -0,0 +1,55 @@
+use anyhow::Result;
+
+use crate::{
+    self as turbo_tasks, macro_helpers::find_cell_by_type, manager::current_task,
+    ConcreteTaskInput, CurrentCellRef, RawVc, TaskId, TaskInput, ValueTypeId, Vc, VcValueType,
+};
+
+#[turbo_tasks::value]
+struct KeyedCell {
+    cell: RawVc,
+    #[turbo_tasks(trace_ignore, debug_ignore)]
+    cell_ref: CurrentCellRef,
+}
+
+impl KeyedCell {}
+
+#[turbo_tasks::value_impl]
+impl KeyedCell {
+    #[turbo_tasks::function]
+    fn new(_task: TaskId, _key: ConcreteTaskInput, value_type_id: ValueTypeId) -> Vc<Self> {
+        let cell_ref = find_cell_by_type(value_type_id);
+        KeyedCell {
+            cell: cell_ref.into(),
+            cell_ref,
+        }
+        .cell()
+    }
+}
+
+/// Cells a value in a cell with a given key. A key MUST only be used once per
+/// function.
+///
+/// Usually calling [Vc::cell] will create cells for a give type based on the
+/// call order of [Vc::cell]. But this can yield to over-invalidation when the
+/// number of cells changes. e. g. not doing the first [Vc::cell] call will move
+/// all remaining values into different cells, causing invalidation of all of
+/// them.
+///
+/// A keyed cell avoids this problem by not using call order, but a key instead.
+///
+/// Internally it creates a new Task based on the key and cells the value into
+/// that task. This is a implementation detail and might change in the future.
+pub async fn keyed_cell<T: PartialEq + Eq + VcValueType, K: TaskInput>(
+    key: K,
+    content: T,
+) -> Result<Vc<T>> {
+    let cell = KeyedCell::new(
+        current_task("keyed_cell"),
+        key.into_concrete(),
+        T::get_value_type_id(),
+    )
+    .await?;
+    cell.cell_ref.compare_and_update_shared(content);
+    Ok(cell.cell.into())
+}

crates/turbo-tasks/src/lib.rs

@@ -48,6 +48,7 @@ mod id;
 mod id_factory;
 mod invalidation;
 mod join_iter_ext;
+mod keyed_cell;
 #[doc(hidden)]
 pub mod macro_helpers;
 mod magic_any;
@@ -85,11 +86,12 @@ pub use invalidation::{
     DynamicEqHash, InvalidationReason, InvalidationReasonKind, InvalidationReasonSet,
 };
 pub use join_iter_ext::{JoinIterExt, TryFlatJoinIterExt, TryJoinIterExt};
+pub use keyed_cell::keyed_cell;
 pub use manager::{
     dynamic_call, emit, get_invalidator, mark_finished, mark_stateful, run_once,
-    run_once_with_reason, spawn_blocking, spawn_thread, trait_call, turbo_tasks, Invalidator,
-    StatsType, TaskIdProvider, TurboTasks, TurboTasksApi, TurboTasksBackendApi, TurboTasksCallApi,
-    Unused, UpdateInfo,
+    run_once_with_reason, spawn_blocking, spawn_thread, trait_call, turbo_tasks, CurrentCellRef,
+    Invalidator, StatsType, TaskIdProvider, TurboTasks, TurboTasksApi, TurboTasksBackendApi,
+    TurboTasksCallApi, Unused, UpdateInfo,
 };
 pub use native_function::NativeFunction;
 use nohash_hasher::BuildNoHashHasher;

crates/turbo-tasks/src/manager.rs

@@ -1159,7 +1159,7 @@ impl<B: Backend + 'static> TaskIdProvider for TurboTasks<B> {
     }
 }
 
-fn current_task(from: &str) -> TaskId {
+pub(crate) fn current_task(from: &str) -> TaskId {
     match CURRENT_TASK_ID.try_with(|id| *id) {
         Ok(id) => id,
         Err(_) => panic!(
@@ -1478,7 +1478,7 @@ pub(crate) async fn read_task_cell(
     }
 }
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
 pub struct CurrentCellRef {
     current_task: TaskId,
     index: CellId,

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

@@ -8,8 +8,8 @@ use anyhow::{anyhow, bail, Result};
 
 use super::concrete_task_input::TransientSharedValue;
 use crate::{
-    magic_any::MagicAny, ConcreteTaskInput, RawVc, SharedValue, TransientInstance, TransientValue,
-    TypedForInput, Value, Vc, VcValueType,
+    magic_any::MagicAny, ConcreteTaskInput, RawVc, SharedValue, TaskId, TransientInstance,
+    TransientValue, TypedForInput, Value, ValueTypeId, Vc, VcValueType,
 };
 
 /// Trait to implement in order for a type to be accepted as a
@@ -21,6 +21,16 @@ pub trait TaskInput: Send + Sync + Clone {
     fn into_concrete(self) -> ConcreteTaskInput;
 }
 
+impl TaskInput for ConcreteTaskInput {
+    fn try_from_concrete(input: &ConcreteTaskInput) -> Result<Self> {
+        Ok(input.clone())
+    }
+
+    fn into_concrete(self) -> ConcreteTaskInput {
+        self
+    }
+}
+
 impl TaskInput for String {
     fn try_from_concrete(input: &ConcreteTaskInput) -> Result<Self> {
         match input {
@@ -148,6 +158,32 @@ impl TaskInput for usize {
     }
 }
 
+impl TaskInput for ValueTypeId {
+    fn try_from_concrete(value: &ConcreteTaskInput) -> Result<Self> {
+        match value {
+            ConcreteTaskInput::Usize(value) => Ok(ValueTypeId::from(*value)),
+            _ => bail!("invalid task input type, expected ValueTypeId"),
+        }
+    }
+
+    fn into_concrete(self) -> ConcreteTaskInput {
+        ConcreteTaskInput::Usize(*self)
+    }
+}
+
+impl TaskInput for TaskId {
+    fn try_from_concrete(value: &ConcreteTaskInput) -> Result<Self> {
+        match value {
+            ConcreteTaskInput::Usize(value) => Ok(TaskId::from(*value)),
+            _ => bail!("invalid task input type, expected TaskId"),
+        }
+    }
+
+    fn into_concrete(self) -> ConcreteTaskInput {
+        ConcreteTaskInput::Usize(*self)
+    }
+}
+
 impl<T> TaskInput for Option<T>
 where
     T: TaskInput,