From fab7a130e79c38d7d2deb1ce75dc4d80676e6af7 Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Sun, 9 Jan 2022 10:08:16 -0800
Subject: [PATCH] Router: preserve break after `while` if `keep`

---
 .../scala/org/scalafmt/internal/FormatOps.scala   | 11 ++++++-----
 .../main/scala/org/scalafmt/internal/Router.scala |  8 ++++++--
 .../resources/scala3/OptionalBraces_keep.stat     | 15 ++++++++++-----
 3 files changed, 22 insertions(+), 12 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 da1c5e5592..e2c9229e9b 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
@@ -2484,7 +2484,7 @@ class FormatOps(
     @inline private def seqLast(seq: Seq[Tree]): Option[T] =
       if (isSeqMulti(seq)) treeLast(seq.last) else None
 
-    def indentAndBreakBeforeCtrl[A](tree: Tree, split: Split)(implicit
+    def indentAndBreakBeforeCtrl[A](tree: Tree, fsplit: => Split)(implicit
         style: ScalafmtConfig,
         classifier: Classifier[T, A]
     ): Option[Split] =
@@ -2495,12 +2495,13 @@ class FormatOps(
         if (kw.is[A]) Some {
           val indent =
             style.indent.ctrlSite.getOrElse(style.indent.getSignificant)
+          val split = fsplit
+          def policy =
+            if (split.isNL) decideNewlinesOnlyBeforeClose(kw)
+            else decideNewlinesOnlyBeforeCloseOnBreak(kw)
           split
             .withIndent(Num(indent), kw, ExpiresOn.Before)
-            .andPolicy(
-              decideNewlinesOnlyBeforeCloseOnBreak(kw),
-              !style.danglingParentheses.ctrlSite
-            )
+            .andPolicy(policy, !style.danglingParentheses.ctrlSite)
         }
         else None
       }
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 d4935c797c..82a12360be 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
@@ -1868,12 +1868,16 @@ class Router(formatOps: FormatOps) {
           .indentAndBreakBeforeCtrl[T.KwThen](owner.cond, mlSplitBase)
         Seq(slb, mlSplitOpt.getOrElse(mlSplitBase))
       case FormatToken(_: T.KwWhile, right, _) =>
-        val splitBase = Split(Space(style.spaces.isSpaceAfterKeyword(right)), 0)
+        def spaceMod = Space(style.spaces.isSpaceAfterKeyword(right))
+        def splitBase = {
+          val onlyNL = style.newlines.source.eq(Newlines.keep) && newlines != 0
+          Split(if (onlyNL) Newline else spaceMod, 0)
+        }
         val split = (formatToken.meta.leftOwner match {
           case t: Term.While =>
             OptionalBraces.indentAndBreakBeforeCtrl[T.KwDo](t.expr, splitBase)
           case _ => None
-        }).getOrElse(splitBase)
+        }).getOrElse(Split(spaceMod, 0))
         Seq(split)
       case FormatToken(_: T.KwFor, right, _) =>
         Seq(Split(Space(style.spaces.isSpaceAfterKeyword(right)), 0))
diff --git a/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat b/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
index 89c17cb58a..dfb7f43d4d 100644
--- a/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
+++ b/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
@@ -1977,7 +1977,8 @@ while
     exitCode
     
 >>>
-while (a1 && a2) == cond1 do
+while
+   (a1 && a2) == cond1 do
    val exitCode = 1
    exitCode
 <<< if redundant parens
@@ -3231,12 +3232,14 @@ object a:
    do matrix(x)(0) += matrix(x - 1)(0)
    while x < range do
       matrix(x)(0) += matrix(x - 1)(0)
-   while x < range
+   while
+      x < range
    do matrix(x)(0) += matrix(x - 1)(0)
    while x < range
    do
       matrix(x)(0) += matrix(x - 1)(0)
-   while x < range
+   while
+      x < range
    do
       matrix(x)(0) += matrix(x - 1)(0)
 <<< #3034 while-do !danglingParentheses
@@ -3263,11 +3266,13 @@ object a:
    do matrix(x)(0) += matrix(x - 1)(0)
    while x < range do
       matrix(x)(0) += matrix(x - 1)(0)
-   while x < range
+   while
+      x < range
    do matrix(x)(0) += matrix(x - 1)(0)
    while x < range
    do
       matrix(x)(0) += matrix(x - 1)(0)
-   while x < range
+   while
+      x < range
    do
       matrix(x)(0) += matrix(x - 1)(0)