From 95250b3c322936345dd3ab9c4ec325885492a537 Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Wed, 9 Oct 2024 14:26:31 -0700
Subject: [PATCH] Router: don't penalize NL before `:` if keeping

Follow-on to #4389.
---
 .../scala/org/scalafmt/internal/Router.scala  | 23 +++++++++++--------
 .../test/scala/org/scalafmt/FormatTests.scala |  2 +-
 2 files changed, 15 insertions(+), 10 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 3d1a2d693..d414ecbcb 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
@@ -1429,16 +1429,21 @@ class Router(formatOps: FormatOps) {
           def penalizeNewlines(extra: Int)(implicit fileLine: FileLine) =
             PenalizeAllNewlines(expire, Constants.BracketPenalty + extra)
           val indent = style.indent.getDefnSite(leftOwner)
-          val nlPenalty = 2 + treeDepth(returnType)
-          Seq(
-            Split(style.newlines.keepBreak(hasBreak()), 0)(sameLineSplit)
-              .withPolicy(penalizeNewlines(0)),
-            // Spark style guide allows this:
-            // https://github.com/databricks/scala-style-guide#indent
-            Split(Newline, Constants.SparkColonNewline + nlPenalty)
-              .withIndent(indent, expire, After)
-              .withPolicy(penalizeNewlines(nlPenalty)),
+          if (style.newlines.keepBreak(hasBreak())) Seq(
+            Split(Newline, 1).withIndent(indent, expire, After)
+              .withPolicy(penalizeNewlines(1)),
           )
+          else {
+            val nlPenalty = 2 + treeDepth(returnType)
+            Seq(
+              Split(sameLineSplit, 0).withPolicy(penalizeNewlines(0)),
+              // Spark style guide allows this:
+              // https://github.com/databricks/scala-style-guide#indent
+              Split(Newline, Constants.SparkColonNewline + nlPenalty)
+                .withIndent(indent, expire, After)
+                .withPolicy(penalizeNewlines(nlPenalty)),
+            )
+          }
         }
       case FormatToken(_: T.Colon, _, ColonDeclTpeLeft(returnType))
           if style.newlines.avoidInResultType =>
diff --git a/scalafmt-tests/shared/src/test/scala/org/scalafmt/FormatTests.scala b/scalafmt-tests/shared/src/test/scala/org/scalafmt/FormatTests.scala
index 6a9196345..20c8e1ca4 100644
--- a/scalafmt-tests/shared/src/test/scala/org/scalafmt/FormatTests.scala
+++ b/scalafmt-tests/shared/src/test/scala/org/scalafmt/FormatTests.scala
@@ -137,7 +137,7 @@ class FormatTests extends FunSuite with CanRunTests with FormatAssertions {
   override def afterAll(): Unit = {
     logger.debug(s"Total explored: ${Debug.explored}")
     if (!onlyUnit && !onlyManual)
-      assertEquals(Debug.explored, 1777089, "total explored")
+      assertEquals(Debug.explored, 1777077, "total explored")
     val results = debugResults.result()
     // TODO(olafur) don't block printing out test results.
     // I don't want to deal with scalaz's Tasks :'(