From fdf74f289e9978bc0e4bad7de646edb020325cae Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Thu, 27 Jan 2022 06:04:51 -0800
Subject: [PATCH] Router: keep newline after `if` if source=keep

---
 .../scala/org/scalafmt/internal/FormatOps.scala |  9 +++++----
 .../scala/org/scalafmt/internal/Router.scala    | 12 ++++++++----
 .../resources/scala3/OptionalBraces_keep.stat   | 17 ++++++++++++-----
 3 files changed, 25 insertions(+), 13 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 04ac145305..94dc16720e 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
@@ -2476,18 +2476,19 @@ class FormatOps(
     @inline private def seqLast(seq: Seq[Tree]): Option[T] =
       if (isSeqMulti(seq)) treeLast(seq.last) else None
 
-    def indentAndBreakBeforeCtrl[A](tree: Tree, fsplit: => Split)(implicit
+    def indentAndBreakBeforeCtrl[A](tree: Tree, split: Split)(implicit
         style: ScalafmtConfig,
         classifier: Classifier[T, A]
     ): Option[Split] =
-      if (!style.dialect.allowSignificantIndentation) None
-      else if (tree.is[Term.Block] && isEnclosedInMatching(tree)) None
+      if (
+        !style.dialect.allowSignificantIndentation ||
+        tree.is[Term.Block] && !split.isNL && isEnclosedInMatching(tree)
+      ) None
       else {
         val kw = tokenAfter(tree).right
         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)
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 aceef355f0..79948a162c 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
@@ -1952,15 +1952,19 @@ class Router(formatOps: FormatOps) {
                   if elses.contains(r) =>
                 d.onlyNewlinesWithFallback(Split(Newline, 0))
             }
-        val spaceMod = Space(style.spaces.isSpaceAfterKeyword(right))
-        val slb = Split(spaceMod, 0).withSingleLine(expire, killOnFail = true)
-        val mlSplitBase = Split(spaceMod, 1).withPolicy(breakOnlyBeforeElse)
+        val mod =
+          if (style.newlines.keepBreak(newlines)) Newline
+          else Space(style.spaces.isSpaceAfterKeyword(right))
+        val slb =
+          Split(mod.isNewline, 0)(mod).withSingleLine(expire, killOnFail = true)
+        val mlSplitBase = Split(mod, if (slb.isIgnored) 0 else 1)
+          .withPolicy(breakOnlyBeforeElse)
         val mlSplitOpt = OptionalBraces
           .indentAndBreakBeforeCtrl[T.KwThen](owner.cond, mlSplitBase)
         Seq(slb, mlSplitOpt.getOrElse(mlSplitBase))
       case FormatToken(_: T.KwWhile | _: T.KwFor, right, _) =>
         def spaceMod = Space(style.spaces.isSpaceAfterKeyword(right))
-        def splitBase = {
+        val splitBase = {
           val onlyNL = style.newlines.keepBreak(newlines)
           Split(if (onlyNL) Newline else spaceMod, 0)
         }
diff --git a/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat b/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
index 128c9dfd18..87e6aa7653 100644
--- a/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
+++ b/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
@@ -1972,7 +1972,8 @@ if
     exitCode
     
 >>>
-if (a1 && a2) == cond1 then
+if
+   (a1 && a2) == cond1 then
    val exitCode = 1
    exitCode
 <<< while indentation danglingParentheses false
@@ -3055,7 +3056,9 @@ object a {
 }
 >>>
 object a {
-  if ({ !Yield(2); false }) then {
+  if
+     ({ !Yield(2); false })
+  then {
     !Yield(3)
   } else {
     !Yield(4)
@@ -3082,7 +3085,9 @@ object a {
 }
 >>>
 object a {
-  while ({ !Yield(2); false }) do {
+  while
+     ({ !Yield(2); false })
+  do {
     !Yield(3)
   }
 }
@@ -3111,7 +3116,8 @@ class C:
    then
       3
    else 4
-   if true
+   if
+      true
    then
       3
    else 4
@@ -3140,7 +3146,8 @@ class C:
    then
       3
    else 4
-   if true
+   if
+      true
    then
       3
    else 4