Skip to content

Commit 50f5fbe

Browse files
sokrasokra
authored
Turbopack: add a batch add method to the storage (#84270)
### What? Makes adding multiple items at once more efficient.
1 parent 763c1f4 commit 50f5fbe

File tree

4 files changed

+129
-2
lines changed

4 files changed

+129
-2
lines changed

‎turbopack/crates/turbo-tasks-backend/src/backend/dynamic_storage.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,19 @@ impl DynamicStorage {
5858
self.get_or_create_map_mut(ty).add(item)
5959
}
6060

61+
pub fn extend(
62+
&mut self,
63+
ty: CachedDataItemType,
64+
item: impl Iterator<Item = CachedDataItem>,
65+
) -> bool {
66+
let map = self.get_or_create_map_mut(ty);
67+
let mut added = true;
68+
for item in item {
69+
added = map.add(item) && added;
70+
}
71+
added
72+
}
73+
6174
pub fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue> {
6275
let ty = item.ty();
6376
self.get_or_create_map_mut(ty).insert(item)

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

Lines changed: 57 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -364,9 +364,26 @@ impl<'e, B: BackingStorage> ChildExecuteContext<'e> for ChildExecuteContextImpl<
364364

365365
pub trait TaskGuard: Debug {
366366
fn id(&self) -> TaskId;
367+
/// Adds a new item to the task if the key is not already present.
368+
/// Returns `true` if the item was added.
369+
/// Returns `false` if an item with the same key was already present.
367370
#[must_use]
368371
fn add(&mut self, item: CachedDataItem) -> bool;
372+
/// Adds a new item to the task. The key must not be already present.
373+
/// Might panic if the key is already present.
369374
fn add_new(&mut self, item: CachedDataItem);
375+
/// Extends the task with items from the iterator.
376+
/// Overwrites existing keys.
377+
/// Returns `true` if all items were new and added.
378+
/// Returns `false` if any item had a key that was already present.
379+
fn extend(
380+
&mut self,
381+
ty: CachedDataItemType,
382+
items: impl Iterator<Item = CachedDataItem>,
383+
) -> bool;
384+
/// Extends the task with items from the iterator.
385+
/// Might panic if any item has a key that is already present.
386+
fn extend_new(&mut self, ty: CachedDataItemType, items: impl Iterator<Item = CachedDataItem>);
370387
fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue>;
371388
fn update(
372389
&mut self,
@@ -489,11 +506,49 @@ impl<B: BackingStorage> TaskGuard for TaskGuardImpl<'_, B> {
489506

490507
#[track_caller]
491508
fn add_new(&mut self, item: CachedDataItem) {
492-
self.check_access(item.category());
493-
let added = self.add(item);
509+
let category = item.category();
510+
self.check_access(category);
511+
if !self.task_id.is_transient() && item.is_persistent() {
512+
self.task.track_modification(category.into_specific());
513+
}
514+
let added = self.task.add(item);
494515
assert!(added, "Item already exists");
495516
}
496517

518+
#[track_caller]
519+
fn extend(
520+
&mut self,
521+
ty: CachedDataItemType,
522+
items: impl Iterator<Item = CachedDataItem>,
523+
) -> bool {
524+
let category = ty.category();
525+
self.check_access(category);
526+
if !self.task_id.is_transient() && ty.is_persistent() {
527+
let mut items = items.peekable();
528+
// Check if the iterator is empty
529+
if items.peek().is_none() {
530+
return true;
531+
}
532+
// TODO this is not optimal as we always track a modification even if nothing is changed
533+
self.task.track_modification(category.into_specific());
534+
self.task.extend(ty, items)
535+
} else {
536+
self.task.extend(ty, items)
537+
}
538+
}
539+
540+
#[track_caller]
541+
fn extend_new(&mut self, ty: CachedDataItemType, items: impl Iterator<Item = CachedDataItem>) {
542+
let category = ty.category();
543+
self.check_access(category);
544+
if !self.task_id.is_transient() && ty.is_persistent() {
545+
self.task.track_modification(category.into_specific());
546+
}
547+
548+
let added = self.task.extend(ty, items);
549+
assert!(added, "At least one item already exists");
550+
}
551+
497552
#[track_caller]
498553
fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue> {
499554
let category = item.category();

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,36 @@ macro_rules! generate_inner_storage_internal {
317317
$crate::generate_inner_storage_internal!(update: $self, $key, $update: $($config)+)
318318
};
319319

320+
// fn extend
321+
(extend: $self:ident, $ty:ident, $items:ident: $tag:ident $key_field:ident => $field:ident,) => {
322+
if let CachedDataItemType::$tag = $ty {
323+
return $self.$field.extend($items.map(|item| {
324+
let pair = turbo_tasks::KeyValuePair::into_key_and_value(item);
325+
if let (CachedDataItemKey::$tag { $key_field }, CachedDataItemValue::$tag { value }) = pair {
326+
($key_field, value)
327+
} else {
328+
unreachable!()
329+
}
330+
}));
331+
}
332+
};
333+
(extend: $self:ident, $ty:ident, $items:ident: $tag:ident => $field:ident,) => {
334+
if let CachedDataItemType::$tag = $ty {
335+
return $self.$field.extend($items.map(|item| {
336+
let pair = turbo_tasks::KeyValuePair::into_key_and_value(item);
337+
if let (_, CachedDataItemValue::$tag { value }) = pair {
338+
((), value)
339+
} else {
340+
unreachable!()
341+
}
342+
}));
343+
}
344+
};
345+
(extend: $self:ident, $ty:ident, $items:ident: $tag:ident $($key_field:ident)? => $field:ident, $($config:tt)+) => {
346+
$crate::generate_inner_storage_internal!(extend: $self, $ty, $items: $tag $($key_field)? => $field,);
347+
$crate::generate_inner_storage_internal!(extend: $self, $ty, $items: $($config)+)
348+
};
349+
320350
// fn get_mut_or_insert_with
321351
(get_mut_or_insert_with: $self:ident, $key:ident, $insert_with:ident: $tag:ident $key_field:ident => $field:ident,) => {
322352
if let CachedDataItemKey::$tag { $key_field } = $key {
@@ -426,6 +456,12 @@ macro_rules! generate_inner_storage {
426456
self.dynamic.add(item)
427457
}
428458

459+
pub fn extend(&mut self, ty: CachedDataItemType, items: impl Iterator<Item = CachedDataItem>) -> bool {
460+
use crate::data_storage::Storage;
461+
$crate::generate_inner_storage_internal!(extend: self, ty, items: $($config)*);
462+
self.dynamic.extend(ty, items)
463+
}
464+
429465
pub fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue> {
430466
use crate::data_storage::Storage;
431467
$crate::generate_inner_storage_internal!(CachedDataItem: self, item, value, option_value, insert(value): $($config)*);

‎turbopack/crates/turbo-tasks-backend/src/data_storage.rs‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,15 @@ pub trait Storage {
1010
where
1111
Self: 'l;
1212

13+
/// Adds a key-value pair to the storage, if the key is not already present.
14+
/// Returns `true` if the key was not present and the value was added.
15+
/// Returns `false` if the key was already present.
1316
fn add(&mut self, key: Self::K, value: Self::V) -> bool;
17+
/// Extends the storage with key-value pairs from the iterator.
18+
/// Overwrites existing keys.
19+
/// Returns `true` if all keys were new and added.
20+
/// Returns `false` if any key was already present.
21+
fn extend(&mut self, items: impl Iterator<Item = (Self::K, Self::V)>) -> bool;
1422
fn insert(&mut self, key: Self::K, value: Self::V) -> Option<Self::V>;
1523
fn remove(&mut self, key: &Self::K) -> Option<Self::V>;
1624
fn contains_key(&self, key: &Self::K) -> bool;
@@ -60,6 +68,14 @@ impl<V> Storage for OptionStorage<V> {
6068
}
6169
}
6270

71+
fn extend(&mut self, items: impl Iterator<Item = (Self::K, Self::V)>) -> bool {
72+
let mut added = true;
73+
for (_, value) in items {
74+
added = self.insert((), value).is_none() && added;
75+
}
76+
added
77+
}
78+
6379
fn insert(&mut self, _: (), value: V) -> Option<V> {
6480
self.value.replace(value)
6581
}
@@ -148,6 +164,13 @@ impl<K: Hash + Eq, V> Storage for AutoMapStorage<K, V> {
148164
}
149165
}
150166

167+
fn extend(&mut self, items: impl Iterator<Item = (Self::K, Self::V)>) -> bool {
168+
let len = self.map.len();
169+
let mut count = 0;
170+
self.map.extend(items.inspect(|_| count += 1));
171+
self.map.len() == len + count
172+
}
173+
151174
fn insert(&mut self, key: K, value: V) -> Option<V> {
152175
self.map.insert(key, value)
153176
}

0 commit comments

Comments
 (0)