with statements: happier comments and handle sad edge cases (#102)

CHANGELOG.md

@@ -4,11 +4,19 @@
 
 ### Improvements
 
-* `with`: rewrite trivial `lhs <- rhs` to `lhs = rhs` (#86)
-* `with`: rewrite with statements if statements when appropriate
-* `with`: switch keyword do to do block when adding clauses to the with body (`, do:` => `do end`)
+#### With Statements
+
+A slew of improvements for the `with` statement in this release:
+
+* rewrite trivial `lhs <- rhs` to `lhs = rhs` (#86)
+* rewrite with statements if statements when appropriate
+* switch keyword do to do block when adding clauses to the with body (`, do:` => `do end`)
+* put more effort into keeping comments in the right spot
+
+#### Other
+
 * Rewrite `{Map|Keyword}.merge(single_key: value)` to use `put/3` instead (#96)
-* Attempt to keep comments in logical places when rewriting trivial `case` and `cond` statments (#97)
+* Attempt to keep comments in logical places when rewriting trivial `case` and `cond` statements (#97)
 
 ## v0.10.5
 

lib/style/blocks.ex

@@ -74,59 +74,80 @@ defmodule Styler.Style.Blocks do
         end)
         |> Enum.split_while(&(not left_arrow?(&1)))
 
-      # the do/else keyword macro of the with statement is the last element of the list
-      [[{do_block, do_body} | elses] | reversed_clauses] = Enum.reverse(children)
-      {postroll, reversed_clauses} = Enum.split_while(reversed_clauses, &(not left_arrow?(&1)))
-      [{:<-, _, [lhs, rhs]} = _final_clause | rest] = reversed_clauses
-
-      # drop singleton identity else clauses like `else foo -> foo end`
-      elses =
-        case elses do
-          [{{_, _, [:else]}, [{:->, _, [[left], right]}]}] -> if nodes_equivalent?(left, right), do: [], else: elses
-          _ -> elses
-        end
-
-      {reversed_clauses, do_body} =
+      # after rewriting `x <- y()` to `x = y()` there are no more arrows.
+      # this never should've been a with statement at all! we can just replace it with assignments
+      if Enum.empty?(children) do
+        {:cont, replace_with_statement(zipper, preroll), ctx}
+      else
+        [[{{_, do_meta, _} = do_block, do_body} | elses] | reversed_clauses] = Enum.reverse(children)
+        {postroll, reversed_clauses} = Enum.split_while(reversed_clauses, &(not left_arrow?(&1)))
+        [{:<-, final_clause_meta, [lhs, rhs]} = _final_clause | rest] = reversed_clauses
+
+        # drop singleton identity else clauses like `else foo -> foo end`
+        elses =
+          case elses do
+            [{{_, _, [:else]}, [{:->, _, [[left], right]}]}] -> if nodes_equivalent?(left, right), do: [], else: elses
+            _ -> elses
+          end
+
+        {reversed_clauses, do_body} =
+          cond do
+            # Put the postroll into the body
+            Enum.any?(postroll) ->
+              {node, do_body_meta, do_children} = do_body
+              do_children = if node == :__block__, do: do_children, else: [do_body]
+              do_body = {:__block__, Keyword.take(do_body_meta, [:line]), Enum.reverse(postroll, do_children)}
+              {reversed_clauses, do_body}
+
+            # Credo.Check.Refactor.RedundantWithClauseResult
+            Enum.empty?(elses) and nodes_equivalent?(lhs, do_body) ->
+              {rest, rhs}
+
+            # no change
+            true ->
+              {reversed_clauses, do_body}
+          end
+
+        do_line = do_meta[:line]
+        final_clause_line = final_clause_meta[:line]
+
+        do_line =
+          cond do
+            do_meta[:format] == :keyword && final_clause_line + 1 >= do_line -> do_line
+            do_meta[:format] == :keyword -> final_clause_line + 1
+            true -> final_clause_line
+          end
+
+        do_block = Macro.update_meta(do_block, &Keyword.put(&1, :line, do_line))
+        # disable keyword `, do:` since there will be multiple statements in the body
+        with_meta =
+          if Enum.any?(postroll),
+            do: Keyword.merge(with_meta, do: [line: with_meta[:line]], end: [line: max_line(children) + 1]),
+            else: with_meta
+
+        with_children = Enum.reverse(reversed_clauses, [[{do_block, do_body} | elses]])
+        zipper = Zipper.replace(zipper, {:with, with_meta, with_children})
+
+        # recurse if the # of `<-` have changed (this `with` could now be eligible for a `case` rewrite)
         cond do
-          # Put the postroll into the body
+          Enum.any?(preroll) ->
+            # put the preroll before the with statement in either a block we create or the existing parent block
+            preroll
+            |> Enum.reduce(Style.ensure_block_parent(zipper), &Zipper.insert_left(&2, &1))
+            |> run(ctx)
+
           Enum.any?(postroll) ->
-            {node, do_body_meta, do_children} = do_body
-            do_children = if node == :__block__, do: do_children, else: [do_body]
-            do_body = {:__block__, Keyword.take(do_body_meta, [:line]), Enum.reverse(postroll, do_children)}
-            {reversed_clauses, do_body}
+            # both of these changed the # of `<-`, so we should have another look at this with statement
+            run(zipper, ctx)
 
-          # Credo.Check.Refactor.RedundantWithClauseResult
-          Enum.empty?(elses) and nodes_equivalent?(lhs, do_body) ->
-            {rest, rhs}
+          Enum.empty?(reversed_clauses) ->
+            # oops! RedundantWithClauseResult removed the final arrow in this. no more need for a with statement!
+            {:cont, replace_with_statement(zipper, with_children), ctx}
 
-          # no change
           true ->
-            {reversed_clauses, do_body}
+            # of clauess didn't change, so don't reecurse or we'll loop FOREEEVEERR
+            {:cont, zipper, ctx}
         end
-
-      # disable keyword `, do:` since there will be multiple clauses
-      with_meta =
-        if Enum.any?(postroll),
-          do: Keyword.merge(with_meta, do: [line: with_meta[:line]], end: [line: max_line(children) + 1]),
-          else: with_meta
-
-      zipper = Zipper.replace(zipper, {:with, with_meta, Enum.reverse(reversed_clauses, [[{do_block, do_body} | elses]])})
-
-      # if there was pre or postroll, the # of `<-` in the statement have changed and so it could be eligible for a `case`
-      # or even `if` rewrite -- so we recurse in both of those cases
-      cond do
-        Enum.any?(preroll) ->
-          # put the preroll before the with statement in either a block we create or the existing parent block
-          preroll
-          |> Enum.reduce(Style.ensure_block_parent(zipper), &Zipper.insert_left(&2, &1))
-          |> run(ctx)
-
-        Enum.any?(postroll) ->
-          run(zipper, ctx)
-
-        # if the # of clauses didn't change, then we don't need to recurse and can continue from here =)
-        true ->
-          {:cont, zipper, ctx}
       end
     else
       # maybe this isn't a with statement - could be a function named `with`
@@ -153,6 +174,30 @@ defmodule Styler.Style.Blocks do
 
   defp style(node), do: node
 
+  # `with a <- b(), c <- d(), do: :ok, else: (_ -> :error)`
+  # =>
+  # `a = b(); c = d(); :ok`
+  defp replace_with_statement(zipper, preroll) do
+    [[{_do, do_body} | _elses] | preroll] = Enum.reverse(preroll)
+    block = Enum.reverse(preroll, [do_body])
+
+    case Zipper.up(zipper) do
+      nil ->
+        Zipper.zip({:__block__, [], block})
+
+      {{:=, _, _}, _} ->
+        Zipper.update(zipper, fn {:with, meta, _} -> {:__block__, Keyword.take(meta, [:line]), block} end)
+
+      {{:__block__, _, _}, _} ->
+        block
+        |> Enum.reduce(zipper, &Zipper.insert_left(&2, &1))
+        |> Zipper.remove()
+
+      ast ->
+        raise "unexpected `with` parent ast: #{inspect(ast)}"
+    end
+  end
+
   defp left_arrow?({:<-, _, _}), do: true
   defp left_arrow?(_), do: false
 

test/style/blocks_test.exs

@@ -240,6 +240,75 @@ defmodule Styler.Style.BlocksTest do
   end
 
   describe "with statements" do
+    test "the saddest edge case of all" do
+      assert_style(
+        """
+        x()
+
+        z =
+          with a <- b(), c <- d(), e <- f() do
+            g
+          else
+            _ -> h
+          end
+
+        y()
+        """,
+        """
+        x()
+
+        z =
+          (
+            a = b()
+            c = d()
+            e = f()
+            g
+          )
+
+        y()
+        """
+      )
+
+      assert_style(
+        """
+        with a <- b(), c <- d(), e <- f() do
+          g
+        else
+          _ -> h
+        end
+        """,
+        """
+        a = b()
+        c = d()
+        e = f()
+        g
+        """
+      )
+
+      assert_style(
+        """
+        x()
+
+        with a <- b(), c <- d(), e <- f() do
+          g
+        else
+          _ -> h
+        end
+
+        y()
+        """,
+        """
+        x()
+
+        a = b()
+        c = d()
+        e = f()
+        g
+        y()
+        """
+      )
+    end
+
     test "doesn't false positive with vars" do
       assert_style("""
       if naming_is_hard, do: with
@@ -293,6 +362,14 @@ defmodule Styler.Style.BlocksTest do
         end
         """)
       end
+
+      assert_style(
+        """
+        with :ok <- foo(),
+          do: :ok
+        """,
+        "foo()"
+      )
     end
 
     test "rewrites non-pattern-matching lhs" do
@@ -319,6 +396,13 @@ defmodule Styler.Style.BlocksTest do
       )
     end
 
+    test "doesn't move keyword do up when it's just one line" do
+      assert_style("""
+      with :ok <- foo(),
+           do: :error
+      """)
+    end
+
     test "rewrites `_ <- rhs` to just rhs" do
       assert_style(
         """
@@ -557,6 +641,33 @@ defmodule Styler.Style.BlocksTest do
       end
       """)
     end
+
+    test "with comments" do
+      # i think the bug here is that the `do` keyword's ast needs its line number moved up
+      # to be equal to the last arrow's line number
+      assert_style(
+        """
+        with :ok <- foo(),
+            :ok <- bar(),
+            # comment 1
+            # comment 2
+            # comment 3
+            _ <- bop() do
+          :ok
+        end
+        """,
+        """
+        with :ok <- foo(),
+             :ok <- bar() do
+          # comment 1
+          # comment 2
+          # comment 3
+          bop()
+          :ok
+        end
+        """
+      )
+    end
   end
 
   test "Credo.Check.Refactor.CondStatements" do