Repository: vercel/next.js
Pull request: Turbopack: add a batch add method to the storage

### What?

Makes adding multiple items at once more efficient.

Changed files:
filename: turbopack/crates/turbo-tasks-backend/src/backend/dynamic_storage.rs
patch: @@ -58,6 +58,19 @@ impl DynamicStorage {
         self.get_or_create_map_mut(ty).add(item)
     }
 
+    pub fn extend(
+        &mut self,
+        ty: CachedDataItemType,
+        item: impl Iterator<Item = CachedDataItem>,
+    ) -> bool {
+        let map = self.get_or_create_map_mut(ty);
+        let mut added = true;
+        for item in item {
+            added = map.add(item) && added;
+        }
+        added
+    }
+
     pub fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue> {
         let ty = item.ty();
         self.get_or_create_map_mut(ty).insert(item)

filename: turbopack/crates/turbo-tasks-backend/src/backend/operation/mod.rs
patch: @@ -364,9 +364,26 @@ impl<'e, B: BackingStorage> ChildExecuteContext<'e> for ChildExecuteContextImpl<
 
 pub trait TaskGuard: Debug {
     fn id(&self) -> TaskId;
+    /// Adds a new item to the task if the key is not already present.
+    /// Returns `true` if the item was added.
+    /// Returns `false` if an item with the same key was already present.
     #[must_use]
     fn add(&mut self, item: CachedDataItem) -> bool;
+    /// Adds a new item to the task. The key must not be already present.
+    /// Might panic if the key is already present.
     fn add_new(&mut self, item: CachedDataItem);
+    /// Extends the task with items from the iterator.
+    /// Overwrites existing keys.
+    /// Returns `true` if all items were new and added.
+    /// Returns `false` if any item had a key that was already present.
+    fn extend(
+        &mut self,
+        ty: CachedDataItemType,
+        items: impl Iterator<Item = CachedDataItem>,
+    ) -> bool;
+    /// Extends the task with items from the iterator.
+    /// Might panic if any item has a key that is already present.
+    fn extend_new(&mut self, ty: CachedDataItemType, items: impl Iterator<Item = CachedDataItem>);
     fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue>;
     fn update(
         &mut self,
@@ -489,11 +506,49 @@ impl<B: BackingStorage> TaskGuard for TaskGuardImpl<'_, B> {
 
     #[track_caller]
     fn add_new(&mut self, item: CachedDataItem) {
-        self.check_access(item.category());
-        let added = self.add(item);
+        let category = item.category();
+        self.check_access(category);
+        if !self.task_id.is_transient() && item.is_persistent() {
+            self.task.track_modification(category.into_specific());
+        }
+        let added = self.task.add(item);
         assert!(added, "Item already exists");
     }
 
+    #[track_caller]
+    fn extend(
+        &mut self,
+        ty: CachedDataItemType,
+        items: impl Iterator<Item = CachedDataItem>,
+    ) -> bool {
+        let category = ty.category();
+        self.check_access(category);
+        if !self.task_id.is_transient() && ty.is_persistent() {
+            let mut items = items.peekable();
+            // Check if the iterator is empty
+            if items.peek().is_none() {
+                return true;
+            }
+            // TODO this is not optimal as we always track a modification even if nothing is changed
+            self.task.track_modification(category.into_specific());
+            self.task.extend(ty, items)
+        } else {
+            self.task.extend(ty, items)
+        }
+    }
+
+    #[track_caller]
+    fn extend_new(&mut self, ty: CachedDataItemType, items: impl Iterator<Item = CachedDataItem>) {
+        let category = ty.category();
+        self.check_access(category);
+        if !self.task_id.is_transient() && ty.is_persistent() {
+            self.task.track_modification(category.into_specific());
+        }
+
+        let added = self.task.extend(ty, items);
+        assert!(added, "At least one item already exists");
+    }
+
     #[track_caller]
     fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue> {
         let category = item.category();

filename: turbopack/crates/turbo-tasks-backend/src/backend/storage.rs
patch: @@ -317,6 +317,36 @@ macro_rules! generate_inner_storage_internal {
         $crate::generate_inner_storage_internal!(update: $self, $key, $update: $($config)+)
     };
 
+    // fn extend
+    (extend: $self:ident, $ty:ident, $items:ident: $tag:ident $key_field:ident => $field:ident,) => {
+        if let CachedDataItemType::$tag = $ty {
+            return $self.$field.extend($items.map(|item| {
+                let pair = turbo_tasks::KeyValuePair::into_key_and_value(item);
+                if let (CachedDataItemKey::$tag { $key_field }, CachedDataItemValue::$tag { value }) = pair {
+                    ($key_field, value)
+                } else {
+                    unreachable!()
+                }
+            }));
+        }
+    };
+    (extend: $self:ident, $ty:ident, $items:ident: $tag:ident => $field:ident,) => {
+        if let CachedDataItemType::$tag = $ty {
+            return $self.$field.extend($items.map(|item| {
+                let pair = turbo_tasks::KeyValuePair::into_key_and_value(item);
+                if let (_, CachedDataItemValue::$tag { value }) = pair {
+                    ((), value)
+                } else {
+                    unreachable!()
+                }
+            }));
+        }
+    };
+    (extend: $self:ident, $ty:ident, $items:ident: $tag:ident $($key_field:ident)? => $field:ident, $($config:tt)+) => {
+        $crate::generate_inner_storage_internal!(extend: $self, $ty, $items: $tag $($key_field)? => $field,);
+        $crate::generate_inner_storage_internal!(extend: $self, $ty, $items: $($config)+)
+    };
+
     // fn get_mut_or_insert_with
     (get_mut_or_insert_with: $self:ident, $key:ident, $insert_with:ident: $tag:ident $key_field:ident => $field:ident,) => {
         if let CachedDataItemKey::$tag { $key_field } = $key {
@@ -426,6 +456,12 @@ macro_rules! generate_inner_storage {
                 self.dynamic.add(item)
             }
 
+            pub fn extend(&mut self, ty: CachedDataItemType, items: impl Iterator<Item = CachedDataItem>) -> bool {
+                use crate::data_storage::Storage;
+                $crate::generate_inner_storage_internal!(extend: self, ty, items: $($config)*);
+                self.dynamic.extend(ty, items)
+            }
+
             pub fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue> {
                 use crate::data_storage::Storage;
                 $crate::generate_inner_storage_internal!(CachedDataItem: self, item, value, option_value, insert(value): $($config)*);

filename: turbopack/crates/turbo-tasks-backend/src/data_storage.rs
patch: @@ -10,7 +10,15 @@ pub trait Storage {
     where
         Self: 'l;
 
+    /// Adds a key-value pair to the storage, if the key is not already present.
+    /// Returns `true` if the key was not present and the value was added.
+    /// Returns `false` if the key was already present.
     fn add(&mut self, key: Self::K, value: Self::V) -> bool;
+    /// Extends the storage with key-value pairs from the iterator.
+    /// Overwrites existing keys.
+    /// Returns `true` if all keys were new and added.
+    /// Returns `false` if any key was already present.
+    fn extend(&mut self, items: impl Iterator<Item = (Self::K, Self::V)>) -> bool;
     fn insert(&mut self, key: Self::K, value: Self::V) -> Option<Self::V>;
     fn remove(&mut self, key: &Self::K) -> Option<Self::V>;
     fn contains_key(&self, key: &Self::K) -> bool;
@@ -60,6 +68,14 @@ impl<V> Storage for OptionStorage<V> {
         }
     }
 
+    fn extend(&mut self, items: impl Iterator<Item = (Self::K, Self::V)>) -> bool {
+        let mut added = true;
+        for (_, value) in items {
+            added = self.insert((), value).is_none() && added;
+        }
+        added
+    }
+
     fn insert(&mut self, _: (), value: V) -> Option<V> {
         self.value.replace(value)
     }
@@ -148,6 +164,13 @@ impl<K: Hash + Eq, V> Storage for AutoMapStorage<K, V> {
         }
     }
 
+    fn extend(&mut self, items: impl Iterator<Item = (Self::K, Self::V)>) -> bool {
+        let len = self.map.len();
+        let mut count = 0;
+        self.map.extend(items.inspect(|_| count += 1));
+        self.map.len() == len + count
+    }
+
     fn insert(&mut self, key: K, value: V) -> Option<V> {
         self.map.insert(key, value)
     }