shrink lists and sets that are stored in Tasks (#2873)

Reduces memory usage by 14%

crates/auto-hash-map/src/map.rs

@@ -179,6 +179,14 @@ impl<K: Eq + Hash, V, H: BuildHasher + Default> AutoMap<K, V, H> {
             },
         }
     }
+
+    /// see [HashMap::shrink_to_fit](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.shrink_to_fit)
+    pub fn shrink_to_fit(&mut self) {
+        match self {
+            AutoMap::List(list) => list.shrink_to_fit(),
+            AutoMap::Map(map) => map.shrink_to_fit(),
+        }
+    }
 }
 
 impl<K: Eq + Hash, V, H: BuildHasher> AutoMap<K, V, H> {
@@ -225,6 +233,13 @@ impl<K, V, H> AutoMap<K, V, H> {
             AutoMap::Map(map) => Iter::Map(map.iter()),
         }
     }
+    /// see [HashMap::iter_mut](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.iter_mut)
+    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
+        match self {
+            AutoMap::List(list) => IterMut::List(list.iter_mut()),
+            AutoMap::Map(map) => IterMut::Map(map.iter_mut()),
+        }
+    }
 
     /// see [HashMap::is_empty](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.is_empty)
     pub fn is_empty(&self) -> bool {
@@ -296,6 +311,22 @@ impl<'a, K, V> Iterator for Iter<'a, K, V> {
     }
 }
 
+pub enum IterMut<'a, K, V> {
+    List(std::slice::IterMut<'a, (K, V)>),
+    Map(std::collections::hash_map::IterMut<'a, K, V>),
+}
+
+impl<'a, K, V> Iterator for IterMut<'a, K, V> {
+    type Item = (&'a K, &'a mut V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            IterMut::List(iter) => iter.next().map(|(k, v)| (&*k, v)),
+            IterMut::Map(iter) => iter.next().map(|(k, v)| (k, v)),
+        }
+    }
+}
+
 pub enum IntoIter<K, V> {
     List(std::vec::IntoIter<(K, V)>),
     Map(std::collections::hash_map::IntoIter<K, V>),

crates/auto-hash-map/src/set.rs

@@ -76,6 +76,11 @@ impl<K: Hash + Eq, H: BuildHasher + Default> AutoSet<K, H> {
         self.map.extend(iter.into_iter().map(|item| (item, ())))
     }
 
+    /// see [HashSet::shrink_to_fit](https://doc.rust-lang.org/std/collections/hash_set/struct.HashSet.html#method.shrink_to_fit)
+    pub fn shrink_to_fit(&mut self) {
+        self.map.shrink_to_fit();
+    }
+
     /// see [HashSet::contains](https://doc.rust-lang.org/std/collections/hash_set/struct.HashSet.html#method.contains)
     pub fn contains(&self, key: &K) -> bool {
         self.map.contains_key(key)

crates/turbo-tasks-memory/src/memory_backend.rs

@@ -421,7 +421,7 @@ impl Backend for MemoryBackend {
 
     fn get_or_create_persistent_task(
         &self,
-        task_type: PersistentTaskType,
+        mut task_type: PersistentTaskType,
         parent_task: TaskId,
         turbo_tasks: &dyn TurboTasksBackendApi,
     ) -> TaskId {
@@ -433,6 +433,9 @@ impl Backend for MemoryBackend {
             // "in progress" until they become active
             task
         } else {
+            // It's important to avoid overallocating memory as this will go into the task
+            // cache and stay there forever. We can to be as small as possible.
+            task_type.shrink_to_fit();
             // slow pass with key lock
             let id = turbo_tasks.get_fresh_task_id();
             let task = match &task_type {

crates/turbo-tasks-memory/src/task.rs

@@ -555,9 +555,14 @@ impl Task {
             match state.state_type {
                 InProgress { ref mut event } => {
                     let event = event.take();
-                    state.state_type = Done {
-                        dependencies: take(&mut dependencies),
-                    };
+                    let mut dependencies = take(&mut dependencies);
+                    // This will stay here for longer, so make sure to not consume too much memory
+                    dependencies.shrink_to_fit();
+                    for cells in state.cells.values_mut() {
+                        cells.shrink_to_fit();
+                    }
+                    state.cells.shrink_to_fit();
+                    state.state_type = Done { dependencies };
                     for scope in state.scopes.iter() {
                         backend.with_scope(scope, |scope| {
                             scope.decrement_unfinished_tasks(backend);

crates/turbo-tasks/src/backend.rs

@@ -76,6 +76,14 @@ pub enum PersistentTaskType {
 }
 
 impl PersistentTaskType {
+    pub fn shrink_to_fit(&mut self) {
+        match self {
+            Self::Native(_, inputs) => inputs.shrink_to_fit(),
+            Self::ResolveNative(_, inputs) => inputs.shrink_to_fit(),
+            Self::ResolveTrait(_, _, inputs) => inputs.shrink_to_fit(),
+        }
+    }
+
     pub fn len(&self) -> usize {
         match self {
             PersistentTaskType::Native(_, v)