From c9815c4b0c5b8151e4aa1eb178a3db673e5d7422 Mon Sep 17 00:00:00 2001
From: Arnout Roemers <1654946+aroemers@users.noreply.github.com>
Date: Sun, 20 Oct 2024 21:25:11 +0200
Subject: [PATCH] Add up-to extension

---
 CHANGELOG.md                     |  7 +++++
 README.md                        |  2 +-
 src/redelay/core.clj             | 11 ++++----
 src/redelay/extensions/up_to.clj | 46 ++++++++++++++++++++++++++++++++
 test/redelay/core_test.clj       |  4 ++-
 5 files changed, 63 insertions(+), 7 deletions(-)
 create mode 100644 src/redelay/extensions/up_to.clj

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ba1e21e..353781b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,12 @@
 # Change Log
 
+## Unreleased
+
+### Added
+
+- The `watchpoint` now receives notifications when a state is referenced (i.e. `deref` or `force`) as well, with the third argument set to `:referred`.
+- A small `up-to` extension has been added, allowing one to stop a state and its dependencies, in reverse order.
+
 ## 2.0.0
 
 ### Added
diff --git a/README.md b/README.md
index 85aa432..8544c7b 100644
--- a/README.md
+++ b/README.md
@@ -218,7 +218,7 @@ Those two functions are actually implemented using the library's extension point
 
 The library has a public `watchpoint` var.
 You can watch this var by using Clojure's `add-watch`.
-The registered watch functions receive `:starting`, `:started`, `:stopping` or `:stopped` and the State object.
+The registered watch functions receive `:referred`, `:starting`, `:started`, `:stopping` or `:stopped` and the State object.
 
 Try the following example:
 
diff --git a/src/redelay/core.clj b/src/redelay/core.clj
index af7b622..fcac3e6 100644
--- a/src/redelay/core.clj
+++ b/src/redelay/core.clj
@@ -6,11 +6,10 @@
 
 ;;; Watchpoint.
 
-(defonce ^{:doc "Add watches to this var to be notified of state changes, using
-  `add-watch`. The third argument to the watch fn will be one of
-  `:starting`, `:started`, `:stopping` or `:stopped`. The fourth
-   argument is the State object."}
-  watchpoint
+(defonce ^{:doc "Add watches to this var to be notified of state
+  changes, using `add-watch`. The third argument to the watch fn will
+  be one of `:referring`, `:starting`, `:started`, `:stopping` or
+  `:stopped`. The fourth argument is the State object."}  watchpoint
   (var watchpoint))
 
 
@@ -24,6 +23,7 @@
 (deftype State [name start-fn stop-fn value meta]
   clojure.lang.IDeref
   (deref [this]
+    (.notifyWatches watchpoint :referring this)
     (when-not (realized? this)
       (locking this
         (when-not (realized? this)
@@ -33,6 +33,7 @@
               (reset! value result)
               (.notifyWatches watchpoint :started this))
             (catch Exception e
+              (.notifyWatches watchpoint :aborted this)
               (throw (ex-info "Exception thrown when starting state" {:state this} e)))))))
     @value)
 
diff --git a/src/redelay/extensions/up_to.clj b/src/redelay/extensions/up_to.clj
new file mode 100644
index 0000000..0c9889d
--- /dev/null
+++ b/src/redelay/extensions/up_to.clj
@@ -0,0 +1,46 @@
+(ns redelay.extensions.up-to
+  "A small extension that allows stopping a state and its dependents (in
+  reverse order).
+
+  Make sure this namespace is loaded before realizing states, in order
+  to hook into the core's extension point (`watchpoint`). The internal
+  dependency graph is determined by watching which states are referred
+  to while starting a state."
+  (:require [redelay.core :as core]))
+
+;;; Internals
+
+(defonce ^:private closeables (atom {}))
+(defonce ^:private starting (atom ()))
+
+(defn- watch [_ _ change state]
+  (case change
+    :referring (when-let [start (peek @starting)]
+                 (swap! closeables update start (fnil conj (hash-set)) state))
+    :starting  (swap! starting conj state)
+    :aborted   (reset! starting ())
+    :started   (swap! starting pop)
+    :stopped   (swap! closeables dissoc state)
+    nil))
+
+(add-watch core/watchpoint ::up-to watch)
+
+(defn- transitive [state]
+  (let [dependencies @closeables]
+    (loop [todo (list state) result (list)]
+      (if-let [head (first todo)]
+        (let [deps (keep (fn [[state deps]] (when (contains? deps head) state)) dependencies)]
+          (recur (reduce conj (pop todo) deps)
+                 (cons head result)))
+        (distinct result)))))
+
+;;; Extension API
+
+(defn stop
+  "Stop the state and its (realized) dependents, in reverse order."
+  [state]
+  (when (realized? state)
+    (let [ordered (transitive state)]
+      (doseq [closeable ordered]
+        (.close closeable))
+      ordered)))
diff --git a/test/redelay/core_test.clj b/test/redelay/core_test.clj
index 853c422..8acd171 100644
--- a/test/redelay/core_test.clj
+++ b/test/redelay/core_test.clj
@@ -90,7 +90,9 @@
     (try
       (is (= 42 @forty-two))
       (stop)
-      (is (= [[::test watchpoint :starting forty-two]
+      (is (= [[::test watchpoint :referring forty-two]
+              [::test watchpoint :starting forty-two]
+              [::test watchpoint :referring two]
               [::test watchpoint :starting two]
               [::test watchpoint :started two]
               [::test watchpoint :started forty-two]