From 803603217139cce6a3807d0c76511de6b688a5d4 Mon Sep 17 00:00:00 2001
From: David Nolen <david.nolen@gmail.com>
Date: Thu, 16 Nov 2023 11:46:11 -0500
Subject: [PATCH] CLJS-3406 implement reset-vals! and swap-vals! through
 protocols (#216)

Co-authored-by: Nikita Prokopov <niki@tonsky.me>
---
 src/main/cljs/cljs/core.cljs      | 36 ++++++++++++++++---------
 src/test/cljs/cljs/core_test.cljs | 44 +++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 13 deletions(-)

diff --git a/src/main/cljs/cljs/core.cljs b/src/main/cljs/cljs/core.cljs
index 620bc8abc..baa502dab 100644
--- a/src/main/cljs/cljs/core.cljs
+++ b/src/main/cljs/cljs/core.cljs
@@ -4569,15 +4569,17 @@ reduces them without incurring seq initialization"
    atom before and after the reset."
   {:added "1.9"}
   [a new-value]
-  (let [validate (.-validator a)]
-    (when-not (nil? validate)
-      (when-not (validate new-value)
-        (throw (js/Error. "Validator rejected reference state"))))
-    (let [old-value (.-state a)]
-      (set! (.-state a) new-value)
-      (when-not (nil? (.-watches a))
-        (-notify-watches a old-value new-value))
-      [old-value new-value])))
+  (if (instance? Atom a)
+    (let [validate (.-validator a)]
+      (when-not (nil? validate)
+        (when-not (validate new-value)
+          (throw (js/Error. "Validator rejected reference state"))))
+      (let [old-value (.-state a)]
+        (set! (.-state a) new-value)
+        (when-not (nil? (.-watches a))
+          (-notify-watches a old-value new-value))
+        [old-value new-value]))
+    [(-deref a) (-reset! a new-value)]))
 
 (defn swap!
   "Atomically swaps the value of atom to be:
@@ -4608,13 +4610,21 @@ reduces them without incurring seq initialization"
   Returns [old new], the value of the atom before and after the swap."
   {:added "1.9"}
   ([a f]
-   (reset-vals! a (f (.-state a))))
+   (if (instance? Atom a)
+     (reset-vals! a (f (.-state a)))
+     [(-deref a) (-swap! a f)]))
   ([a f x]
-   (reset-vals! a (f (.-state a) x)))
+   (if (instance? Atom a)
+     (reset-vals! a (f (.-state a) x))
+     [(-deref a) (-swap! a f x)]))
   ([a f x y]
-   (reset-vals! a (f (.-state a) x y)))
+   (if (instance? Atom a)
+     (reset-vals! a (f (.-state a) x y))
+     [(-deref a) (-swap! a f x y)]))
   ([a f x y & more]
-   (reset-vals! a (apply f (.-state a) x y more))))
+   (if (instance? Atom a)
+     (reset-vals! a (apply f (.-state a) x y more))
+     [(-deref a) (-swap! a f x y more)])))
 
 (defn compare-and-set!
   "Atomically sets the value of atom to newval if and only if the
diff --git a/src/test/cljs/cljs/core_test.cljs b/src/test/cljs/cljs/core_test.cljs
index 812d3a74a..625c14c3f 100644
--- a/src/test/cljs/cljs/core_test.cljs
+++ b/src/test/cljs/cljs/core_test.cljs
@@ -1998,3 +1998,47 @@
     (let [a #js {}]
       (set! a -x false)
       (is (false? (.-x a))))))
+
+(deftest test-cljs-3406
+  (testing "ISwap/IReset protocols"
+    (let [a (atom {:x 0})
+          c (reify
+              IDeref
+              (-deref [_]
+                (:x @a))
+              
+              ISwap
+              (-swap! [o f]
+                (:x (swap! a update :x f)))
+              (-swap! [o f x]
+                (:x (swap! a update :x f x)))
+              (-swap! [o f x y]
+                (:x (swap! a update :x f x y)))
+              (-swap! [o f x y zs]
+                (:x (swap! a #(apply update % :x f x y zs))))
+              
+              IReset
+              (-reset! [o new-value]
+                (:x (swap! a assoc :x new-value))))]
+      (is (= 0 @c))
+      (is (= 1 (swap! c inc)))
+      (is (= 1 @c))
+      (is (= 2 (swap! c + 1)))
+      (is (= 2 @c))
+      (is (= 5 (swap! c + 1 2)))
+      (is (= 5 @c))
+      (is (= 11 (swap! c + 1 2 3)))
+      (is (= 11 @c))
+      (is (= 0 (reset! c 0)))
+      (is (= 0 @c))
+      
+      (is (= [0 1] (swap-vals! c inc)))
+      (is (= 1 @c))
+      (is (= [1 2] (swap-vals! c + 1)))
+      (is (= 2 @c))
+      (is (= [2 5] (swap-vals! c + 1 2)))
+      (is (= 5 @c))
+      (is (= [5 11] (swap-vals! c + 1 2 3)))
+      (is (= 11 @c))
+      (is (= [11 0] (reset-vals! c 0)))
+      (is (= 0 @c)))))