Split: add implicit penalizeNL and penalizeIf

scalameta · Nov 25, 2024 · 1d4cc1c · 1d4cc1c
1 parent b7cc8ec
commit 1d4cc1c
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 10 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
@@ -1410,7 +1410,7 @@ class FormatOps(
     private object CallSite {
 
       private val penalizeOpenNL: Policy.Pf = { case Decision(_, s) =>
-        s.map(x => if (x.isNL) x.withPenalty(1) else x)
+        s.penalizeNL(1)
       }
 
       @tailrec

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
@@ -87,15 +87,13 @@ class Router(formatOps: FormatOps) {
             }
           else if (!style.newlines.sourceIgnored && !isTripleQuote(m.left.text))
             Policy.onLeft(end, "INTERP-KEEP-NONL") {
-              case Decision(x, ss) if x.noBreak =>
-                ss.map(s => if (s.isNL) s.withPenalty(penalty) else s)
+              case Decision(x, ss) if x.noBreak => ss.penalizeNL(penalty)
             }
           else Policy ?
             (style.newlines.inInterpolation eq
               Newlines.InInterpolation.allow) &&
             Policy.onLeft(end, "INTERP-ALLOW-NL", rank = -1) {
-              case Decision(_, ss) => ss
-                  .map(s => if (s.isNL) s.withPenalty(1) else s)
+              case Decision(_, ss) => ss.penalizeNL(1)
             }
         }
         val split = Split(NoSplit, 0, policy = policy)
@@ -1276,8 +1274,7 @@ class Router(formatOps: FormatOps) {
                       if isParamClauseSite(m.leftOwner) &&
                         styleMap.at(o).binPack.bracketDefnSite
                           .exists(_ != BinPack.Site.Never) =>
-                    if (isRightCommentThenBreak(ftd)) s
-                    else s.map(x => if (x.isNL) x.withPenalty(p) else x)
+                    if (isRightCommentThenBreak(ftd)) s else s.penalizeNL(p)
                 }
               }
               getNoSplit(nextComma.map(getSlbEndOnLeft))

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Split.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Split.scala
@@ -334,4 +334,11 @@ object Split {
       fileLineStack: FileLineStack,
   ): Split = if (mod eq null) ignored else Split(mod, cost)
 
+  implicit class ImplicitSeqSplit(private val obj: Seq[Split]) extends AnyVal {
+    def penalize(penalty: Int): Seq[Split] = obj.map(_.withPenalty(penalty))
+    def penalizeIf(penalty: Int)(f: Split => Boolean): Seq[Split] = obj
+      .map(x => if (f(x)) x.withPenalty(penalty) else x)
+    def penalizeNL(penalty: Int): Seq[Split] = penalizeIf(penalty)(_.isNL)
+  }
+
 }
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/util/PolicyOps.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/util/PolicyOps.scala
@@ -23,8 +23,8 @@ object PolicyOps {
     private val checkSyntax = noSyntaxNL || !style.newlines.ignoreInSyntax
     override val f: Policy.Pf = {
       case Decision(ft, s) if penalizeLambdas || !ft.left.is[T.RightArrow] =>
-        if (checkSyntax && ft.leftHasNewline) s.map(_.withPenalty(penalty))
-        else s.map(x => if (x.isNL) x.withPenalty(penalty) else x)
+        if (checkSyntax && ft.leftHasNewline) s.penalize(penalty)
+        else s.penalizeNL(penalty)
     }
     override def prefix: String = s"PNL+$penalty"
   }
@@ -56,7 +56,7 @@ object PolicyOps {
       val nonBoolPenalty = if (TokenOps.isBoolOperator(l)) 0 else 5
       val penalty = TreeOps.nestedSelect(m.leftOwner) +
         TreeOps.nestedApplies(m.rightOwner) + nonBoolPenalty
-      s.map(x => if (x.isNL) x.withPenalty(penalty) else x)
+      s.penalizeNL(penalty)
   }
 
   /** Forces all splits up to including expire to be on a single line.