keep children count instead of iterating children every time

Author: sokra

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

@@ -13,7 +13,7 @@ use crate::{
             invalidate::make_task_dirty,
             AggregatedDataUpdate, ExecuteContext, Operation,
         },
-        storage::update_count,
+        storage::{update, update_count},
         TaskDataCategory,
     },
     data::{CachedDataItemKey, CellRef, CollectibleRef, CollectiblesRef},
@@ -82,6 +82,11 @@ impl Operation for CleanupOldEdgesOperation {
                                 for &child_id in children.iter() {
                                     task.remove(&CachedDataItemKey::Child { task: child_id });
                                 }
+                                update!(task, ChildrenCount, |count: Option<u32>| {
+                                    // If this underflows, we messed up counting somewhere
+                                    let count = count.unwrap_or_default() - children.len() as u32;
+                                    (count != 0).then_some(count)
+                                });
                                 if is_aggregating_node(get_aggregation_number(&task)) {
                                     queue.push(AggregationUpdateJob::InnerLostFollowers {
                                         upper_ids: vec![task_id],

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

@@ -11,10 +11,10 @@ use crate::{
             },
             is_root_node, ExecuteContext, Operation,
         },
-        storage::get,
+        storage::{get, update},
         TaskDataCategory,
     },
-    data::{CachedDataItem, CachedDataItemIndex, CachedDataItemKey},
+    data::{CachedDataItem, CachedDataItemKey},
 };
 
 const AGGREGATION_NUMBER_BUFFER_SPACE: u32 = 2;
@@ -44,6 +44,13 @@ impl ConnectChildOperation {
             task: child_task_id,
             value: (),
         }) {
+            // Update the children count
+            let mut children_count = 0;
+            update!(parent_task, ChildrenCount, |count: Option<u32>| {
+                children_count = count.unwrap_or_default() + 1;
+                Some(children_count)
+            });
+
             // Update the task aggregation
             let mut queue = AggregationUpdateQueue::new();
 
@@ -54,16 +61,6 @@ impl ConnectChildOperation {
             let parent_aggregation = if is_root_node(current_parent_aggregation.base) {
                 u32::MAX
             } else {
-                let children_count = parent_task
-                    .iter(CachedDataItemIndex::Children)
-                    .filter(|(k, _)| {
-                        matches!(
-                            *k,
-                            CachedDataItemKey::Child { .. }
-                                | CachedDataItemKey::OutdatedChild { .. }
-                        )
-                    })
-                    .count();
                 let target_distance = children_count.ilog2() * 2;
                 let parent_aggregation = current_parent_aggregation
                     .base

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

@@ -313,6 +313,9 @@ pub enum CachedDataItem {
         task: TaskId,
         value: (),
     },
+    ChildrenCount {
+        value: u32,
+    },
 
     // Cells
     CellData {
@@ -438,6 +441,7 @@ impl CachedDataItem {
             }
             CachedDataItem::Dirty { .. } => true,
             CachedDataItem::Child { task, .. } => !task.is_transient(),
+            CachedDataItem::ChildrenCount { .. } => true,
             CachedDataItem::CellData { .. } => true,
             CachedDataItem::CellTypeMaxIndex { .. } => true,
             CachedDataItem::OutputDependency { target, .. } => !target.is_transient(),
@@ -502,6 +506,7 @@ impl CachedDataItemKey {
             }
             CachedDataItemKey::Dirty { .. } => true,
             CachedDataItemKey::Child { task, .. } => !task.is_transient(),
+            CachedDataItemKey::ChildrenCount {} => true,
             CachedDataItemKey::CellData { .. } => true,
             CachedDataItemKey::CellTypeMaxIndex { .. } => true,
             CachedDataItemKey::OutputDependency { target, .. } => !target.is_transient(),
@@ -534,6 +539,7 @@ impl CachedDataItemKey {
         match self {
             CachedDataItemKey::Collectible { .. }
             | CachedDataItemKey::Child { .. }
+            | CachedDataItemKey::ChildrenCount { .. }
             | CachedDataItemKey::CellData { .. }
             | CachedDataItemKey::CellTypeMaxIndex { .. }
             | CachedDataItemKey::OutputDependency { .. }
@@ -678,6 +684,7 @@ impl CachedDataItemValue {
 #[derive(Debug)]
 pub struct CachedDataUpdate {
     pub task: TaskId,
+    // TODO generate CachedDataItemUpdate to avoid repeating the variant field 3 times
     pub key: CachedDataItemKey,
     pub value: Option<CachedDataItemValue>,
     pub old_value: Option<CachedDataItemValue>,