BestFirstSearch: optimal has low-penalty NL split

scalafmt-core/shared/src/main/scala/org/scalafmt/internal/BestFirstSearch.scala

@@ -192,11 +192,16 @@ private class BestFirstSearch private (range: Set[Range])(implicit
           case Left(x) => return Left(killOnFail(opt, x))
         }
       }
-    val furtherState = traverseSameLine(nextNextState)
-    if (null eq furtherState) Left(killOnFail(opt, nextNextState))
-    else if (furtherState.appliedPenalty > nextNextState.appliedPenalty)
-      Left(nextNextState)
-    else Either.cond(!opt.recurseOnly, furtherState, furtherState)
+    def checkPenalty(state: State, orElse: => Either[State, State]) =
+      if (state.appliedPenalty > nextNextState.appliedPenalty)
+        Left(nextNextState)
+      else orElse
+    traverseSameLine(nextNextState) match {
+      case x @ Left(s) =>
+        if (s eq null) Left(killOnFail(opt, nextNextState))
+        else checkPenalty(s, x)
+      case x @ Right(s) => checkPenalty(s, if (opt.recurseOnly) Left(s) else x)
+    }
   }
 
   private def getActiveSplits(ft: FormatToken, state: State, maxCost: Int)(
@@ -218,20 +223,27 @@ private class BestFirstSearch private (range: Set[Range])(implicit
     splits.sortBy(_.costWithPenalty)
   }
 
+  private def stateAsOptimal(state: State, splits: Seq[Split]) = {
+    val isOptimal = splits.exists { s =>
+      s.isNL && s.penalty < Constants.ShouldBeNewline
+    } || tokens.isRightCommentWithBreak(tokens(state.depth))
+    if (isOptimal) Some(state) else None
+  }
+
   /** Follow states having single active non-newline split
     */
   @tailrec
   private def traverseSameLine(
       state: State,
-  )(implicit queue: StateQueue): State =
-    if (state.depth >= tokens.length) state
+  )(implicit queue: StateQueue): Either[State, State] =
+    if (state.depth >= tokens.length) Right(state)
     else {
       val splitToken = tokens(state.depth)
       implicit val style: ScalafmtConfig = styleMap.at(splitToken)
       getActiveSplits(splitToken, state, Int.MaxValue) match {
-        case Seq() => null // dead end if empty
+        case Seq() => Left(null) // dead end if empty
         case Seq(split) =>
-          if (split.isNL) state
+          if (split.isNL) Right(state)
           else {
             implicit val nextState: State =
               getNext(state, split, allAltAreNL = false)
@@ -240,29 +252,34 @@ private class BestFirstSearch private (range: Set[Range])(implicit
         case ss
             if state.appliedPenalty == 0 &&
               RightParenOrBracket(splitToken.right) =>
-          traverseSameLineZeroCost(ss.filter(_.costWithPenalty == 0), state)
-        case _ => state
+          traverseSameLineZeroCost(ss, state)
+        case ss => stateAsOptimal(state, ss).toRight(state)
       }
     }
 
   @tailrec
   private def traverseSameLineZeroCost(
       splits: Seq[Split],
       state: State,
-  )(implicit style: ScalafmtConfig, queue: StateQueue): State = splits match {
-    case Seq(split) if !split.isNL =>
-      val nextState: State = getNext(state, split, allAltAreNL = false)
-      if (nextState.split.costWithPenalty > 0) state
-      else if (nextState.depth >= tokens.length) nextState
-      else {
-        val nextToken = tokens(nextState.depth)
-        if (RightParenOrBracket(nextToken.right)) {
-          implicit val style: ScalafmtConfig = styleMap.at(nextToken)
-          val nextSplits = getActiveSplits(nextToken, nextState, maxCost = 0)
-          traverseSameLineZeroCost(nextSplits, nextState)
-        } else nextState
-      }
-    case _ => state
+      nlState: => Option[State] = None,
+  )(implicit style: ScalafmtConfig, queue: StateQueue): Either[State, State] = {
+    def newNlState: Option[State] = stateAsOptimal(state, splits).orElse(nlState)
+    splits.filter(_.costWithPenalty == 0) match {
+      case Seq(split) if !split.isNL =>
+        val nextState: State = getNext(state, split, allAltAreNL = false)
+        if (nextState.split.costWithPenalty > 0) newNlState.toRight(state)
+        else if (nextState.depth >= tokens.length) Right(nextState)
+        else {
+          val nextToken = tokens(nextState.depth)
+          if (RightParenOrBracket(nextToken.right)) {
+            implicit val style: ScalafmtConfig = styleMap.at(nextToken)
+            val nextSplits =
+              getActiveSplits(nextToken, nextState, maxCost = Int.MaxValue)
+            traverseSameLineZeroCost(nextSplits, nextState, newNlState)
+          } else newNlState.toRight(nextState)
+        }
+      case _ => newNlState.toRight(state)
+    }
   }
 
   def getBestPath: SearchResult = {

scalafmt-tests/shared/src/test/resources/default/String.stat

@@ -504,9 +504,9 @@ object a {
 }
 >>>
 object a {
-  s"foo1 ${bar} foo2 ${baz(bar)} foo3 ${
-      qux(baz, bar)
-    } foo4"
+  s"foo1 ${
+      bar
+    } foo2 ${baz(bar)} foo3 ${qux(baz, bar)} foo4"
   s"""
     |foo1 ${bar} foo2 ${baz(bar)} foo3 ${
       qux(baz, bar)

scalafmt-tests/shared/src/test/resources/scalajs/DefDef.stat

@@ -240,8 +240,8 @@ object a {
 >>>
 object a {
   object b {
-    def applyDynamicNamed(name: String)(
-        fields: (String, Any)*): Object with Dynamic = sys.error("stub")
+    def applyDynamicNamed(name: String)(fields: (String,
+            Any)*): Object with Dynamic = sys.error("stub")
   }
 }
 <<< weird breaking #252, oneline
@@ -255,8 +255,8 @@ object a {
 >>>
 object a {
   object b {
-    def applyDynamicNamed(name: String)(
-        fields: (String, Any)*): Object with Dynamic = sys.error("stub")
+    def applyDynamicNamed(name: String)(fields: (String,
+            Any)*): Object with Dynamic = sys.error("stub")
   }
 }
 <<< toFunction8 #255 avoidInResultType, sometimesBeforeColonInMethodReturnType
@@ -648,8 +648,9 @@ newlines.source = keep
 ===
 @inline implicit def enrichAsScalaFromToIntFunction[T](jf: java.util.function.ToIntFunction[T]): RichToIntFunctionAsFunction1[T] = new RichToIntFunctionAsFunction1[T](jf)
 >>>
-@inline implicit def enrichAsScalaFromToIntFunction[T](jf: java.util
-      .function.ToIntFunction[T]): RichToIntFunctionAsFunction1[T] =
+@inline implicit def enrichAsScalaFromToIntFunction[T](
+    jf: java.util.function.ToIntFunction[
+        T]): RichToIntFunctionAsFunction1[T] =
   new RichToIntFunctionAsFunction1[T](jf)
 <<< #4133 overflow selects with apply-type 2
 maxColumn = 70
@@ -658,8 +659,8 @@ newlines.source = keep
 def asScalaFromDoubleFunction[R](jf: java.util.function.DoubleFunction[R]): scala.Function1[java.lang.Double, R]
 >>>
 def asScalaFromDoubleFunction[R](
-    jf: java.util.function.DoubleFunction[R]): scala.Function1[
-    java.lang.Double, R]
+    jf: java.util.function.DoubleFunction[
+        R]): scala.Function1[java.lang.Double, R]
 <<< #4133 complex nested applies with long infix args
 maxColumn = 80
 newlines.source = keep
@@ -981,3 +982,55 @@ val template: Elems = List(
                         deprValueMembers))
         )
     ))
+<<< #4133 implicit with bracket and comment in type
+maxColumn = 78
+newlines.source = keep
+===
+override def flatten[B](implicit asTraversable: A => /*<:<!!!*/ GenTraversableOnce[B]): Stream[B] = {
+  // foo
+}
+>>>
+override def flatten[B](
+    implicit asTraversable: A => /*<:<!!!*/ GenTraversableOnce[
+        B]): Stream[B] = {
+  // foo
+}
+<<< #4133 no implicit with bracket and comment in type
+maxColumn = 78
+newlines.source = keep
+===
+override def flatten[B](implicit_asTraversable: A => /*<:<!!!*/ GenTraversableOnce[B]): Stream[B] = {
+  // foo
+}
+>>>
+override def flatten[B](
+    implicit_asTraversable: A => /*<:<!!!*/ GenTraversableOnce[
+        B]): Stream[B] = {
+  // foo
+}
+<<< #4133 no implicit with bracket and no comment in type
+maxColumn = 78
+newlines.source = keep
+===
+override def flatten[B](implicit_asTraversable: A => hadComment_GenTraversableOnce[B]): Stream[B] = {
+  // foo
+}
+>>>
+override def flatten[B](
+    implicit_asTraversable: A => hadComment_GenTraversableOnce[
+        B]): Stream[B] = {
+  // foo
+}
+<<< #4133 no implicit with bracket and no comment in type, no arrow
+maxColumn = 78
+newlines.source = keep
+===
+override def flatten[B](implicit_asTraversable: A_aw_hadComment_GenTraversableOnce[B]): Stream[B] = {
+  // foo
+}
+>>>
+override def flatten[B](
+    implicit_asTraversable: A_aw_hadComment_GenTraversableOnce[
+        B]): Stream[B] = {
+  // foo
+}