From 321fd4f45ed3a4c18334e353c9fef64f6a61e566 Mon Sep 17 00:00:00 2001
From: Chris Truter <chris@metabase.com>
Date: Wed, 28 Feb 2024 11:52:28 -0800
Subject: [PATCH] Do visit deduplication in our walk helpers

---
 src/macaw/rewrite.clj |  8 +++-----
 src/macaw/walk.clj    | 13 +++++++++++--
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/macaw/rewrite.clj b/src/macaw/rewrite.clj
index 42aae56..656fd47 100644
--- a/src/macaw/rewrite.clj
+++ b/src/macaw/rewrite.clj
@@ -55,12 +55,10 @@
         replace-name (fn [->s]
                        (fn [acc ^ASTNodeAccess visitable]
                          (let [node (.getASTNode visitable)]
-                           ;; work around ast walker repeatedly visiting the same expressions (bug ?!)
-                           ;; also not sure why sometimes we get a phantom visitable without an underlying node
-                           (if (or (nil? node) (contains? @seen visitable))
+                           ;; not sure why sometimes we get a phantom visitable without an underlying node
+                           (if (nil? node)
                              acc
-                             (do (vswap! seen conj visitable)
-                                 (conj acc [(node->idx-range node sql) (->s visitable)]))))))]
+                             (conj acc [(node->idx-range node sql) (->s visitable)])))))]
     (splice-replacements
      sql
      (mw/fold-query
diff --git a/src/macaw/walk.clj b/src/macaw/walk.clj
index ef59d13..bf915a6 100644
--- a/src/macaw/walk.clj
+++ b/src/macaw/walk.clj
@@ -17,17 +17,26 @@
     (f v)
     acc))
 
+;; work around ast walker repeatedly visiting the same expressions (bug ?!)
+(defn- deduplicate-visits [f]
+  (let [seen (volatile! #{})]
+    (fn [acc visitable]
+      (if (contains? @seen visitable)
+        acc
+        (do (vswap! seen conj visitable)
+            (f acc visitable))))))
+
 (defn- update-keys-vals [m key-f val-f]
   (into {} (map (fn [[k v]] [(key-f k) (val-f v)])) m))
 
 (defn walk-query
   "Walk over the query's AST, using the callbacks for their side-effects, for example to mutate the AST itself."
   [parsed-query callbacks]
-  (let [callbacks (update-keys-vals callbacks ->callback-key preserve)]
+  (let [callbacks (update-keys-vals callbacks ->callback-key (comp deduplicate-visits preserve))]
     (.walk (AstWalker. callbacks ::ignored) parsed-query)))
 
 (defn fold-query
   "Fold over the query's AST, using the callbacks to update the accumulator."
   [parsed-query callbacks init-val]
-  (let [callbacks (update-keys callbacks ->callback-key)]
+  (let [callbacks (update-keys-vals callbacks ->callback-key deduplicate-visits)]
     (.fold (AstWalker. callbacks init-val) parsed-query)))