[Turbopack] add count method to storage (#74830)

add count method to storage

avoid tracking children count

<!-- Thanks for opening a PR! Your contribution is much appreciated.
To make sure your PR is handled as smoothly as possible we request that you follow the checklist sections below.
Choose the right checklist for the change(s) that you're making:

## For Contributors

### Improving Documentation

- Run `pnpm prettier-fix` to fix formatting issues before opening the PR.
- Read the Docs Contribution Guide to ensure your contribution follows the docs guidelines: https://nextjs.org/docs/community/contribution-guide

### Adding or Updating Examples

- The "examples guidelines" are followed from our contributing doc https://github.com/vercel/next.js/blob/canary/contributing/examples/adding-examples.md
- Make sure the linting passes by running `pnpm build && pnpm lint`. See https://github.com/vercel/next.js/blob/canary/contributing/repository/linting.md

### Fixing a bug

- Related issues linked using `fixes #number`
- Tests added. See: https://github.com/vercel/next.js/blob/canary/contributing/core/testing.md#writing-tests-for-nextjs
- Errors have a helpful link attached, see https://github.com/vercel/next.js/blob/canary/contributing.md

### Adding a feature

- Implements an existing feature request or RFC. Make sure the feature request has been accepted for implementation before opening a PR. (A discussion must be opened, see https://github.com/vercel/next.js/discussions/new?category=ideas)
- Related issues/discussions are linked using `fixes #number`
- e2e tests added (https://github.com/vercel/next.js/blob/canary/contributing/core/testing.md#writing-tests-for-nextjs)
- Documentation added
- Telemetry added. In case of a feature if it's used or not.
- Errors have a helpful link attached, see https://github.com/vercel/next.js/blob/canary/contributing.md


## For Maintainers

- Minimal description (aim for explaining to someone not on the team to understand the PR)
- When linking to a Slack thread, you might want to share details of the conclusion
- Link both the Linear (Fixes NEXT-xxx) and the GitHub issues
- Add review comments if necessary to explain to the reviewer the logic behind a change

### What?

### Why?

### How?

Closes NEXT-
Fixes #

-->

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

@@ -13,7 +13,7 @@ use crate::{
             invalidate::{make_task_dirty, TaskDirtyCause},
             AggregatedDataUpdate, ExecuteContext, Operation, TaskGuard,
         },
-        storage::{update_count, update_ucount_and_get},
+        storage::update_count,
         TaskDataCategory,
     },
     data::{CachedDataItemKey, CellRef, CollectibleRef, CollectiblesRef},
@@ -82,8 +82,6 @@ impl Operation for CleanupOldEdgesOperation {
                                 for &child_id in children.iter() {
                                     task.remove(&CachedDataItemKey::Child { task: child_id });
                                 }
-                                let remove_children_count = u32::try_from(children.len()).unwrap();
-                                update_ucount_and_get!(task, ChildrenCount, -remove_children_count);
                                 if is_aggregating_node(get_aggregation_number(&task)) {
                                     queue.push(AggregationUpdateJob::InnerOfUpperLostFollowers {
                                         upper_id: task_id,

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

@@ -12,7 +12,7 @@ use crate::{
             },
             is_root_node, ExecuteContext, Operation, TaskGuard,
         },
-        storage::{get, update_ucount_and_get},
+        storage::{count, get},
         TaskDataCategory,
     },
     data::{CachedDataItem, CachedDataItemKey},
@@ -60,8 +60,8 @@ impl ConnectChildOperation {
         }) {
             let mut queue = AggregationUpdateQueue::new();
 
-            // Update the children count
-            let children_count = update_ucount_and_get!(parent_task, ChildrenCount, 1);
+            // Get the children count
+            let children_count = count!(parent_task, Child);
 
             // Compute future parent aggregation number based on the number of children
             let current_parent_aggregation = get!(parent_task, AggregationNumber)

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

@@ -371,6 +371,7 @@ pub trait TaskGuard: Debug {
     fn get(&self, key: &CachedDataItemKey) -> Option<CachedDataItemValueRef<'_>>;
     fn get_mut(&mut self, key: &CachedDataItemKey) -> Option<CachedDataItemValueRefMut<'_>>;
     fn has_key(&self, key: &CachedDataItemKey) -> bool;
+    fn count(&self, ty: CachedDataItemType) -> usize;
     fn iter(
         &self,
         ty: CachedDataItemType,
@@ -600,6 +601,11 @@ impl<B: BackingStorage> TaskGuard for TaskGuardImpl<'_, B> {
         self.task.has_key(key)
     }
 
+    fn count(&self, ty: CachedDataItemType) -> usize {
+        self.check_access(ty.category());
+        self.task.count(ty)
+    }
+
     fn iter(
         &self,
         ty: CachedDataItemType,

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

@@ -199,6 +199,10 @@ impl InnerStorage {
             .unwrap_or_default()
     }
 
+    pub fn count(&self, ty: CachedDataItemType) -> usize {
+        self.get_map(ty).map(|m| m.len()).unwrap_or_default()
+    }
+
     pub fn iter(
         &self,
         ty: CachedDataItemType,
@@ -321,6 +325,12 @@ impl DerefMut for StorageWriteGuard<'_> {
     }
 }
 
+macro_rules! count {
+    ($task:ident, $key:ident) => {{
+        $task.count($crate::data::CachedDataItemType::$key)
+    }};
+}
+
 macro_rules! get {
     ($task:ident, $key:ident $input:tt) => {{
         #[allow(unused_imports)]
@@ -516,6 +526,7 @@ macro_rules! remove {
     };
 }
 
+pub(crate) use count;
 pub(crate) use get;
 pub(crate) use get_many;
 pub(crate) use get_mut;

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

@@ -577,6 +577,12 @@ impl CachedDataItemKey {
         matches!(self, CachedDataItemKey::CellData { .. })
     }
 
+    pub fn category(&self) -> TaskDataCategory {
+        self.ty().category()
+    }
+}
+
+impl CachedDataItemType {
     pub fn category(&self) -> TaskDataCategory {
         match self {
             Self::Collectible { .. }

turbopack/crates/turbo-tasks-macros/src/derive/key_value_pair_macro.rs

@@ -162,6 +162,11 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream {
                                 storage.is_none()
                             }
                         },
+                        len: quote! {
+                            #storage_name::#variant_name { storage } => {
+                                if storage.is_some() { 1 } else { 0 }
+                            }
+                        },
                         iter: quote! {
                             #storage_name::#variant_name { storage } => {
                                 #iter_name::#variant_name(storage.iter())
@@ -234,6 +239,11 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream {
                                 storage.is_empty()
                             }
                         },
+                        len: quote! {
+                            #storage_name::#variant_name { storage } => {
+                                storage.len()
+                            }
+                        },
                         iter: quote! {
                             #storage_name::#variant_name { storage } => {
                                 #iter_name::#variant_name(storage.iter())
@@ -307,6 +317,11 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream {
                                 storage.is_empty()
                             }
                         },
+                        len: quote! {
+                            #storage_name::#variant_name { storage } => {
+                                storage.len()
+                            }
+                        },
                         iter: quote! {
                             #storage_name::#variant_name { storage } => {
                                 #iter_name::#variant_name(storage.iter())
@@ -344,6 +359,7 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream {
         .iter()
         .map(|decl| &decl.is_empty)
         .collect::<Vec<_>>();
+    let storage_len = storage.iter().map(|decl| &decl.len).collect::<Vec<_>>();
     let storage_iter = storage.iter().map(|decl| &decl.iter).collect::<Vec<_>>();
     let storage_iterator = storage
         .iter()
@@ -590,6 +606,14 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream {
                 }
             }
 
+            pub fn len(&self) -> usize {
+                match self {
+                    #(
+                        #storage_len
+                    )*
+                }
+            }
+
             pub fn iter(&self) -> #iter_name {
                 match self {
                     #(
@@ -764,6 +788,7 @@ struct StorageDecl {
     get_mut: proc_macro2::TokenStream,
     shrink_to_fit: proc_macro2::TokenStream,
     is_empty: proc_macro2::TokenStream,
+    len: proc_macro2::TokenStream,
     iter: proc_macro2::TokenStream,
 
     iterator: proc_macro2::TokenStream,