turbo-tasks: task-storage memory wins

turbopack/crates/turbo-tasks-backend/src/backend/counter_map.rs

@@ -12,34 +12,37 @@ use bincode::{
 };
 use rustc_hash::FxHasher;
 
-type InnerMap<K, V> = AutoMap<K, V, BuildHasherDefault<FxHasher>, 1>;
+type InnerMap<K, V, const I: usize> = AutoMap<K, V, BuildHasherDefault<FxHasher>, I>;
 
 /// A map optimized for reference counting, backed by AutoMap.
 ///
 /// Entries are automatically removed when their count reaches zero.
 /// This provides memory-efficient storage for sparse counter data.
+///
+/// The `I` const generic forwards the inline capacity to the backing `AutoMap`
+/// — see the schema field-by-field sizing for the chosen values.
 #[derive(Debug, Clone)]
-pub struct CounterMap<K, V>(InnerMap<K, V>);
+pub struct CounterMap<K, V, const I: usize>(InnerMap<K, V, I>);
 
-impl<K, V> Default for CounterMap<K, V> {
+impl<K, V, const I: usize> Default for CounterMap<K, V, I> {
     fn default() -> Self {
         Self(InnerMap::default())
     }
 }
 
-impl<K: Eq + Hash, V: Eq> PartialEq for CounterMap<K, V> {
+impl<K: Eq + Hash, V: Eq, const I: usize> PartialEq for CounterMap<K, V, I> {
     fn eq(&self, other: &Self) -> bool {
         self.0 == other.0
     }
 }
 
-impl<K: Encode, V: Encode> Encode for CounterMap<K, V> {
+impl<K: Encode, V: Encode, const I: usize> Encode for CounterMap<K, V, I> {
     fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
         self.0.encode(encoder)
     }
 }
 
-impl<Context, K, V> Decode<Context> for CounterMap<K, V>
+impl<Context, K, V, const I: usize> Decode<Context> for CounterMap<K, V, I>
 where
     K: Decode<Context> + Eq + Hash,
     V: Decode<Context>,
@@ -80,7 +83,7 @@ impl CounterValue for i32 {
     }
 }
 
-impl<K, V> CounterMap<K, V> {
+impl<K, V, const I: usize> CounterMap<K, V, I> {
     pub fn new() -> Self {
         Self(AutoMap::default())
     }
@@ -138,16 +141,16 @@ impl<K, V> CounterMap<K, V> {
     }
 }
 
-impl<K, V> IntoIterator for CounterMap<K, V> {
+impl<K, V, const I: usize> IntoIterator for CounterMap<K, V, I> {
     type Item = (K, V);
-    type IntoIter = <InnerMap<K, V> as IntoIterator>::IntoIter;
+    type IntoIter = <InnerMap<K, V, I> as IntoIterator>::IntoIter;
 
     fn into_iter(self) -> Self::IntoIter {
         self.0.into_iter()
     }
 }
 
-impl<K: Hash + Eq, V: CounterValue> CounterMap<K, V> {
+impl<K: Hash + Eq, V: CounterValue, const I: usize> CounterMap<K, V, I> {
     /// Insert a key-value pair. Panics if value is zero (invariant: zero values are not stored).
     pub fn insert(&mut self, key: K, value: V) -> Option<V> {
         debug_assert!(
@@ -297,15 +300,15 @@ mod tests {
 
     #[test]
     fn test_update_count_new_entry() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         // Adding new entry crosses zero (from nothing to something)
         assert!(map.update_count(1, 5));
         assert_eq!(map.get(&1), Some(&5));
     }
 
     #[test]
     fn test_update_count_increment() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.update_count(1, 5);
         // Incrementing existing entry doesn't cross zero
         assert!(!map.update_count(1, 3));
@@ -314,7 +317,7 @@ mod tests {
 
     #[test]
     fn test_update_count_removal_on_zero() {
-        let mut map: CounterMap<u32, i32> = CounterMap::new();
+        let mut map: CounterMap<u32, i32, 1> = CounterMap::new();
         map.update_count(1, 5);
         // Subtracting to zero removes entry and crosses zero
         assert!(map.update_count(1, -5));
@@ -324,69 +327,69 @@ mod tests {
 
     #[test]
     fn test_update_count_zero_delta_on_empty() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         // Adding zero to non-existent entry doesn't create it
         assert!(!map.update_count(1, 0));
         assert!(map.is_empty());
     }
 
     #[test]
     fn test_update_and_get_new_entry() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         assert_eq!(map.update_and_get(1, 5), 5);
         assert_eq!(map.get(&1), Some(&5));
     }
 
     #[test]
     fn test_update_and_get_increment() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.update_and_get(1, 5);
         assert_eq!(map.update_and_get(1, 3), 8);
         assert_eq!(map.get(&1), Some(&8));
     }
 
     #[test]
     fn test_update_and_get_removal() {
-        let mut map: CounterMap<u32, i32> = CounterMap::new();
+        let mut map: CounterMap<u32, i32, 1> = CounterMap::new();
         map.update_and_get(1, 5);
         assert_eq!(map.update_and_get(1, -5), 0);
         assert!(map.is_empty());
     }
 
     #[test]
     fn test_add_entry() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.add_entry(1, 10);
         assert_eq!(map.get(&1), Some(&10));
     }
 
     #[test]
     #[should_panic(expected = "Entry already exists")]
     fn test_add_entry_panics_on_duplicate() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.add_entry(1, 10);
         map.add_entry(1, 20); // Should panic
     }
 
     #[test]
     fn test_update_positive_crossing_new_positive() {
-        let mut map: CounterMap<u32, i32> = CounterMap::new();
+        let mut map: CounterMap<u32, i32, 1> = CounterMap::new();
         // From nothing to positive - crosses positive boundary
         assert!(map.update_positive_crossing(1, 5));
         assert_eq!(map.get(&1), Some(&5));
     }
 
     #[test]
     fn test_update_positive_crossing_new_negative() {
-        let mut map: CounterMap<u32, i32> = CounterMap::new();
+        let mut map: CounterMap<u32, i32, 1> = CounterMap::new();
         // From nothing to negative - doesn't cross positive boundary
         assert!(!map.update_positive_crossing(1, -5));
         assert_eq!(map.get(&1), Some(&-5));
     }
 
     #[test]
     fn test_update_positive_crossing_stay_positive() {
-        let mut map: CounterMap<u32, i32> = CounterMap::new();
+        let mut map: CounterMap<u32, i32, 1> = CounterMap::new();
         map.update_positive_crossing(1, 5);
         // Staying positive doesn't cross boundary
         assert!(!map.update_positive_crossing(1, 3));
@@ -395,7 +398,7 @@ mod tests {
 
     #[test]
     fn test_update_positive_crossing_to_non_positive() {
-        let mut map: CounterMap<u32, i32> = CounterMap::new();
+        let mut map: CounterMap<u32, i32, 1> = CounterMap::new();
         map.update_positive_crossing(1, 5);
         // Crossing to non-positive
         assert!(map.update_positive_crossing(1, -8));
@@ -404,7 +407,7 @@ mod tests {
 
     #[test]
     fn test_update_positive_crossing_to_zero_removes() {
-        let mut map: CounterMap<u32, i32> = CounterMap::new();
+        let mut map: CounterMap<u32, i32, 1> = CounterMap::new();
         map.update_positive_crossing(1, 5);
         // Crossing to zero removes and crosses boundary
         assert!(map.update_positive_crossing(1, -5));
@@ -413,37 +416,37 @@ mod tests {
 
     #[test]
     fn test_update_with_create() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.update_with(1, |_| Some(10));
         assert_eq!(map.get(&1), Some(&10));
     }
 
     #[test]
     fn test_update_with_modify() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.update_with(1, |_| Some(10));
         map.update_with(1, |v| v.map(|x| x + 5));
         assert_eq!(map.get(&1), Some(&15));
     }
 
     #[test]
     fn test_update_with_remove() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.update_with(1, |_| Some(10));
         map.update_with(1, |_| None);
         assert!(map.is_empty());
     }
 
     #[test]
     fn test_update_with_no_op() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.update_with(1, |_| None);
         assert!(map.is_empty());
     }
 
     #[test]
     fn test_len_and_is_empty() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         assert!(map.is_empty());
         assert_eq!(map.len(), 0);
 
@@ -457,7 +460,7 @@ mod tests {
 
     #[test]
     fn test_iter() {
-        let mut map: CounterMap<u32, u32> = CounterMap::new();
+        let mut map: CounterMap<u32, u32, 1> = CounterMap::new();
         map.update_count(1, 5);
         map.update_count(2, 10);
 

turbopack/crates/turbo-tasks-backend/src/backend/mod.rs

@@ -34,8 +34,8 @@ use turbo_tasks::{
     TaskId, TaskPersistence, TaskPriority, TraitTypeId, TurboTasksBackendApi, TurboTasksPanic,
     ValueTypeId,
     backend::{
-        Backend, CachedTaskType, CellContent, CellHash, TaskExecutionSpec, TransientTaskType,
-        TurboTaskContextError, TurboTaskLocalContextError, TurboTasksError,
+        Backend, CachedTaskType, CachedTaskTypeArc, CellContent, CellHash, TaskExecutionSpec,
+        TransientTaskType, TurboTaskContextError, TurboTaskLocalContextError, TurboTasksError,
         TurboTasksExecutionError, TurboTasksExecutionErrorMessage, TypedCellContent,
         VerificationMode,
     },
@@ -70,8 +70,8 @@ use crate::{
     },
     backing_storage::{BackingStorage, SnapshotItem, compute_task_type_hash},
     data::{
-        ActivenessState, CellRef, CollectibleRef, CollectiblesRef, Dirtyness, InProgressCellState,
-        InProgressState, InProgressStateInner, OutputValue, TransientTask,
+        ActivenessState, CellDependency, CellRef, CollectibleRef, CollectiblesRef, Dirtyness,
+        InProgressCellState, InProgressState, InProgressStateInner, OutputValue, TransientTask,
     },
     error::TaskError,
     utils::{
@@ -785,7 +785,8 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                 && (!task.immutable() || cfg!(feature = "verify_immutable"))
             {
                 let reader = reader.unwrap();
-                let _ = task.add_cell_dependents((cell, key, reader));
+                let _ = task
+                    .add_cell_dependents(CellDependency::new(CellRef { task: reader, cell }, key));
                 drop(task);
 
                 // Note: We use `task_pair` earlier to lock the task and its reader at the same
@@ -797,8 +798,9 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                     task: task_id,
                     cell,
                 };
-                if !reader_task.remove_outdated_cell_dependencies(&(target, key)) {
-                    let _ = reader_task.add_cell_dependencies((target, key));
+                let dep = CellDependency::new(target, key);
+                if !reader_task.remove_outdated_cell_dependencies(&dep) {
+                    let _ = reader_task.add_cell_dependencies(dep);
                 }
                 drop(reader_task);
             }
@@ -1526,7 +1528,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                     // Only now do we force the allocation.
                     // NOTE: if our caller had to perform resolution, then this will have already
                     // been boxed and take_box just takes it.
-                    let task_type = Arc::new(CachedTaskType {
+                    let task_type = CachedTaskTypeArc::new(CachedTaskType {
                         native_fn,
                         this,
                         arg: arg.take_box(),
@@ -1757,7 +1759,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         }
     }
 
-    fn debug_get_cached_task_type(&self, task_id: TaskId) -> Option<Arc<CachedTaskType>> {
+    fn debug_get_cached_task_type(&self, task_id: TaskId) -> Option<CachedTaskTypeArc> {
         let task = self.storage.access_mut(task_id);
         task.get_persistent_task_type().cloned()
     }
@@ -2197,7 +2199,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             Some(
                 // Collect all dependencies on tasks to check if all dependencies are immutable
                 task.iter_output_dependencies()
-                    .chain(task.iter_cell_dependencies().map(|(target, _key)| target.task))
+                    .chain(task.iter_cell_dependencies().map(|dep| dep.cell_ref().task))
                     .collect::<FxHashSet<_>>(),
             )
         } else {
@@ -2236,7 +2238,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             // breaking dependency tracking.
             old_edges.extend(
                 task.iter_outdated_cell_dependencies()
-                    .map(|(target, key)| OutdatedEdge::CellDependency(target, key)),
+                    .map(OutdatedEdge::CellDependency),
             );
             old_edges.extend(
                 task.iter_outdated_output_dependencies()

turbopack/crates/turbo-tasks-backend/src/backend/operation/cleanup_old_edges.rs

@@ -17,7 +17,7 @@ use crate::{
         },
         storage_schema::TaskStorageAccessors,
     },
-    data::{CellRef, CollectibleRef, CollectiblesRef},
+    data::{CellDependency, CellRef, CollectibleRef, CollectiblesRef},
 };
 
 #[derive(Encode, Decode, Clone)]
@@ -48,7 +48,7 @@ impl Default for CleanupOldEdgesOperation {
 pub enum OutdatedEdge {
     Child(TaskId),
     Collectible(CollectibleRef, i32),
-    CellDependency(CellRef, Option<u64>),
+    CellDependency(CellDependency),
     OutputDependency(TaskId),
     CollectiblesDependency(CollectiblesRef),
 }
@@ -166,27 +166,28 @@ impl CleanupOldEdgesOperation {
                                     AggregatedDataUpdate::new().collectibles_update(collectibles),
                                 ));
                             }
-                            OutdatedEdge::CellDependency(
-                                CellRef {
-                                    task: cell_task_id,
-                                    cell,
-                                },
-                                key,
-                            ) => {
+                            OutdatedEdge::CellDependency(dep) => {
+                                let (
+                                    CellRef {
+                                        task: cell_task_id,
+                                        cell,
+                                    },
+                                    key,
+                                ) = dep.into_parts();
                                 {
                                     let mut task = ctx.task(cell_task_id, TaskDataCategory::Data);
-                                    task.remove_cell_dependents(&(cell, key, task_id));
-                                }
-                                {
-                                    let mut task = ctx.task(task_id, TaskDataCategory::Data);
-                                    task.remove_cell_dependencies(&(
+                                    task.remove_cell_dependents(&CellDependency::new(
                                         CellRef {
-                                            task: cell_task_id,
+                                            task: task_id,
                                             cell,
                                         },
                                         key,
                                     ));
                                 }
+                                {
+                                    let mut task = ctx.task(task_id, TaskDataCategory::Data);
+                                    task.remove_cell_dependencies(&dep);
+                                }
                             }
                             OutdatedEdge::OutputDependency(output_task_id) => {
                                 #[cfg(feature = "trace_task_output_dependencies")]