From 9820f5856b4b042a06c3f2680162ffc7a4819eb5 Mon Sep 17 00:00:00 2001
From: Matt Enlow <enlow@adobe.com>
Date: Mon, 11 Sep 2023 09:16:15 -0400
Subject: [PATCH] rewrite left arrows to put pattern matches on right

---
 CHANGELOG.md                    |  5 +++++
 lib/style/single_node.ex        | 16 +++++++++++-----
 lib/styler.ex                   | 12 +++++-------
 test/style/single_node_test.exs |  9 +++++++--
 4 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2f4fc912..f58bdb9e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,11 @@
 
 ## main
 
+### Improvements
+
+* Added right-hand-pattern-matching rewrites to `for` and `with` left arrow expressions `<-`
+  (ex: `with map = %{} <- foo()` => `with %{} = map <- foo`)
+
 ## v0.8.5
 
 ### Fixes
diff --git a/lib/style/single_node.ex b/lib/style/single_node.ex
index d6dee2a8..573fb3bd 100644
--- a/lib/style/single_node.ex
+++ b/lib/style/single_node.ex
@@ -141,16 +141,22 @@ defmodule Styler.Style.SingleNode do
   defp style(trivial_case(head, {:__block__, _, [true]}, do_body, {:_, _, _}, else_body)),
     do: if_ast(head, do_body, else_body)
 
-  defp style({:case, cm, [head, [{do_block, arrows}]]}), do: {:case, cm, [head, [{do_block, r_align_matches(arrows)}]]}
-
-  defp style({:fn, m, arrows}), do: {:fn, m, r_align_matches(arrows)}
+  # ARROW REWRITES
+  # `with`, `for` left arrow - if only we could write something this trivial for `->`!
+  defp style({:<-, cm, [lhs, rhs]}), do: {:<-, cm, [put_matches_on_right(lhs), rhs]}
+  # there's complexity to `:->` due to `cond` also utilizing the symbol but with different semantics.
+  # thus, we have to have a clause for each place that `:->` can show up
+  # `with` elses
+  defp style({{:__block__, _, [:else]} = else_, arrows}), do: {else_, rewrite_arrows(arrows)}
+  defp style({:case, cm, [head, [{do_, arrows}]]}), do: {:case, cm, [head, [{do_, rewrite_arrows(arrows)}]]}
+  defp style({:fn, m, arrows}), do: {:fn, m, rewrite_arrows(arrows)}
 
   defp style(node), do: node
 
-  defp r_align_matches(arrows) when is_list(arrows),
+  defp rewrite_arrows(arrows) when is_list(arrows),
     do: Enum.map(arrows, fn {:->, m, [lhs, rhs]} -> {:->, m, [put_matches_on_right(lhs), rhs]} end)
 
-  defp r_align_matches(macros_or_something_crazy_oh_no_abooort), do: macros_or_something_crazy_oh_no_abooort
+  defp rewrite_arrows(macros_or_something_crazy_oh_no_abooort), do: macros_or_something_crazy_oh_no_abooort
 
   defp put_matches_on_right(ast) do
     ast
diff --git a/lib/styler.ex b/lib/styler.ex
index 098a7493..78e69810 100644
--- a/lib/styler.ex
+++ b/lib/styler.ex
@@ -60,15 +60,14 @@ defmodule Styler do
 
     {ast, comments} =
       input
-      |> Styler.string_to_quoted_with_comments(to_string(file))
+      |> string_to_quoted_with_comments(to_string(file))
       |> style(file, opts)
 
     quoted_to_string(ast, comments, formatter_opts)
   end
 
-  @doc """
-  Wraps `Code.string_to_quoted_with_comments` with our desired options
-  """
+  @doc false
+  # Wrap `Code.string_to_quoted_with_comments` with our desired options
   def string_to_quoted_with_comments(code, file \\ "nofile") when is_binary(code) do
     Code.string_to_quoted_with_comments!(code,
       literal_encoder: &__MODULE__.literal_encoder/2,
@@ -81,9 +80,8 @@ defmodule Styler do
   @doc false
   def literal_encoder(a, b), do: {:ok, {:__block__, b, [a]}}
 
-  @doc """
-  Turns an ast and comments back into code, formatting it along the way.
-  """
+  @doc false
+  # Turns an ast and comments back into code, formatting it along the way.
   def quoted_to_string(ast, comments, formatter_opts \\ []) do
     opts = [{:comments, comments}, {:escape, false} | formatter_opts]
     {line_length, opts} = Keyword.pop(opts, :line_length, 122)
diff --git a/test/style/single_node_test.exs b/test/style/single_node_test.exs
index 581cc729..0cd69207 100644
--- a/test/style/single_node_test.exs
+++ b/test/style/single_node_test.exs
@@ -21,12 +21,12 @@ defmodule Styler.Style.SingleNodeTest do
     assert_style("Logger.warn(foo, bar)", "Logger.warning(foo, bar)")
   end
 
-  test "Timex.now -> DateTime.utc_now/now!" do
+  test "Timex.now => DateTime.utc_now/now!" do
     assert_style("Timex.now()", "DateTime.utc_now()")
     assert_style(~S|Timex.now("Some/Timezone")|, ~S|DateTime.now!("Some/Timezone")|)
   end
 
-  test "Timex.today -> Date.utc_today" do
+  test "Timex.today => Date.utc_today" do
     assert_style("Timex.today()", "Date.utc_today()")
     assert_style(~S|Timex.today("Some/Timezone")|, ~S|Timex.today("Some/Timezone")|)
   end
@@ -133,6 +133,11 @@ defmodule Styler.Style.SingleNodeTest do
   end
 
   describe "RHS pattern matching" do
+    test "left arrows" do
+      assert_style("with {:ok, result = %{}} <- foo, do: result", "with {:ok, %{} = result} <- foo, do: result")
+      assert_style("for map = %{} <- maps, do: map[:key]", "for %{} = map <- maps, do: map[:key]")
+    end
+
     test "case statements" do
       assert_style(
         """