From 53cb3a4429cbce6417fbf1dc07c87f108f8453f2 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Fri, 6 Jan 2023 11:33:50 -0800
Subject: [PATCH] Remove GC step for callback collections, always drop
 callbacks asap

---
 crates/gpui/src/app.rs                     | 10 ----
 crates/gpui/src/app/callback_collection.rs | 65 ++++++++++++++--------
 2 files changed, 42 insertions(+), 33 deletions(-)

diff --git a/crates/gpui/src/app.rs b/crates/gpui/src/app.rs
index 4a5b6f54be..5568155cf7 100644
--- a/crates/gpui/src/app.rs
+++ b/crates/gpui/src/app.rs
@@ -2123,16 +2123,6 @@ impl MutableAppContext {
                     self.pending_notifications.clear();
                     self.remove_dropped_entities();
                 } else {
-                    self.focus_observations.gc();
-                    self.global_subscriptions.gc();
-                    self.global_observations.gc();
-                    self.subscriptions.gc();
-                    self.observations.gc();
-                    self.window_activation_observations.gc();
-                    self.window_fullscreen_observations.gc();
-                    self.keystroke_observations.gc();
-                    self.release_observations.gc();
-
                     self.remove_dropped_entities();
 
                     if refreshing {
diff --git a/crates/gpui/src/app/callback_collection.rs b/crates/gpui/src/app/callback_collection.rs
index 44db224967..45479eabad 100644
--- a/crates/gpui/src/app/callback_collection.rs
+++ b/crates/gpui/src/app/callback_collection.rs
@@ -16,7 +16,17 @@ pub struct Subscription<K: Clone + Hash + Eq, F> {
 
 struct Mapping<K, F> {
     callbacks: HashMap<K, BTreeMap<usize, F>>,
-    dropped_subscriptions: HashSet<(K, usize)>,
+    dropped_subscriptions: HashMap<K, HashSet<usize>>,
+}
+
+impl<K: Hash + Eq, F> Mapping<K, F> {
+    fn clear_dropped_state(&mut self, key: &K, subscription_id: usize) -> bool {
+        if let Some(subscriptions) = self.dropped_subscriptions.get_mut(&key) {
+            subscriptions.remove(&subscription_id)
+        } else {
+            false
+        }
+    }
 }
 
 impl<K, F> Default for Mapping<K, F> {
@@ -50,6 +60,14 @@ impl<K: Clone + Hash + Eq + Copy, F> CallbackCollection<K, F> {
         self.internal.lock().callbacks.is_empty()
     }
 
+    pub fn count(&self, key: K) -> usize {
+        self.internal
+            .lock()
+            .callbacks
+            .get(&key)
+            .map_or(0, |callbacks| callbacks.len())
+    }
+
     pub fn subscribe(&mut self, key: K, subscription_id: usize) -> Subscription<K, F> {
         Subscription {
             key,
@@ -63,7 +81,7 @@ impl<K: Clone + Hash + Eq + Copy, F> CallbackCollection<K, F> {
 
         // If this callback's subscription was dropped before the callback was
         // added, then just drop the callback.
-        if this.dropped_subscriptions.remove(&(key, subscription_id)) {
+        if this.clear_dropped_state(&key, subscription_id) {
             return;
         }
 
@@ -74,10 +92,12 @@ impl<K: Clone + Hash + Eq + Copy, F> CallbackCollection<K, F> {
     }
 
     pub fn remove(&mut self, key: K) {
-        let callbacks = self.internal.lock().callbacks.remove(&key);
-
         // Drop these callbacks after releasing the lock, in case one of them
         // owns a subscription to this callback collection.
+        let mut this = self.internal.lock();
+        let callbacks = this.callbacks.remove(&key);
+        this.dropped_subscriptions.remove(&key);
+        drop(this);
         drop(callbacks);
     }
 
@@ -91,29 +111,26 @@ impl<K: Clone + Hash + Eq + Copy, F> CallbackCollection<K, F> {
         if let Some(callbacks) = callbacks {
             for (subscription_id, mut callback) in callbacks {
                 // If this callback's subscription was dropped while invoking an
-                // earlier callback, then just drop this callback.
-                if self
-                    .internal
-                    .lock()
-                    .dropped_subscriptions
-                    .remove(&(key, subscription_id))
-                {
+                // earlier callback, then just drop the callback.
+                let mut this = self.internal.lock();
+                if this.clear_dropped_state(&key, subscription_id) {
                     continue;
                 }
 
-                if call_callback(&mut callback, cx) {
-                    self.add_callback(key, subscription_id, callback);
+                drop(this);
+                let alive = call_callback(&mut callback, cx);
+
+                // If this callback's subscription was dropped while invoking the callback
+                // itself, or if the callback returns false, then just drop the callback.
+                let mut this = self.internal.lock();
+                if this.clear_dropped_state(&key, subscription_id) || !alive {
+                    continue;
                 }
-            }
-        }
-    }
 
-    pub fn gc(&mut self) {
-        let mut this = self.internal.lock();
-
-        for (key, id) in std::mem::take(&mut this.dropped_subscriptions) {
-            if let Some(callbacks) = this.callbacks.get_mut(&key) {
-                callbacks.remove(&id);
+                this.callbacks
+                    .entry(key)
+                    .or_default()
+                    .insert(subscription_id, callback);
             }
         }
     }
@@ -150,7 +167,9 @@ impl<K: Clone + Hash + Eq, F> Drop for Subscription<K, F> {
             // removed later.
             mapping
                 .dropped_subscriptions
-                .insert((self.key.clone(), self.id));
+                .entry(self.key.clone())
+                .or_default()
+                .insert(self.id);
         }
     }
 }