From 6db2198a8f56f49ee2678b9f42a387073695ffac Mon Sep 17 00:00:00 2001
From: "reut.sharabani" <reut.sharabani@appsflyer.com>
Date: Wed, 22 Mar 2023 22:11:21 +0200
Subject: [PATCH] align maps

---
 README.md                         |  14 +++
 cljfmt/src/cljfmt/core.cljc       |  77 +++++++++++++
 cljfmt/test/cljfmt/core_test.cljc | 174 ++++++++++++++++++++++++++++++
 3 files changed, 265 insertions(+)

diff --git a/README.md b/README.md
index d5a0211..0425204 100644
--- a/README.md
+++ b/README.md
@@ -154,6 +154,20 @@ selectively enabled or disabled:
   other references in the `ns` forms at the top of your namespaces.
   Defaults to false.
 
+* `:align-maps?` -
+  true if cljfmt should left align the values of maps
+  This will convert:
+  ```clojure
+  {:foo 1
+   :barbaz 2}
+  ```
+  To:
+  ```clojure
+  {:foo    1
+   :barbaz 2}
+  ```
+  Defaults to false.
+
 You can also configure the behavior of cljfmt:
 
 * `:paths` - determines which directories to include in the
diff --git a/cljfmt/src/cljfmt/core.cljc b/cljfmt/src/cljfmt/core.cljc
index 70f3df3..1ecf23f 100644
--- a/cljfmt/src/cljfmt/core.cljc
+++ b/cljfmt/src/cljfmt/core.cljc
@@ -486,6 +486,80 @@
 (defn sort-ns-references [form]
   (transform form edit-all ns-reference? sort-arguments))
 
+(defn- node-width [zloc]
+  (-> zloc z/node n/string count))
+
+(defn- node-column [zloc]
+  (loop [zloc (z/left* zloc), n 0]
+    (if (or (nil? zloc) (line-break? zloc))
+      n
+      (recur (z/left* zloc)
+             (if (clojure-whitespace? zloc) n (inc n))))))
+
+(defn- group-separator? [zloc]
+  (= (z/string zloc) "\n\n"))
+
+(defn- node-group [zloc]
+  (loop [zloc (z/left* zloc), n 0]
+    (if (nil? zloc)
+      n
+      (recur (z/left* zloc)
+             (if (group-separator? zloc) (inc n) n)))))
+
+(defn- comma-after? [zloc]
+  (let [right (z/right* zloc)]
+    (or (comma? right)
+        (and (z/whitespace? right) (comma? (z/right* right))))))
+
+(defn- max-group-column-widths [zloc]
+  (loop [zloc (z/down zloc), max-widths {}]
+    (if (nil? zloc)
+      max-widths
+      (let [width  (if (comma-after? zloc)
+                     (inc (node-width zloc))
+                     (node-width zloc))
+            column (node-column zloc)
+            group  (node-group zloc)]
+        (recur (z/right zloc)
+               (update-in max-widths [group column] (fnil max 0) width))))))
+
+(defn- remove-space-right [zloc]
+  (let [right (z/right* zloc)]
+    (if (space? right) (z/remove* right) zloc)))
+
+(defn- insert-space-right [zloc n]
+  (let [right (z/right* zloc)]
+    (if (comma? right)
+      (insert-space-right (remove-space-right right) (dec n))
+      (z/insert-space-right zloc n))))
+
+(defn- set-spacing-right [zloc n]
+  (-> zloc (remove-space-right) (insert-space-right n)))
+
+(defn- map-children [zloc f]
+  (if-let [zloc (z/down zloc)]
+    (loop [zloc zloc]
+      (let [zloc (f zloc)]
+        (if-let [zloc (z/right zloc)]
+          (recur zloc)
+          (z/up zloc))))
+    zloc))
+
+(defn- pad-node [zloc width]
+  (set-spacing-right zloc (- width (node-width zloc))))
+
+(defn- end-of-line? [zloc]
+  (line-break? (skip-whitespace-and-commas (z/right* zloc))))
+
+(defn- align-form-columns [zloc]
+  (let [max-widths (max-group-column-widths zloc)]
+    (map-children zloc #(cond-> %
+                          (and (z/right %) (not (end-of-line? %)))
+                          (pad-node (inc (get-in max-widths [(node-group %) (node-column %)])))))))
+
+(defn align-maps [form]
+  (transform form edit-all z/map? align-form-columns))
+
 (def default-options
   {:indentation?                          true
    :insert-missing-whitespace?            true
@@ -495,6 +569,7 @@
    :remove-trailing-whitespace?           true
    :split-keypairs-over-multiple-lines?   false
    :sort-ns-references?                   false
+   :align-maps?                           false
    :indents   default-indents
    :alias-map {}})
 
@@ -516,6 +591,8 @@
            insert-missing-whitespace)
          (cond-> (:remove-multiple-non-indenting-spaces? opts)
            remove-multiple-non-indenting-spaces)
+         (cond-> (:align-maps? opts)
+           align-maps)
          (cond-> (:indentation? opts)
            (reindent (:indents opts) (:alias-map opts)))
          (cond-> (:remove-trailing-whitespace? opts)
diff --git a/cljfmt/test/cljfmt/core_test.cljc b/cljfmt/test/cljfmt/core_test.cljc
index 9bed12b..4aaa4bd 100644
--- a/cljfmt/test/cljfmt/core_test.cljc
+++ b/cljfmt/test/cljfmt/core_test.cljc
@@ -1336,3 +1336,177 @@
         "   ^{:x 1} b"
         "   [c]))"]
        {:sort-ns-references? true})))
+
+(deftest test-align-maps
+  (testing "straightforward test cases"
+    (testing "sanity"
+      (is (reformats-to?
+           ["(def x 1)"]
+           ["(def x 1)"]
+           {:align-maps? true})))
+    (testing "no op 1"
+      (is (reformats-to?
+           ["{:a 1}"]
+           ["{:a 1}"]
+           {:align-maps? true})))
+    (testing "no op 2"
+      (is (reformats-to?
+           ["{:a 1"
+            " :b 2}"]
+           ["{:a 1"
+            " :b 2}"]
+           {:align-maps? true})))
+    (testing "empty"
+      (is (reformats-to?
+           ["{}"]
+           ["{}"]
+           {:align-maps? true})))
+    (testing "simple"
+      (is (reformats-to?
+           ["{:x 1"
+            " :longer 2}"]
+           ["{:x      1"
+            " :longer 2}"]
+           {:align-maps? true})))
+    (testing "nested simple"
+      (is (reformats-to?
+           ["{:x {:x 1}"
+            " :longer 2}"]
+           ["{:x      {:x 1}"
+            " :longer 2}"]
+           {:align-maps? true})))
+    (testing "nested align"
+      (is (reformats-to?
+           ["{:x {:x 1"
+            "     :longer 2}"
+            " :longer 2}"]
+           ["{:x      {:x      1"
+            "          :longer 2}"
+            " :longer 2}"]
+           {:align-maps? true})))
+    (testing "align many"
+      (is (reformats-to?
+           ["{:a 1"
+            " :longer 2"
+            " :b 3}"]
+           ["{:a      1"
+            " :longer 2"
+            " :b      3}"]
+           {:align-maps? true})))
+    (testing "preserves comments"
+      (is (reformats-to?
+           ["{:a 1 ;; comment"
+            " :longer 2}"]
+           ["{:a      1 ;; comment"
+            " :longer 2}"]
+           {:align-maps? true}))))
+  (testing "non-trivial test cases"
+    (testing "idnentation after align"
+      (is (reformats-to?
+           ["(def m {{:a 1"
+            ":b 2} [x"
+            "y]"
+            ":d [z]})"]
+           ["(def m {{:a 1"
+            "         :b 2} [x"
+            "                y]"
+            "        :d [z]})"])))
+    (testing "cljs map values"
+      (is (reformats-to?
+           ["{:indents {'thing.core/defthing [[:inner 0]]"
+            "'let [[:inner 0]]}"
+            "#?@(:cljs [:alias-map {}])}"]
+           ["{:indents                   {'thing.core/defthing [[:inner 0]]"
+            "                             'let                 [[:inner 0]]}"
+            " #?@(:cljs [:alias-map {}])}"]
+           {:align-maps? true})))
+    (testing "indentation off #1"
+      (is (reformats-to?
+           ["{  :a 1"
+            " :longer 2}"]
+           ["{:a      1"
+            " :longer 2}"]
+           {:align-maps? true})))
+    (testing "indentation off #2"
+      (is (reformats-to?
+           ["{  :a     1"
+            " :longer 2}"]
+           ["{:a      1"
+            " :longer 2}"]
+           {:align-maps? true})))
+    (testing "indentation off #3"
+      (is (reformats-to?
+           ["{:a           1"
+            " :longer 2}"]
+           ["{:a      1"
+            " :longer 2}"]
+           {:align-maps? true})))
+    (testing "columns"
+      (testing "multi-value line"
+        (is (reformats-to?
+             ["{:a 1 :b 2"
+              " :longer 3}"]
+             ["{:a      1 :b 2"
+              " :longer 3}"]
+             {:align-maps? true})))
+      (testing "multi-value line"
+        (is (reformats-to?
+             ["{:a 1 :longer-a 2"
+              " :longer-b 3 :c 4}"]
+             ["{:a        1 :longer-a 2"
+              " :longer-b 3 :c        4}"]
+             {:align-maps? true})))
+      (testing "multi-value commas"
+        (is (reformats-to?
+             ["{:a 1, :longer-a 2"
+              " :longer-b 3 , :c 4}"]
+             ["{:a        1, :longer-a 2"
+              " :longer-b 3, :c        4}"]
+             {:align-maps? true})))
+      (testing "multi-value uneven"
+        (is (reformats-to?
+             ["{:a 1 :longer-a 2 :c 3"
+              " :longer-b 4  :d 5}"]
+             ["{:a        1 :longer-a 2 :c 3"
+              " :longer-b 4 :d        5}"]
+             {:align-maps? true})))
+      (testing "multi-value groups 1"
+        (is (reformats-to?
+             ["{:a 1 :longer-a 2"
+              " :longer-b 3 :c 4"
+              ""
+              " :d 5 :e 6"
+              " :fg 7 :h 8}"]
+             ["{:a        1 :longer-a 2"
+              " :longer-b 3 :c        4"
+              ""
+              " :d  5 :e 6"
+              " :fg 7 :h 8}"]
+             {:align-maps? true})))
+      (testing "multi-value groups 2"
+        (is (reformats-to?
+             ["{:a 1 :longer-a 2"
+              " :longer-b 3 :c 4"
+              ""
+              ""
+              " :d 5 :e 6"
+              " :fg 7 :h 8"
+              ""
+              " :i 9 :jklmno 10"
+              " :p 11 :q :value}"]
+             ["{:a        1 :longer-a 2"
+              " :longer-b 3 :c        4"
+              ""
+              " :d  5 :e 6"
+              " :fg 7 :h 8"
+              ""
+              " :i 9  :jklmno 10"
+              " :p 11 :q      :value}"]
+             {:align-maps? true})))
+      (testing "multi-value partial commas"
+        (is (reformats-to?
+             ["{:a 1 :longer-a 2"
+              " :longer-b 3 , :c 4}"]
+             ["{:a        1  :longer-a 2"
+              " :longer-b 3, :c        4}"]
+             {:align-maps? true}))))))