From 45166d8e6b586d2d50ab219b86b219208bdb2c6f Mon Sep 17 00:00:00 2001
From: Tobias Koppers <tobias.koppers@googlemail.com>
Date: Mon, 5 Dec 2022 14:32:24 +0100
Subject: [PATCH] make sure that removing collectibles will cleanup helper
 entries (vercel/turbo#2872)

---
 .../turbo-tasks-memory/src/count_hash_set.rs  |  9 +++
 crates/turbo-tasks-memory/src/scope.rs        | 73 ++++++++++++++-----
 2 files changed, 64 insertions(+), 18 deletions(-)

diff --git a/crates/turbo-tasks-memory/src/count_hash_set.rs b/crates/turbo-tasks-memory/src/count_hash_set.rs
index 1afb13024481b..8f089cf3fe454 100644
--- a/crates/turbo-tasks-memory/src/count_hash_set.rs
+++ b/crates/turbo-tasks-memory/src/count_hash_set.rs
@@ -51,13 +51,22 @@ impl<T, H: Default> CountHashSet<T, H> {
 }
 
 impl<T, H> CountHashSet<T, H> {
+    /// Get the number of positive entries
     pub fn len(&self) -> usize {
         self.inner.len() - self.negative_entries
     }
 
+    /// Checks if the set looks empty from outside. It might still have negative
+    /// entries, but they should be treated as not existing.
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
+
+    /// Checks if this set is equal to a fresh created set, meaning it has no
+    /// positive but also no negative entries.
+    pub fn is_unset(&self) -> bool {
+        self.inner.is_empty()
+    }
 }
 
 impl<T: Eq + Hash, H: BuildHasher + Default> CountHashSet<T, H> {
diff --git a/crates/turbo-tasks-memory/src/scope.rs b/crates/turbo-tasks-memory/src/scope.rs
index da8c77fbaa2e5..22fa0c5fd9457 100644
--- a/crates/turbo-tasks-memory/src/scope.rs
+++ b/crates/turbo-tasks-memory/src/scope.rs
@@ -7,7 +7,7 @@ use std::{
     sync::atomic::{AtomicIsize, AtomicUsize, Ordering},
 };
 
-use auto_hash_map::{AutoMap, AutoSet};
+use auto_hash_map::{map::Entry, AutoMap, AutoSet};
 use nohash_hasher::BuildNoHashHasher;
 use parking_lot::Mutex;
 use turbo_tasks::{
@@ -387,8 +387,12 @@ impl TaskScope {
         reader: TaskId,
     ) {
         let mut state = self.state.lock();
-        if let Some((_, dependent_tasks)) = state.collectibles.get_mut(&trait_type) {
+        if let Entry::Occupied(mut entry) = state.collectibles.entry(trait_type) {
+            let (collectibles, dependent_tasks) = entry.get_mut();
             dependent_tasks.remove(&reader);
+            if collectibles.is_unset() && dependent_tasks.is_empty() {
+                entry.remove();
+            }
         }
     }
 }
@@ -535,14 +539,32 @@ impl TaskScopeState {
         collectible: RawVc,
         count: usize,
     ) -> Option<ScopeCollectibleChangeEffect> {
-        let (collectibles, dependent_tasks) = self.collectibles.entry(trait_id).or_default();
-        if collectibles.add_count(collectible, count) {
-            log_scope_update!("add_collectible {} -> {}", *self.id, collectible);
-            Some(ScopeCollectibleChangeEffect {
-                notify: take(dependent_tasks),
-            })
-        } else {
-            None
+        match self.collectibles.entry(trait_id) {
+            Entry::Occupied(mut entry) => {
+                let (collectibles, dependent_tasks) = entry.get_mut();
+                if collectibles.add_count(collectible, count) {
+                    log_scope_update!("add_collectible {} -> {}", *self.id, collectible);
+                    Some(ScopeCollectibleChangeEffect {
+                        notify: take(dependent_tasks),
+                    })
+                } else {
+                    if collectibles.is_unset() && dependent_tasks.is_empty() {
+                        entry.remove();
+                    }
+                    None
+                }
+            }
+            Entry::Vacant(entry) => {
+                let result = entry
+                    .insert(Default::default())
+                    .0
+                    .add_count(collectible, count);
+                debug_assert!(result, "this must be always a new entry");
+                log_scope_update!("add_collectible {} -> {}", *self.id, collectible);
+                Some(ScopeCollectibleChangeEffect {
+                    notify: AutoSet::new(),
+                })
+            }
         }
     }
 
@@ -568,14 +590,29 @@ impl TaskScopeState {
         collectible: RawVc,
         count: usize,
     ) -> Option<ScopeCollectibleChangeEffect> {
-        let (collectibles, dependent_tasks) = self.collectibles.entry(trait_id).or_default();
-        if collectibles.remove_count(collectible, count) {
-            log_scope_update!("remove_collectible {} -> {}", *self.id, collectible);
-            Some(ScopeCollectibleChangeEffect {
-                notify: take(dependent_tasks),
-            })
-        } else {
-            None
+        match self.collectibles.entry(trait_id) {
+            Entry::Occupied(mut entry) => {
+                let (collectibles, dependent_tasks) = entry.get_mut();
+                if collectibles.remove_count(collectible, count) {
+                    let notify = take(dependent_tasks);
+                    if collectibles.is_unset() {
+                        entry.remove();
+                    }
+                    log_scope_update!("remove_collectible {} -> {}", *self.id, collectible);
+                    Some(ScopeCollectibleChangeEffect { notify })
+                } else {
+                    None
+                }
+            }
+            Entry::Vacant(e) => {
+                let result = e
+                    .insert(Default::default())
+                    .0
+                    .remove_count(collectible, count);
+
+                debug_assert!(!result, "this must never be visible from outside");
+                None
+            }
         }
     }