perf(turbo-tasks-backend): use DefaultStorage for AggregationNumber to save memory (#88336)

This drops the size of the `aggregation_number` field from 16 to 12 bytes which allows the full `InnerStorage` struct to do from 136 bytes to 128 bytes (thanks alignment!)

Due to allocator behavior this is probably saving more like 32 bytes per InnerStorage allocation, see [mimalloc size classes](https://github.com/microsoft/mimalloc/blob/f0cd5505aa102cee991be0367b82506638a16281/src/init.c#L60)

Author: lukesandberg

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

@@ -14,7 +14,7 @@ use crate::{
         AggregationNumber, CachedDataItem, CachedDataItemKey, CachedDataItemType,
         CachedDataItemValue, CachedDataItemValueRef, CachedDataItemValueRefMut, OutputValue,
     },
-    data_storage::{AutoMapStorage, OptionStorage},
+    data_storage::{AutoMapStorage, DefaultStorage, OptionStorage},
     utils::{
         dash_map_drop_contents::drop_contents,
         dash_map_multi::{RefMut, get_multiple_mut},
@@ -138,7 +138,7 @@ impl InnerStorageState {
 }
 
 pub struct InnerStorageSnapshot {
-    aggregation_number: OptionStorage<AggregationNumber>,
+    aggregation_number: DefaultStorage<AggregationNumber>,
     output_dependent: AutoMapStorage<TaskId, ()>,
     output: OptionStorage<OutputValue>,
     upper: AutoMapStorage<TaskId, u32>,
@@ -206,7 +206,7 @@ impl InnerStorageSnapshot {
 
 #[derive(Debug, Clone)]
 pub struct InnerStorage {
-    aggregation_number: OptionStorage<AggregationNumber>,
+    aggregation_number: DefaultStorage<AggregationNumber>,
     output_dependent: AutoMapStorage<TaskId, ()>,
     output: OptionStorage<OutputValue>,
     upper: AutoMapStorage<TaskId, u32>,
@@ -1220,3 +1220,24 @@ where
         None
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_inner_storage_size() {
+        // InnerStorage size affects memory usage per task.
+        // We track this to catch unexpected bloat from type changes.
+        assert_eq!(
+            std::mem::size_of::<InnerStorage>(),
+            128,
+            "InnerStorage size changed - please review if this is intentional"
+        );
+        assert_eq!(
+            std::mem::size_of::<InnerStorageSnapshot>(),
+            128,
+            "InnerStorageSnapshot size changed - please review if this is intentional"
+        );
+    }
+}

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

@@ -133,6 +133,120 @@ impl<V> Storage for OptionStorage<V> {
     }
 }
 
+/// Storage for a single value that uses `Default::default()` to initialize the value "empty".
+#[derive(Debug, Clone)]
+pub struct DefaultStorage<V> {
+    value: V,
+}
+
+impl<V: Default> Default for DefaultStorage<V> {
+    fn default() -> Self {
+        Self {
+            value: V::default(),
+        }
+    }
+}
+
+impl<V: Default + PartialEq> Storage for DefaultStorage<V> {
+    type K = ();
+    type V = V;
+    type Iterator<'l>
+        = std::option::IntoIter<(&'l (), &'l V)>
+    where
+        Self: 'l;
+
+    fn add(&mut self, _: (), value: V) -> bool {
+        if self.value == V::default() {
+            self.value = value;
+            true
+        } else {
+            false
+        }
+    }
+
+    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> {
+        let old = std::mem::replace(&mut self.value, value);
+        if old == V::default() { None } else { Some(old) }
+    }
+
+    fn remove(&mut self, _: &()) -> Option<V> {
+        let old = std::mem::take(&mut self.value);
+        if old == V::default() { None } else { Some(old) }
+    }
+
+    fn contains_key(&self, _: &()) -> bool {
+        self.value != V::default()
+    }
+
+    fn get(&self, _: &()) -> Option<&V> {
+        if self.value != V::default() {
+            Some(&self.value)
+        } else {
+            None
+        }
+    }
+
+    fn get_mut(&mut self, _: &()) -> Option<&mut V> {
+        if self.value != V::default() {
+            Some(&mut self.value)
+        } else {
+            None
+        }
+    }
+
+    fn get_mut_or_insert_with(&mut self, _: (), f: impl FnOnce() -> V) -> &mut V {
+        if self.value == V::default() {
+            self.value = f();
+        }
+        &mut self.value
+    }
+
+    fn shrink_to_fit(&mut self) {
+        // Nothing to do
+    }
+
+    fn is_empty(&self) -> bool {
+        self.value == V::default()
+    }
+
+    fn len(&self) -> usize {
+        if self.value != V::default() { 1 } else { 0 }
+    }
+
+    fn iter(&self) -> Self::Iterator<'_> {
+        if self.value != V::default() {
+            Some(value_to_key_value(&self.value)).into_iter()
+        } else {
+            None.into_iter()
+        }
+    }
+
+    fn extract_if<'l, F>(&'l mut self, mut f: F) -> impl Iterator<Item = (Self::K, Self::V)>
+    where
+        F: for<'a, 'b> FnMut(&'a Self::K, &'b mut Self::V) -> bool + 'l,
+    {
+        if self.value != V::default() && f(&(), &mut self.value) {
+            let old = std::mem::take(&mut self.value);
+            return Some(((), old)).into_iter();
+        }
+        None.into_iter()
+    }
+
+    fn update(&mut self, _: (), update: impl FnOnce(Option<V>) -> Option<V>) {
+        let old = std::mem::take(&mut self.value);
+        let old_opt = if old == V::default() { None } else { Some(old) };
+        self.value = update(old_opt).unwrap_or_default();
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct AutoMapStorage<K, V> {
     map: AutoMap<K, V, BuildHasherDefault<FxHasher>, 1>,