From 856b40627df102b2af0b7964ff26b9d837793d5f Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Thu, 26 Sep 2024 16:40:03 -0700
Subject: [PATCH] Router: extract single stat from rewritten block

---
 .../org/scalafmt/internal/FormatOps.scala     | 17 +++++++-------
 .../scala/org/scalafmt/internal/Router.scala  |  9 ++++----
 .../resources/rewrite/RedundantBraces.stat    | 22 ++++++-------------
 3 files changed, 21 insertions(+), 27 deletions(-)

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
index 951245a336..7af1c80b65 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
@@ -1577,14 +1577,6 @@ class FormatOps(
         getNlSplit(if (policy.isEmpty) 0 else 1),
       )
       def hasStateColumn = spaceIndents.exists(_.hasStateColumn)
-      @tailrec
-      def getBlockStat(t: Tree): Tree = t match {
-        case b: Term.Block => getSingleStatExceptEndMarker(b.stats) match {
-            case Some(s) if !isEnclosedInMatching(b) => getBlockStat(s)
-            case _ => t
-          }
-        case _ => t
-      }
       val adjustedBody = getBlockStat(body)
       val (spaceSplit, nlSplit) = adjustedBody match {
         case t: Term.If if isKeep || ifWithoutElse(t) || hasStateColumn =>
@@ -1756,6 +1748,15 @@ class FormatOps(
     ): Split = asInfixApp(body)
       .fold(withIndent(nlSplit, endFt))(InfixSplits.withNLIndent(nlSplit))
 
+    @tailrec
+    def getBlockStat(t: Tree): Tree = t match {
+      case b: Term.Block => getSingleStatExceptEndMarker(b.stats) match {
+          case Some(s) if !isEnclosedInMatching(b) => getBlockStat(s)
+          case _ => t
+        }
+      case _ => t
+    }
+
   }
 
   def withNLPolicy(endFt: FormatToken)(nlSplit: Split): Split = nlSplit
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala
index d95c926d71..f134f1eacf 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala
@@ -2611,10 +2611,10 @@ class Router(formatOps: FormatOps) {
     if (style.newlines.getBeforeMultiline eq Newlines.classic) classicSplits
     else CtrlBodySplits.getWithIndent(body, endFt)(null)(Split(Newline2x(ft), _))
 
-  private def getSplitsValEqualsClassic(body: Tree, endFt: FormatToken)(implicit
-      style: ScalafmtConfig,
-      ft: FormatToken,
-  ): Seq[Split] = {
+  private def getSplitsValEqualsClassic(
+      rawBody: Tree,
+      endFt: FormatToken,
+  )(implicit style: ScalafmtConfig, ft: FormatToken): Seq[Split] = {
     def wouldDangle = ft.meta.leftOwner.parent
       .exists(style.danglingParentheses.atSite(_, false))
 
@@ -2648,6 +2648,7 @@ class Router(formatOps: FormatOps) {
           PenalizeAllNewlines(expire, Constants.ShouldBeSingleLine),
         )
       }
+    val body = CtrlBodySplits.getBlockStat(rawBody)
     val spaceSplit = body match {
       case _ if ft.hasBreak && ft.meta.leftOwner.is[Defn] => Split.ignored
       case _: Term.If => twoBranches
diff --git a/scalafmt-tests/shared/src/test/resources/rewrite/RedundantBraces.stat b/scalafmt-tests/shared/src/test/resources/rewrite/RedundantBraces.stat
index eea503795b..cc3f3d18f2 100644
--- a/scalafmt-tests/shared/src/test/resources/rewrite/RedundantBraces.stat
+++ b/scalafmt-tests/shared/src/test/resources/rewrite/RedundantBraces.stat
@@ -2054,18 +2054,10 @@ hostLocalDirManager = {
   }
 }
 >>>
-Idempotency violated
-=> Diff (- obtained, + expected)
--hostLocalDirManager =
--  if (
--    (conf.get(config.SHUFFLE_HOST_LOCAL_DISK_READING_ENABLED) &&
--      !conf.get(config.SHUFFLE_USE_OLD_FETCH_PROTOCOL)) ||
--    Utils.isPushBasedShuffleEnabled(conf, isDriver)
--  )
--    None
-+hostLocalDirManager = if (
-+  (conf.get(config.SHUFFLE_HOST_LOCAL_DISK_READING_ENABLED) &&
-+    !conf.get(config.SHUFFLE_USE_OLD_FETCH_PROTOCOL)) ||
-+  Utils.isPushBasedShuffleEnabled(conf, isDriver)
-+)
-+  None
+hostLocalDirManager =
+  if (
+    (conf.get(config.SHUFFLE_HOST_LOCAL_DISK_READING_ENABLED) &&
+      !conf.get(config.SHUFFLE_USE_OLD_FETCH_PROTOCOL)) ||
+    Utils.isPushBasedShuffleEnabled(conf, isDriver)
+  )
+    None