Performance Improvements (1+2) review follow-up (vercel/turborepo#6061)

### Description

Code cleanup


Closes WEB-1683

crates/turbo-tasks-memory/src/aggregation_tree/bottom_connection.rs

@@ -2,12 +2,11 @@ use std::{hash::Hash, ops::ControlFlow, sync::Arc};
 
 use auto_hash_map::{map::RawEntry, AutoMap};
 use nohash_hasher::{BuildNoHashHasher, IsEnabled};
-use smallvec::SmallVec;
 
 use super::{
     bottom_tree::BottomTree,
     inner_refs::{BottomRef, ChildLocation},
-    AggregationContext,
+    AggregationContext, StackVec,
 };
 
 struct BottomRefInfo {
@@ -197,7 +196,7 @@ impl<T, I: IsEnabled + Eq + Hash + Clone> BottomConnection<T, I> {
 
 pub enum BottomUppers<T, I: IsEnabled> {
     Left(Arc<BottomTree<T, I>>),
-    Inner(SmallVec<[(BottomRef<T, I>, u8); 16]>),
+    Inner(StackVec<(BottomRef<T, I>, u8)>),
 }
 
 impl<T, I: IsEnabled + Eq + Hash + Clone> BottomUppers<T, I> {

crates/turbo-tasks-memory/src/aggregation_tree/bottom_tree.rs

@@ -3,7 +3,6 @@ use std::{hash::Hash, ops::ControlFlow, sync::Arc};
 use nohash_hasher::{BuildNoHashHasher, IsEnabled};
 use parking_lot::{Mutex, MutexGuard};
 use ref_cast::RefCast;
-use smallvec::SmallVec;
 
 use super::{
     bottom_connection::BottomConnection,
@@ -13,7 +12,7 @@ use super::{
         remove_left_upper_from_item,
     },
     top_tree::TopTree,
-    AggregationContext, CHILDREN_INNER_THRESHOLD, CONNECTIVITY_LIMIT,
+    AggregationContext, StackVec, CHILDREN_INNER_THRESHOLD, CONNECTIVITY_LIMIT,
 };
 use crate::count_hash_set::{CountHashSet, RemoveIfEntryResult};
 
@@ -82,15 +81,15 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
         }
     }
 
-    fn add_children_of_child_if_following(&self, children: &mut SmallVec<[&I; 16]>) {
+    fn add_children_of_child_if_following(&self, children: &mut StackVec<&I>) {
         let mut state = self.state.lock();
         children.retain(|&mut child| !state.following.add_if_entry(child));
     }
 
     fn add_children_of_child_following<C: AggregationContext<Info = T, ItemRef = I>>(
         self: &Arc<Self>,
         aggregation_context: &C,
-        mut children: SmallVec<[&I; 16]>,
+        mut children: StackVec<&I>,
     ) {
         let mut state = self.state.lock();
         children.retain(|&mut child| state.following.add_clonable(child));
@@ -114,7 +113,7 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
     ) where
         I: 'a,
     {
-        let mut following = SmallVec::default();
+        let mut following = StackVec::default();
         if self.height == 0 {
             for child in children {
                 let can_be_inner =
@@ -254,7 +253,7 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
         children: &mut Vec<&'a I>,
     ) {
         let mut state = self.state.lock();
-        let mut removed = SmallVec::<[_; 16]>::default();
+        let mut removed = StackVec::default();
         children.retain(|&child| match state.following.remove_if_entry(child) {
             RemoveIfEntryResult::PartiallyRemoved => false,
             RemoveIfEntryResult::NotPresent => true,
@@ -285,7 +284,7 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
     fn remove_children_of_child_following<C: AggregationContext<Info = T, ItemRef = I>>(
         self: &Arc<Self>,
         aggregation_context: &C,
-        mut children: SmallVec<[&I; 16]>,
+        mut children: StackVec<&I>,
     ) {
         let mut state = self.state.lock();
         children.retain(|&mut child| state.following.remove_clonable(child));
@@ -315,7 +314,7 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
     ) where
         I: 'a,
     {
-        let unremoveable: SmallVec<[_; 16]> = if self.height == 0 {
+        let unremoveable: StackVec<_> = if self.height == 0 {
             children
                 .into_iter()
                 .filter(|&child| !remove_inner_upper_from_item(aggregation_context, child, self))
@@ -342,11 +341,7 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
         let mut state = self.state.lock();
         let old_inner = state.bottom_upper.set_left_upper(upper);
         let add_change = aggregation_context.info_to_add_change(&state.data);
-        let children = state
-            .following
-            .iter()
-            .cloned()
-            .collect::<SmallVec<[_; 16]>>();
+        let children = state.following.iter().cloned().collect::<StackVec<_>>();
 
         let remove_change = (!old_inner.is_unset())
             .then(|| aggregation_context.info_to_remove_change(&state.data))
@@ -443,11 +438,7 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
             if let Some(change) = aggregation_context.info_to_add_change(&state.data) {
                 upper.child_change(aggregation_context, &change);
             }
-            let children = state
-                .following
-                .iter()
-                .cloned()
-                .collect::<SmallVec<[_; 16]>>();
+            let children = state.following.iter().cloned().collect::<StackVec<_>>();
             drop(state);
             if !children.is_empty() {
                 upper.add_children_of_child(
@@ -471,11 +462,7 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
         if let Some(change) = aggregation_context.info_to_remove_change(&state.data) {
             upper.child_change(aggregation_context, &change);
         }
-        let following = state
-            .following
-            .iter()
-            .cloned()
-            .collect::<SmallVec<[_; 16]>>();
+        let following = state.following.iter().cloned().collect::<StackVec<_>>();
         if state.top_upper.is_empty() {
             drop(state);
             self.remove_self_from_lower(aggregation_context);
@@ -498,11 +485,7 @@ impl<T, I: Clone + Eq + Hash + IsEnabled> BottomTree<T, I> {
         let removed = inner.remove_clonable(BottomRef::ref_cast(upper));
         if removed {
             let remove_change = aggregation_context.info_to_remove_change(&state.data);
-            let following = state
-                .following
-                .iter()
-                .cloned()
-                .collect::<SmallVec<[_; 16]>>();
+            let following = state.following.iter().cloned().collect::<StackVec<_>>();
             drop(state);
             if let Some(change) = remove_change {
                 upper.child_change(aggregation_context, &change);
@@ -600,11 +583,7 @@ fn propagate_lost_following_to_uppers<C: AggregationContext>(
     child_of_child: &C::ItemRef,
 ) {
     let bottom_uppers = state.bottom_upper.as_cloned_uppers();
-    let top_upper = state
-        .top_upper
-        .iter()
-        .cloned()
-        .collect::<SmallVec<[_; 16]>>();
+    let top_upper = state.top_upper.iter().cloned().collect::<StackVec<_>>();
     drop(state);
     for TopRef { upper } in top_upper {
         upper.remove_child_of_child(aggregation_context, child_of_child);

crates/turbo-tasks-memory/src/aggregation_tree/leaf.rs

@@ -3,15 +3,14 @@ use std::{hash::Hash, sync::Arc};
 use auto_hash_map::AutoSet;
 use nohash_hasher::IsEnabled;
 use ref_cast::RefCast;
-use smallvec::SmallVec;
 use tracing::Level;
 
 use super::{
     bottom_connection::{BottomConnection, DistanceCountMap},
     bottom_tree::BottomTree,
     inner_refs::{BottomRef, ChildLocation},
     top_tree::TopTree,
-    AggregationContext, AggregationItemLock, CHILDREN_INNER_THRESHOLD,
+    AggregationContext, AggregationItemLock, LargeStackVec, CHILDREN_INNER_THRESHOLD,
 };
 
 /// The leaf of the aggregation tree. It's usually stored inside of the nodes
@@ -245,7 +244,7 @@ pub fn add_inner_upper_to_item<C: AggregationContext>(
                 change,
                 item.children()
                     .map(|r| r.into_owned())
-                    .collect::<SmallVec<[_; 128]>>(),
+                    .collect::<LargeStackVec<_>>(),
             )
         } else {
             return true;
@@ -267,7 +266,7 @@ pub fn add_inner_upper_to_item<C: AggregationContext>(
 
 struct AddLeftUpperIntermediateResult<C: AggregationContext>(
     Option<C::ItemChange>,
-    SmallVec<[C::ItemRef; 128]>,
+    LargeStackVec<C::ItemRef>,
     DistanceCountMap<BottomRef<C::Info, C::ItemRef>>,
     Option<C::ItemChange>,
 );
@@ -352,7 +351,7 @@ pub fn remove_inner_upper_from_item<C: AggregationContext>(
     let children = item
         .children()
         .map(|r| r.into_owned())
-        .collect::<SmallVec<[_; 128]>>();
+        .collect::<LargeStackVec<_>>();
     drop(item);
 
     if let Some(change) = change {

crates/turbo-tasks-memory/src/aggregation_tree/mod.rs

@@ -35,6 +35,7 @@ mod top_tree;
 use std::{borrow::Cow, hash::Hash, ops::ControlFlow, sync::Arc};
 
 use nohash_hasher::IsEnabled;
+use smallvec::SmallVec;
 
 use self::{leaf::top_tree, top_tree::TopTree};
 pub use self::{
@@ -49,6 +50,9 @@ const CONNECTIVITY_LIMIT: u8 = 7;
 /// When reached the parent of the children will form a new bottom tree.
 const CHILDREN_INNER_THRESHOLD: usize = 2000;
 
+type StackVec<I> = SmallVec<[I; 16]>;
+type LargeStackVec<I> = SmallVec<[I; 32]>;
+
 /// The context trait which defines how the aggregation tree should behave.
 pub trait AggregationContext {
     type ItemLock<'a>: AggregationItemLock<

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

@@ -4,11 +4,10 @@ use std::{
 };
 
 use auto_hash_map::AutoSet;
-use nohash_hasher::BuildNoHashHasher;
 use turbo_tasks::{
     backend::CellContent,
     event::{Event, EventListener},
-    TaskId, TurboTasksBackendApi,
+    TaskId, TaskIdSet, TurboTasksBackendApi,
 };
 
 use crate::MemoryBackend;
@@ -25,20 +24,18 @@ pub(crate) enum Cell {
     /// tracking is still active. Any update will invalidate dependent tasks.
     /// Assigning a value will transition to the Value state.
     /// Reading this cell will transition to the Recomputing state.
-    TrackedValueless {
-        dependent_tasks: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
-    },
+    TrackedValueless { dependent_tasks: TaskIdSet },
     /// Someone wanted to read the content and it was not available. The content
     /// is now being recomputed.
     /// Assigning a value will transition to the Value state.
     Recomputing {
-        dependent_tasks: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+        dependent_tasks: TaskIdSet,
         event: Event,
     },
     /// The content was set only once and is tracked.
     /// GC operation will transition to the TrackedValueless state.
     Value {
-        dependent_tasks: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+        dependent_tasks: TaskIdSet,
         content: CellContent,
     },
 }
@@ -95,11 +92,10 @@ impl Cell {
     }
 
     /// Returns the list of dependent tasks.
-    pub fn dependent_tasks(&self) -> &AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2> {
+    pub fn dependent_tasks(&self) -> &TaskIdSet {
         match self {
             Cell::Empty => {
-                static EMPTY: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2> =
-                    AutoSet::with_hasher();
+                static EMPTY: TaskIdSet = AutoSet::with_hasher();
                 &EMPTY
             }
             Cell::Value {
@@ -117,7 +113,7 @@ impl Cell {
     /// Switch the cell to recomputing state.
     fn recompute(
         &mut self,
-        dependent_tasks: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+        dependent_tasks: TaskIdSet,
         description: impl Fn() -> String + Sync + Send + 'static,
         note: impl Fn() -> String + Sync + Send + 'static,
     ) -> EventListener {

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

@@ -15,7 +15,6 @@ use std::{
 use anyhow::{bail, Result};
 use auto_hash_map::{AutoMap, AutoSet};
 use dashmap::{mapref::entry::Entry, DashMap};
-use nohash_hasher::BuildNoHashHasher;
 use rustc_hash::FxHasher;
 use tokio::task::futures::TaskLocalFuture;
 use tracing::trace_span;
@@ -26,7 +25,7 @@ use turbo_tasks::{
     },
     event::EventListener,
     util::{IdFactory, NoMoveVec},
-    CellId, RawVc, TaskId, TraitTypeId, TurboTasksBackendApi, Unused,
+    CellId, RawVc, TaskId, TaskIdSet, TraitTypeId, TurboTasksBackendApi, Unused,
 };
 
 use crate::{
@@ -221,7 +220,7 @@ impl MemoryBackend {
 
     pub(crate) fn schedule_when_dirty_from_aggregation(
         &self,
-        set: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+        set: TaskIdSet,
         turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
     ) {
         for task in set {
@@ -262,7 +261,7 @@ impl Backend for MemoryBackend {
 
     fn invalidate_tasks_set(
         &self,
-        tasks: &AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+        tasks: &TaskIdSet,
         turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
     ) {
         for &task in tasks {

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

@@ -16,7 +16,6 @@ use anyhow::{anyhow, Result};
 use auto_hash_map::{AutoMap, AutoSet};
 use concurrent_queue::ConcurrentQueue;
 use dashmap::{mapref::entry::Entry, DashMap, DashSet};
-use nohash_hasher::BuildNoHashHasher;
 use turbo_tasks::{
     backend::{
         Backend, BackendJobId, CellContent, PersistentTaskType, TaskExecutionSpec,
@@ -28,7 +27,7 @@ use turbo_tasks::{
         PersistedGraphApi, ReadTaskState, TaskCell, TaskData,
     },
     util::{IdFactory, NoMoveVec, SharedError},
-    CellId, RawVc, TaskId, TraitTypeId, TurboTasksBackendApi, Unused,
+    CellId, RawVc, TaskId, TaskIdSet, TraitTypeId, TurboTasksBackendApi, Unused,
 };
 
 type RootTaskFn =
@@ -65,11 +64,11 @@ struct MemoryTaskState {
     need_persist: bool,
     has_changes: bool,
     freshness: TaskFreshness,
-    cells: HashMap<CellId, (TaskCell, AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>)>,
+    cells: HashMap<CellId, (TaskCell, TaskIdSet)>,
     output: Option<Result<RawVc, SharedError>>,
-    output_dependent: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+    output_dependent: TaskIdSet,
     dependencies: AutoSet<RawVc>,
-    children: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+    children: TaskIdSet,
     event: Event,
     event_cells: Event,
 }
@@ -1040,7 +1039,7 @@ impl<P: PersistedGraph> Backend for MemoryBackendWithPersistedGraph<P> {
 
     fn invalidate_tasks_set(
         &self,
-        tasks: &AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+        tasks: &TaskIdSet,
         turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackendWithPersistedGraph<P>>,
     ) {
         for &task in tasks {

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

@@ -5,17 +5,15 @@ use std::{
 };
 
 use anyhow::{anyhow, Error, Result};
-use auto_hash_map::AutoSet;
-use nohash_hasher::BuildNoHashHasher;
-use turbo_tasks::{util::SharedError, RawVc, TaskId, TurboTasksBackendApi};
+use turbo_tasks::{util::SharedError, RawVc, TaskId, TaskIdSet, TurboTasksBackendApi};
 
 use crate::MemoryBackend;
 
 #[derive(Default, Debug)]
 pub struct Output {
     pub(crate) content: OutputContent,
     updates: u32,
-    pub(crate) dependent_tasks: AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2>,
+    pub(crate) dependent_tasks: TaskIdSet,
 }
 
 #[derive(Clone, Debug, Default)]
@@ -97,7 +95,7 @@ impl Output {
         }
     }
 
-    pub fn dependent_tasks(&self) -> &AutoSet<TaskId, BuildNoHashHasher<TaskId>, 2> {
+    pub fn dependent_tasks(&self) -> &TaskIdSet {
         &self.dependent_tasks
     }