From 8cb117ce4623f8e9abe59cf018f302c7487e476f Mon Sep 17 00:00:00 2001
From: Christian Paul Dehli <cp@deh.li>
Date: Thu, 11 Jul 2024 15:58:48 -0400
Subject: [PATCH] fix unmount bug

---
 core/src/refx/hooks.cljs      |  2 +-
 core/src/refx/subs.cljc       | 37 ++++++++++++++++++++++++-----------
 core/test/refx/subs_test.cljs | 12 ++++++------
 3 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/core/src/refx/hooks.cljs b/core/src/refx/hooks.cljs
index 5db5a23..cf0eab9 100644
--- a/core/src/refx/hooks.cljs
+++ b/core/src/refx/hooks.cljs
@@ -28,7 +28,7 @@
         (react/useMemo
          (fn []
            [(fn [callback]
-              (let [key {:index (swap! use-sub-counter inc)}]
+              (let [key {::subs/index (swap! use-sub-counter inc)}]
                 (subs/-add-listener sub key callback)
                 #(subs/-remove-listener sub key)))
             (fn []
diff --git a/core/src/refx/subs.cljc b/core/src/refx/subs.cljc
index 0f6264f..bb3d09d 100644
--- a/core/src/refx/subs.cljc
+++ b/core/src/refx/subs.cljc
@@ -5,6 +5,13 @@
 
 (def kind :sub)
 
+(defn ui-sub?
+  "UI subscriptions are batched together and invoked after all subs
+  finish computing their new values. This is done to reduce renders and
+  prevent consumers from having an inconsistent view of the DB."
+  [sub-key]
+  (some? (::index sub-key)))
+
 ;; --- signals ----------------------------------------------------------------
 
 (defprotocol ISignal
@@ -78,7 +85,7 @@
 
 (defonce ^:private listeners-state
   (letfn [(comparator [a b]
-            (compare (:index a) (:index b)))]
+            (compare (::index a) (::index b)))]
     (atom {:counter 0 :pending (sorted-map-by comparator)})))
 
 (defn- invoke-listener
@@ -103,16 +110,20 @@
 
     (interop/next-tick
      (fn []
-       (let [listener-fns (atom nil)]
-         (swap! listeners-state (fn [state]
-                                  (let [{:keys [counter pending] :as new-state}
-                                        (update state :counter dec)]
-                                    (if (zero? counter)
-                                      (do (reset! listener-fns pending)
-                                          (update new-state :pending empty))
-                                      new-state))))
-         (doseq [[_ f] @listener-fns]
-           (f)))))))
+       (let [{:keys [counter pending]}
+             (swap! listeners-state update :counter dec)]
+
+         (when (zero? counter)
+           (doseq [[listener-key _] pending]
+             ;; Triggering a listener-fn can result in a subsequent sub's
+             ;; remove-listener to be called (which will remove it from pending).
+             ;; This check ensure it's still pending.
+             (let [listener-fn (atom nil)]
+               (swap! listeners-state (fn [state]
+                                        (reset! listener-fn (get-in state [:pending listener-key]))
+                                        (update state :pending dissoc listener-key)))
+               (when-let [f @listener-fn]
+                 (f))))))))))
 
 (deftype Listeners [^:mutable listeners]
   Object
@@ -160,6 +171,8 @@
   (-add-listener [_ k f]
     (.add listeners k f))
   (-remove-listener [this k]
+    (when (ui-sub? k)
+      (swap! listeners-state update :pending dissoc k))
     (.remove listeners k)
     (when (.empty? listeners)
       (sub-orphaned this)))
@@ -238,6 +251,8 @@
   (-add-listener [_ k f]
     (.add listeners k f))
   (-remove-listener [this k]
+    (when (ui-sub? k)
+      (swap! listeners-state update :pending dissoc k))
     (.remove listeners k)
     (when (.empty? listeners)
       (sub-orphaned this)))
diff --git a/core/test/refx/subs_test.cljs b/core/test/refx/subs_test.cljs
index 83c2371..ca6cc20 100644
--- a/core/test/refx/subs_test.cljs
+++ b/core/test/refx/subs_test.cljs
@@ -119,7 +119,7 @@
           listener-calls (atom [])
           remove-listener-fns (atom '())
           add-listener! (fn [sub]
-                          (let [key {:index (swap! listener-count inc)}]
+                          (let [key {::subs/index (swap! listener-count inc)}]
                             (subs/-add-listener sub key #(swap! listener-calls conj key))
                             (swap! remove-listener-fns conj #(subs/-remove-listener sub key))))
           remove-listeners! (fn []
@@ -130,13 +130,13 @@
       (add-listener! sub-c)
       (add-listener! sub-d)
       (reset! source 1)
-      (remove-listeners!)
       (async done
              (js/setTimeout (fn []
+                              (remove-listeners!)
                               (is (= @listener-calls
-                                     [{:index 1}
-                                      {:index 2}
-                                      {:index 3}
-                                      {:index 4}]))
+                                     [{::subs/index 1}
+                                      {::subs/index 2}
+                                      {::subs/index 3}
+                                      {::subs/index 4}]))
                               (done))
                             10)))))