From 95fba171b1ddc945c4ae624f3b3d7e7b425f41a1 Mon Sep 17 00:00:00 2001
From: Matt Enlow <enlow@adobe.com>
Date: Mon, 15 Jul 2024 16:35:29 -0600
Subject: [PATCH] handle block bodies in with statement replacement

---
 lib/style/blocks.ex        | 26 ++++++++++++++------------
 test/style/blocks_test.exs |  4 ++--
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/lib/style/blocks.ex b/lib/style/blocks.ex
index a6f5cff4..dbfe5ab8 100644
--- a/lib/style/blocks.ex
+++ b/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
diff --git a/test/style/blocks_test.exs b/test/style/blocks_test.exs
index 79d8ddca..8444666f 100644
--- a/test/style/blocks_test.exs
+++ b/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