From 20733bc4483533cc0b3e6dd16d92309afc27eeed Mon Sep 17 00:00:00 2001
From: James Reeves <jreeves@weavejester.com>
Date: Tue, 14 Mar 2023 15:52:59 +0000
Subject: [PATCH] Add :align-map-columns? option

Add the :align-map-columns? option. This will align the keys and values
of maps such that they line up in columns. Defaults to false.

This option only works on maps, not binding vectors (such as in 'let').

Partially solves #35. Closes #299. Closes #370.
---
 README.md                         |   4 +
 cljfmt/src/cljfmt/core.cljc       |  87 +++++++++++++++++
 cljfmt/test/cljfmt/core_test.cljc | 149 ++++++++++++++++++++++++++++++
 3 files changed, 240 insertions(+)

diff --git a/README.md b/README.md
index 73db8a7a..c92102e9 100644
--- a/README.md
+++ b/README.md
@@ -225,6 +225,10 @@ In order to load the standard configuration file from Leiningen, add the
 
 ### Formatting Options
 
+* `:align-map-columns?` -
+  true if cljfmt should align the keys and values of maps such that they
+  line up in columns. Defaults to false.
+
 * `:indentation?` -
   true if cljfmt should correct the indentation of your code.
   Defaults to true.
diff --git a/cljfmt/src/cljfmt/core.cljc b/cljfmt/src/cljfmt/core.cljc
index c8fb545f..ac1685c8 100644
--- a/cljfmt/src/cljfmt/core.cljc
+++ b/cljfmt/src/cljfmt/core.cljc
@@ -363,6 +363,7 @@
    :split-keypairs-over-multiple-lines?   false
    :sort-ns-references?                   false
    :function-arguments-indentation        :community
+   :align-map-columns?                    false
    :indents       default-indents
    :extra-indents {}
    :alias-map     {}})
@@ -572,6 +573,90 @@
 (defn sort-ns-references [form]
   (transform form edit-all ns-reference? sort-arguments))
 
+(defn- skip-to-linebreak-or-element [zloc]
+  (z/skip z/right* (some-fn space? comma?) zloc))
+
+(defn- reduce-columns [zloc f init]
+  (loop [zloc zloc, col 0, acc init]
+    (if-some [zloc (skip-to-linebreak-or-element zloc)]
+      (if (line-break? zloc)
+        (recur (z/right* zloc) 0 acc)
+        (recur (z/right* zloc) (inc col) (f zloc col acc)))
+      acc)))
+
+(defn- count-columns [zloc]
+  (inc (reduce-columns zloc #(max %2 %3) 0)))
+
+(defn- trailing-commas [zloc]
+  (let [right (z/right* zloc)]
+    (if (and right (comma? right))
+      (-> right z/node n/string)
+      "")))
+
+(defn- node-end-position [zloc]
+  (let [lines (str (prior-line-string zloc)
+                   (n/string (z/node zloc))
+                   (trailing-commas zloc))]
+    (transduce (comp (remove #(str/starts-with? % ";"))
+                     (map count))
+               max 0 (str/split lines #"\r?\n"))))
+
+(defn- max-column-end-position [zloc col]
+  (reduce-columns zloc
+                  (fn [zloc c max-pos]
+                    (if (= c col)
+                      (max max-pos (node-end-position zloc))
+                      max-pos))
+                  0))
+
+(defn- update-space-left [zloc delta]
+  (let [left (z/left* zloc)]
+    (cond
+      (space? left) (let [n (-> left z/node n/string count)]
+                      (z/right* (z/replace* left (n/spaces (+ n delta)))))
+      (pos? delta)  (z/insert-space-left zloc delta)
+      :else         zloc)))
+
+(defn- skip-to-next-line-in-form [zloc]
+  (z/right (z/skip z/right* (complement line-break?) (z/right* zloc))))
+
+(defn- pad-node [zloc start-position]
+  (let [padding (- start-position (margin zloc))
+        zloc    (update-space-left zloc padding)]
+    (if-some [zloc (z/down zloc)]
+      (loop [zloc zloc]
+        (if-some [zloc (skip-to-next-line-in-form zloc)]
+          (let [zloc (update-space-left zloc padding)]
+            (if-some [zloc (z/right* zloc)]
+              (recur zloc)
+              (z/up zloc)))
+          (z/up zloc)))
+      zloc)))
+
+(defn- edit-column [zloc column f]
+  (loop [zloc zloc, col 0]
+    (if-some [zloc (skip-to-linebreak-or-element zloc)]
+      (let [zloc (if (and (= col column) (not (line-break? zloc)))
+                   (f zloc)
+                   zloc)
+            col  (if (line-break? zloc) 0 (inc col))]
+        (if-some [zloc (z/right* zloc)]
+          (recur zloc col)
+          zloc))
+      zloc)))
+
+(defn- align-column [zloc col]
+  (if-some [zloc (z/down zloc)]
+    (let [start-position (inc (max-column-end-position zloc (dec col)))]
+      (z/up (edit-column zloc col #(pad-node % start-position))))
+    zloc))
+
+(defn- align-form-columns [zloc]
+  (reduce align-column zloc (-> zloc z/down count-columns range rest)))
+
+(defn align-map-columns [form]
+  (transform form edit-all z/map? align-form-columns))
+
 (defn reformat-form
   ([form]
    (reformat-form form {}))
@@ -594,6 +679,8 @@
            (reindent (merge (:indents opts) (:extra-indents opts))
                      (:alias-map opts)
                      opts))
+         (cond-> (:align-map-columns? opts)
+           align-map-columns)
          (cond-> (:remove-trailing-whitespace? opts)
            remove-trailing-whitespace)))))
 
diff --git a/cljfmt/test/cljfmt/core_test.cljc b/cljfmt/test/cljfmt/core_test.cljc
index ff022894..3576a4bf 100644
--- a/cljfmt/test/cljfmt/core_test.cljc
+++ b/cljfmt/test/cljfmt/core_test.cljc
@@ -2032,3 +2032,152 @@
         "                 \"Help\"))"
         " (str \"🔢 \" (str \"email\""
         "                 \"Leverage\"))]"])))
+
+(deftest test-align-map-columns
+  (testing "empty maps"
+    (is (reformats-to?
+         ["{}"]
+         ["{}"]
+         {:align-map-columns? true})))
+  (testing "basic aligning"
+    (is (reformats-to?
+         ["{:x 1"
+          " :longer 2}"]
+         ["{:x      1"
+          " :longer 2}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{:longer 1"
+          " :x 2}"]
+         ["{:longer 1"
+          " :x      2}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{:x 1 :longer 2}"]
+         ["{:x 1 :longer 2}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{:x 1 :y 2"
+          " :longer 2}"]
+         ["{:x      1 :y 2"
+          " :longer 2}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{:a 1 :b 2 :cc 3"
+          ":dd 4 :eee 5 :f 6"
+          ":ggg 7 :hh 8 :iii 9}"]
+         ["{:a   1 :b   2 :cc  3"
+          " :dd  4 :eee 5 :f   6"
+          " :ggg 7 :hh  8 :iii 9}"]
+         {:align-map-columns? true})))
+  (testing "wrong alignment"
+    (is (reformats-to?
+         ["{:f 1"
+          " :bar    2}"]
+         ["{:f   1"
+          " :bar 2}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{:foo     1"
+          " :b 2}"]
+         ["{:foo 1"
+          " :b   2}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{   :foo 1"
+          ":b       2}"]
+         ["{:foo 1"
+          " :b   2}"]
+         {:align-map-columns? true})))
+  (testing "commas"
+    (is (reformats-to?
+         ["{:a 1, :b 2, :cc 3"
+          ":dd 4, :eee 5, :f 6"
+          ":ggg 7, :hh 8, :iii 9}"]
+         ["{:a   1, :b   2, :cc  3"
+          " :dd  4, :eee 5, :f   6"
+          " :ggg 7, :hh  8, :iii 9}"]
+         {:align-map-columns? true})))
+  (testing "nested maps"
+    (is (reformats-to?
+         ["{:a {:b 1"
+          "     :c 2}"
+          " :ddd {:e 3}}"]
+         ["{:a   {:b 1"
+          "       :c 2}"
+          " :ddd {:e 3}}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{:aaa {:b 1"
+          "       :c 2}"
+          " :d {:e 3}}"]
+         ["{:aaa {:b 1"
+          "       :c 2}"
+          " :d   {:e 3}}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{{:a 1"
+          "  :b 2} 3"
+          " {:ccc 4} 5}"]
+         ["{{:a 1"
+          "  :b 2}   3"
+          " {:ccc 4} 5}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{{:a 1"
+          "  :b 2} 3"
+          " :c 5}"]
+         ["{{:a 1"
+          "  :b 2} 3"
+          " :c     5}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["{:a {:b 1"
+          " :c 2} :ddd 3"
+          " :eee {:ff 3} :e 4}"]
+         ["{:a   {:b 1"
+          "       :c 2}  :ddd 3"
+          " :eee {:ff 3} :e   4}"]
+         {:align-map-columns? true})))
+  (testing "nested forms"
+    (is (reformats-to?
+         ["{:x (let [x 1]"
+          "      (+ x 1))"
+          " :yyy (let [y 2]"
+          "        (+ y 2))}"]
+         ["{:x   (let [x 1]"
+          "        (+ x 1))"
+          " :yyy (let [y 2]"
+          "        (+ y 2))}"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["(def m {:x 1"
+          ":longer 2})"]
+         ["(def m {:x      1"
+          "        :longer 2})"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["(def m {{:a 1"
+          ":b 2} 3"
+          ":d 4})"]
+         ["(def m {{:a 1"
+          "         :b 2} 3"
+          "        :d     4})"]
+         {:align-map-columns? true}))
+    (is (reformats-to?
+         ["(def m {{:a 1"
+          ":b 2} [x"
+          "y]"
+          ":d [z]})"]
+         ["(def m {{:a 1"
+          "         :b 2} [x"
+          "                y]"
+          "        :d     [z]})"]
+         {:align-map-columns? true})))
+  (testing "comments"
+    (is (reformats-to?
+         ["{:x 1   ; a comment"
+          " :longer 2}"]
+         ["{:x      1   ; a comment"
+          " :longer 2}"]
+         {:align-map-columns? true}))))