Fix iced-refactor-thread-first/last to work with macros

src/iced/nrepl/refactor/thread.clj

@@ -1,13 +1,17 @@
 (ns iced.nrepl.refactor.thread
   (:require [clojure.string :as str]))
 
+(def ^:private replace-prefix "__ICED__")
+
 (defmulti expand-sexp (fn [sym _] sym))
 
 (defmethod expand-sexp '->
   [_ x]
   (loop [x x, expanded []]
     (let [[head target & tail] (if (sequential? x) x [])]
-      (if target
+      (if (and target
+               (not (vector? x))
+               (not (str/starts-with? (str head) replace-prefix)))
         (recur target (conj expanded {:head head :tail tail}))
         (conj expanded {:value x})))))
 
@@ -17,7 +21,9 @@
     (let [[head & tail] (if (sequential? x) x [])
           target (last tail)
           tail (drop-last tail)]
-      (if target
+      (if (and target
+               (not (vector? x))
+               (not (str/starts-with? (str head) replace-prefix)))
         (recur target (conj expanded {:head head :tail tail}))
         (conj expanded {:value x})))))
 
@@ -30,13 +36,37 @@
                     [(:value value)] bodies))
       (:value value))))
 
+(defn- lambda-replace-pair [code]
+  (reduce (fn [res x]
+            (assoc res x (str "(" replace-prefix (subs x 2))))
+          {} (re-seq #"#\([^ )]+" code)))
+
+(defn- deref-replace-pair [code]
+  (reduce (fn [res x]
+            (assoc res x (if (str/starts-with? x "@(")
+                           (str "(" replace-prefix (subs x 2))
+                           (str replace-prefix (subs x 1)))))
+          {} (re-seq #"@[^ )]+" code)))
+
+(defn- apply-replace-pairs [code pairs]
+  (reduce (fn [res [before after]] (str/replace res before after))
+          code pairs))
+
+(defn- rollback-replace-pairs [code pairs]
+  (reduce (fn [res [before after]] (str/replace res after before))
+          code pairs))
+
 (defn- thread* [sym code]
-  (let [sexp (read-string code)
+  (let [replace-pairs (merge (lambda-replace-pair code)
+                             (deref-replace-pair code))
+        code' (apply-replace-pairs code replace-pairs)
+        sexp (read-string code')
         expanded (expand-sexp sym sexp)]
     (if (> (count expanded) 2)
       (-> (construct sym expanded)
           str
-          (str/replace "," ""))
+          (str/replace "," "")
+          (rollback-replace-pairs replace-pairs))
       code)))
 
 (def thread-first (partial thread* '->))

test/iced/nrepl/refactor/thread_test.clj

@@ -4,36 +4,70 @@
 
 (t/deftest thread-first-test
   (t/are [in out] (= out (sut/thread-first in))
-    "(foo (bar baz))" "(-> baz bar foo)"
-    "(foo (bar baz x))" "(-> baz (bar x) foo)"
+    "(foo (bar baz))"     "(-> baz bar foo)"
+    "(foo (bar baz x))"   "(-> baz (bar x) foo)"
     "(foo (bar baz x) y)" "(-> baz (bar x) (foo y))"
-    "(foo (bar (baz)))" "(-> (baz) bar foo)"
-    "(foo)" "(foo)"
-    "foo" "foo"
-    "(foo bar)" "(foo bar)"
-    "(foo bar baz)" "(foo bar baz)"))
+    "(foo (bar (baz)))"   "(-> (baz) bar foo)"
+    "(foo)"               "(foo)"
+    "foo"                 "foo"
+    "(foo bar)"           "(foo bar)"
+    "(foo bar baz)"       "(foo bar baz)"))
 
 (t/deftest thread-first-with-map-test
   (t/is (= "(-> {:a 1 :b 2} bar foo)"
            (sut/thread-first "(foo (bar {:a 1 :b 2}))"))))
 
+(t/deftest thread-first-with-vector-test
+  (t/is (= "(-> [1 2 3] (bar baz) foo)"
+           (sut/thread-first "(foo (bar [1 2 3] baz))"))))
+
+(t/deftest thread-first-with-macro-test
+  (t/are [in out] (= out (sut/thread-first in))
+    "(foo (bar 123 #(baz %)))"      "(-> 123 (bar #(baz %)) foo)"
+    "(foo (bar 1 #(bar 2 %)))"      "(-> 1 (bar #(bar 2 %)) foo)"
+    "(foo (bar #(baz %) 123))"      "(-> #(baz %) (bar 123) foo)"
+    "(foo (bar 1 #(b %)) #(bar %))" "(-> 1 (bar #(b %)) (foo #(bar %)))"
+    "(foo (bar 1 #(bar %)) #(b %))" "(-> 1 (bar #(bar %)) (foo #(b %)))"))
+
+(t/deftest thread-first-with-deref-test
+  (t/are [in out] (= out (sut/thread-first in))
+    "(foo (bar @baz))"   "(-> @baz bar foo)"
+    "(foo (bar @(baz)))" "(-> @(baz) bar foo)"))
+
 (t/deftest thread-first-failure-test
   (t/is (thrown? Exception (sut/thread-first ""))))
 
 (t/deftest thread-last-test
   (t/are [in out] (= out (sut/thread-last in))
-    "(foo (bar baz))" "(->> baz bar foo)"
-    "(foo (bar x baz))" "(->> baz (bar x) foo)"
+    "(foo (bar baz))"     "(->> baz bar foo)"
+    "(foo (bar x baz))"   "(->> baz (bar x) foo)"
     "(foo y (bar x baz))" "(->> baz (bar x) (foo y))"
-    "(foo (bar (baz)))" "(->> (baz) bar foo)"
-    "(foo)" "(foo)"
-    "foo" "foo"
-    "(foo bar)" "(foo bar)"
-    "(foo bar baz)" "(foo bar baz)"))
+    "(foo (bar (baz)))"   "(->> (baz) bar foo)"
+    "(foo)"               "(foo)"
+    "foo"                 "foo"
+    "(foo bar)"           "(foo bar)"
+    "(foo bar baz)"       "(foo bar baz)"))
 
 (t/deftest thread-last-with-map-test
   (t/is (= "(->> {:a 1 :b 2} bar foo)"
            (sut/thread-last "(foo (bar {:a 1 :b 2}))"))))
 
+(t/deftest thread-last-with-vector-test
+  (t/is (= "(->> [1 2 3] (bar baz) foo)"
+           (sut/thread-last "(foo (bar baz [1 2 3]))"))))
+
+(t/deftest thread-last-with-macro-test
+  (t/are [in out] (= out (sut/thread-last in))
+    "(foo (bar #(baz %) 123))" "(->> 123 (bar #(baz %)) foo)"
+    "(foo (bar #(bar %) 123))" "(->> 123 (bar #(bar %)) foo)"
+    "(foo (bar 123 #(baz %)))" "(->> #(baz %) (bar 123) foo)"
+    "(foo (bar #(+ 1 %) 23))"  "(->> 23 (bar #(+ 1 %)) foo)"
+    "(foo (bar #(bar %) 123))" "(->> 123 (bar #(bar %)) foo)"))
+
+(t/deftest thread-last-with-deref-test
+  (t/are [in out] (= out (sut/thread-last in))
+    "(foo (bar @baz))"   "(->> @baz bar foo)"
+    "(foo (bar @(baz)))" "(->> @(baz) bar foo)"))
+
 (t/deftest thread-last-failure-test
   (t/is (thrown? Exception (sut/thread-last ""))))