From b58db35f3ffe1f2f0a54e2301a6fcb179c918916 Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Fri, 7 Jan 2022 22:20:33 -0800
Subject: [PATCH] Router: handle then and do similar to else/yield

---
 .../scala/org/scalafmt/internal/Router.scala  |  9 +++
 .../test/resources/scala3/OptionalBraces.stat | 60 +++++++++++--------
 .../resources/scala3/OptionalBraces_fold.stat | 24 ++++----
 .../resources/scala3/OptionalBraces_keep.stat | 48 ++++++++++-----
 4 files changed, 89 insertions(+), 52 deletions(-)

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 60d0752cc3..c793c88d65 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
@@ -1932,6 +1932,15 @@ class Router(formatOps: FormatOps) {
             .withSingleLineNoOptimal(expire, exclude, noSyntaxNL = noSyntaxNL),
           Split(Newline, 1)
         )
+      case ft @ FormatToken(_, _: T.KwThen | _: T.KwDo, _) =>
+        if (style.newlines.sourceIgnored || newlines == 0)
+          Seq(
+            Split(Space, 0)
+              .withOptimalToken(nextNonCommentSameLine(next(ft)).left),
+            Split(Newline, 1)
+          )
+        else
+          Seq(Split(Newline, 0))
       // Last else branch
       case FormatToken(_: T.KwElse, _, _) if (leftOwner match {
             case t: Term.If => !t.elsep.is[Term.If]
diff --git a/scalafmt-tests/src/test/resources/scala3/OptionalBraces.stat b/scalafmt-tests/src/test/resources/scala3/OptionalBraces.stat
index d8b96024f0..f6fb36c735 100644
--- a/scalafmt-tests/src/test/resources/scala3/OptionalBraces.stat
+++ b/scalafmt-tests/src/test/resources/scala3/OptionalBraces.stat
@@ -3115,14 +3115,14 @@ class C:
   else 4
 >>>
 class C:
-   if true then
-      3
+   if true
+   then 3
    else 4
-   if true then
-      3
+   if true
+   then 3
    else 4
-   if true then
-      3
+   if true
+   then 3
    else 4
 <<< #3033 if-then !danglingParentheses
 maxColumn = 15
@@ -3143,14 +3143,14 @@ class C:
   else 4
 >>>
 class C:
-   if true then
-      3
+   if true
+   then 3
    else 4
-   if true then
-      3
+   if true
+   then 3
    else 4
-   if true then
-      3
+   if true
+   then 3
    else 4
 <<< #3034 for-yield danglingParentheses
 danglingParentheses.ctrlSite = true
@@ -3197,9 +3197,12 @@ object a:
    matrix(x)(0) += matrix(x - 1)(0)
 >>>
 object a:
-   for x <- range do matrix(x)(0) += matrix(x - 1)(0)
-   for x <- range do matrix(x)(0) += matrix(x - 1)(0)
-   for x <- range do matrix(x)(0) += matrix(x - 1)(0)
+   for x <- range
+   do matrix(x)(0) += matrix(x - 1)(0)
+   for x <- range
+   do matrix(x)(0) += matrix(x - 1)(0)
+   for x <- range
+   do matrix(x)(0) += matrix(x - 1)(0)
 <<< #3034 for-do !danglingParentheses
 danglingParentheses.ctrlSite = false
 ===
@@ -3215,9 +3218,12 @@ object a:
    matrix(x)(0) += matrix(x - 1)(0)
 >>>
 object a:
-   for x <- range do matrix(x)(0) += matrix(x - 1)(0)
-   for x <- range do matrix(x)(0) += matrix(x - 1)(0)
-   for x <- range do matrix(x)(0) += matrix(x - 1)(0)
+   for x <- range
+   do matrix(x)(0) += matrix(x - 1)(0)
+   for x <- range
+   do matrix(x)(0) += matrix(x - 1)(0)
+   for x <- range
+   do matrix(x)(0) += matrix(x - 1)(0)
 <<< #3034 while-do danglingParentheses
 danglingParentheses.ctrlSite = true
 ===
@@ -3233,9 +3239,12 @@ object a:
     matrix(x)(0) += matrix(x - 1)(0)
 >>>
 object a:
-   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 do matrix(x)(0) += matrix(x - 1)(0)
+   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
+   do matrix(x)(0) += matrix(x - 1)(0)
 <<< #3034 while-do !danglingParentheses
 danglingParentheses.ctrlSite = false
 ===
@@ -3251,6 +3260,9 @@ object a:
     matrix(x)(0) += matrix(x - 1)(0)
 >>>
 object a:
-   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 do matrix(x)(0) += matrix(x - 1)(0)
+   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
+   do matrix(x)(0) += matrix(x - 1)(0)
diff --git a/scalafmt-tests/src/test/resources/scala3/OptionalBraces_fold.stat b/scalafmt-tests/src/test/resources/scala3/OptionalBraces_fold.stat
index 8b64f2ee94..45ef256dd1 100644
--- a/scalafmt-tests/src/test/resources/scala3/OptionalBraces_fold.stat
+++ b/scalafmt-tests/src/test/resources/scala3/OptionalBraces_fold.stat
@@ -2948,14 +2948,14 @@ class C:
   else 4
 >>>
 class C:
-   if true then
-      3
+   if true
+   then 3
    else 4
-   if true then
-      3
+   if true
+   then 3
    else 4
-   if true then
-      3
+   if true
+   then 3
    else 4
 <<< #3033 if-then short !danglingParentheses
 maxColumn = 13
@@ -2976,14 +2976,14 @@ class C:
   else 4
 >>>
 class C:
-   if true then
-      3
+   if true
+   then 3
    else 4
-   if true then
-      3
+   if true
+   then 3
    else 4
-   if true then
-      3
+   if true
+   then 3
    else 4
 <<< #3034 for-yield danglingParentheses
 maxColumn = 50
diff --git a/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat b/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
index 9d94804f0a..a312249b8e 100644
--- a/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
+++ b/scalafmt-tests/src/test/resources/scala3/OptionalBraces_keep.stat
@@ -3081,12 +3081,15 @@ class C:
   else 4
 >>>
 class C:
-   if true then 3
+   if true
+   then 3
    else 4
-   if true then
+   if true
+   then
       3
    else 4
-   if true then
+   if true
+   then
       3
    else 4
 <<< #3033 if-then !danglingParentheses
@@ -3107,12 +3110,15 @@ class C:
   else 4
 >>>
 class C:
-   if true then 3
+   if true
+   then 3
    else 4
-   if true then
+   if true
+   then
       3
    else 4
-   if true then
+   if true
+   then
       3
    else 4
 <<< #3034 for-yield danglingParentheses
@@ -3160,8 +3166,10 @@ object a:
    matrix(x)(0) += matrix(x - 1)(0)
 >>>
 object a:
-   for x <- range do matrix(x)(0) += matrix(x - 1)(0)
-   for x <- range do
+   for x <- range
+   do matrix(x)(0) += matrix(x - 1)(0)
+   for x <- range
+   do
       matrix(x)(0) += matrix(x - 1)(0)
    for
       x <- range
@@ -3182,8 +3190,10 @@ object a:
    matrix(x)(0) += matrix(x - 1)(0)
 >>>
 object a:
-   for x <- range do matrix(x)(0) += matrix(x - 1)(0)
-   for x <- range do
+   for x <- range
+   do matrix(x)(0) += matrix(x - 1)(0)
+   for x <- range
+   do
       matrix(x)(0) += matrix(x - 1)(0)
    for
       x <- range
@@ -3204,10 +3214,13 @@ object a:
     matrix(x)(0) += matrix(x - 1)(0)
 >>>
 object a:
-   while x < range do matrix(x)(0) += matrix(x - 1)(0)
-   while x < range do
+   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 do
+   while x < range
+   do
       matrix(x)(0) += matrix(x - 1)(0)
 <<< #3034 while-do !danglingParentheses
 danglingParentheses.ctrlSite = false
@@ -3224,8 +3237,11 @@ object a:
     matrix(x)(0) += matrix(x - 1)(0)
 >>>
 object a:
-   while x < range do matrix(x)(0) += matrix(x - 1)(0)
-   while x < range do
+   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 do
+   while x < range
+   do
       matrix(x)(0) += matrix(x - 1)(0)