handle block bodies in with statement replacement

lib/style/blocks.ex

@@ -227,23 +227,25 @@ defmodule Styler.Style.Blocks do
   defp replace_with_statement(zipper, preroll) do
     [[{_do, do_body} | _elses] | preroll] = Enum.reverse(preroll)
 
-    # RedundantWithClauseResult except we rewrote the `<-` to an `=`
-    # with a, b, x <- y(), do: x
-    # =>
-    # a; b; y
     block =
-      case preroll do
-        [{:=, _, [lhs, rhs]} | rest] ->
-          if nodes_equivalent?(lhs, do_body),
-            do: [rhs | rest],
-            else: [do_body | preroll]
+      case do_body do
+        {:__block__, _, [{_, _, _} | _] = children} ->
+          Enum.reverse(preroll, children)
 
         _ ->
-          [do_body | preroll]
+          # RedundantWithClauseResult except we rewrote the `<-` to an `=`
+          # `with a, b, x <- y(), do: x` => `a; b; y`
+          case preroll do
+            [{:=, _, [lhs, rhs]} | rest] ->
+              if nodes_equivalent?(lhs, do_body),
+                do: Enum.reverse(rest, [rhs]),
+                else: Enum.reverse(preroll, [do_body])
+
+            _ ->
+              Enum.reverse(preroll, [do_body])
+          end
       end
 
-    block = Enum.reverse(block)
-
     case Style.ensure_block_parent(zipper) do
       {:ok, zipper} ->
         zipper

test/style/blocks_test.exs

@@ -371,6 +371,7 @@ defmodule Styler.Style.BlocksTest do
       assert_style(
         """
         with x = y, a = b do
+          w
           z
         else
          _ -> whatever
@@ -379,6 +380,7 @@ defmodule Styler.Style.BlocksTest do
         """
         x = y
         a = b
+        w
         z
         """
       )
@@ -554,9 +556,7 @@ defmodule Styler.Style.BlocksTest do
         end
         """
       )
-    end
 
-    test "with to if" do
       assert_style(
         """
         with true <- foo do