BestFirstSearch: don't give up too soon

scalafmt-core/shared/src/main/scala/org/scalafmt/Scalafmt.scala

@@ -69,8 +69,7 @@ object Scalafmt {
         val formatOps = new FormatOps(tree, style)
         runner.event(CreateFormatOps(formatOps))
         val formatWriter = new FormatWriter(formatOps)
-        val search = new BestFirstSearch(formatOps, range, formatWriter)
-        val partial = search.getBestPath
+        val partial = BestFirstSearch(formatOps, range, formatWriter)
         val formattedString = formatWriter.mkString(partial.state)
         val correctedFormattedString =
           if ((style.lineEndings == preserve && isWindows) ||

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

@@ -18,7 +18,7 @@ import org.scalafmt.util.TreeOps
 /**
   * Implements best first search to find optimal formatting.
   */
-class BestFirstSearch(
+private class BestFirstSearch private (
     val formatOps: FormatOps,
     range: Set[Range],
     formatWriter: FormatWriter
@@ -105,12 +105,6 @@ class BestFirstSearch(
     state.column << 8 | state.indentation
   }
 
-  def hasReachedEof(state: State, stop: Token): Boolean =
-    state.depth >= tokens.length || {
-      val token = tokens(state.depth)
-      token.left.start >= stop.start && token.left.end > token.left.start
-    }
-
   val memo = mutable.Map.empty[(Int, StateHash), State]
 
   def shortestPathMemo(start: State, stop: Token, depth: Int, maxCost: Int)(
@@ -122,7 +116,7 @@ class BestFirstSearch(
     else {
       // Only update state if it reached stop.
       val nextState = shortestPath(start, stop, depth, maxCost)
-      if (tokens(nextState.depth).left != stop) None
+      if (null == nextState) None
       else {
         memo.update(key, nextState)
         Some(nextState)
@@ -139,11 +133,18 @@ class BestFirstSearch(
       depth: Int = 0,
       maxCost: Int = Integer.MAX_VALUE
   ): State = {
-    val Q = new mutable.PriorityQueue[State]()
-    var result = start
+    def newGeneration = new mutable.PriorityQueue[State]()
+    var Q = newGeneration
+    var generations: List[mutable.PriorityQueue[State]] = Nil
+    def addGeneration =
+      if (Q.nonEmpty) {
+        generations = Q :: generations
+        Q = newGeneration
+      }
     Q += start
+
     // TODO(olafur) this while loop is waaaaaaaaaaaaay tooo big.
-    while (Q.nonEmpty) {
+    while (true) {
       val curr = Q.dequeue()
       explored += 1
       runner.event(Explored(explored, depth, Q.size))
@@ -153,11 +154,15 @@ class BestFirstSearch(
           formatWriter.mkString(deepestYet),
           tokens(deepestYet.depth).left
         )
-      if (hasReachedEof(curr, stop)) {
-        result = curr
-        Q.clear()
-      } else if (shouldEnterState(curr)) {
-        val splitToken = tokens(curr.depth)
+      val splitToken =
+        if (curr.depth >= tokens.length) null else tokens(curr.depth)
+      if (null == splitToken ||
+        splitToken.left.start >= stop.start &&
+        splitToken.left.start < splitToken.left.end) {
+        return curr
+      }
+
+      if (shouldEnterState(curr)) {
         val style = styleMap.at(splitToken)
         if (curr.depth > deepestYet.depth) {
           deepestYet = curr
@@ -171,7 +176,8 @@ class BestFirstSearch(
             dequeueOnNewStatements &&
             dequeueSpots.contains(tokenHash) &&
             (depth > 0 || !isInsideNoOptZone(splitToken)))
-            Q.clear()
+            if (depth == 0) addGeneration
+            else Q.clear()
         }
 
         val blockClose = shouldRecurseOnBlock(splitToken, stop, style)
@@ -180,8 +186,6 @@ class BestFirstSearch(
             .foreach(Q.enqueue(_))
         else if (escapeInPathologicalCases &&
           visits(curr.depth) > maxVisitsPerToken) {
-          Q.clear()
-          best.clear()
           runner.event(CompleteFormat(explored, deepestYet))
           throw SearchStateExploded(
             deepestYet,
@@ -216,7 +220,7 @@ class BestFirstSearch(
                     nextState.split.cost == 0 =>
                 val nextNextState =
                   shortestPath(nextState, token, depth + 1, maxCost = 0)
-                if (hasReachedEof(nextNextState, token)) {
+                if (null != nextNextState) {
                   optimalNotFound = false
 //                  logger.elem(split, splitToken, formatWriter.mkString(nextNextState.splits), tokens(nextNextState.splits.length))
                   Q.enqueue(nextNextState)
@@ -235,13 +239,23 @@ class BestFirstSearch(
           }
         }
       }
+
+      if (Q.isEmpty) {
+        if (generations.isEmpty)
+          return null
+
+        Q = generations.head
+        generations = generations.tail
+      }
     }
-    result
+
+    // unreachable
+    null
   }
 
   def getBestPath: SearchResult = {
     val state = shortestPath(State.start, tree.tokens.last)
-    if (state.depth == tokens.length) {
+    if (null != state) {
       runner.event(CompleteFormat(explored, state))
       SearchResult(state, reachedEOF = true)
     } else {
@@ -251,7 +265,6 @@ class BestFirstSearch(
         deepestYet.policy.execute(Decision(tok, nextSplits))
       val msg = s"""UNABLE TO FORMAT,
                    |tok=$tok
-                   |state.depth=${state.depth}
                    |toks.length=${tokens.length}
                    |deepestYet.length=${deepestYet.depth}
                    |policies=${deepestYet.policy.policies}
@@ -268,3 +281,14 @@ class BestFirstSearch(
 }
 
 case class SearchResult(state: State, reachedEOF: Boolean)
+
+object BestFirstSearch {
+
+  def apply(
+      formatOps: FormatOps,
+      range: Set[Range],
+      formatWriter: FormatWriter
+  ): SearchResult =
+    new BestFirstSearch(formatOps, range, formatWriter).getBestPath
+
+}

scalafmt-tests/src/test/resources/default/Lambda.stat

@@ -386,3 +386,62 @@ testAsync("Resource.make is equivalent to a partially applied bracket") { implic
     }
   }
 }
+<<< #448
+object Foo {
+  code.collect {
+    case d: Defn.Def if (d match {
+      case _ => false
+
+
+    }) &&
+      d=>
+  }
+}
+>>>
+object Foo {
+  code.collect {
+    case d: Defn.Def
+        if (d match {
+          case _ => false
+
+        }) &&
+          d =>
+  }
+}
+<<< #1485
+style = IntelliJ
+maxColumn = 120
+===
+class A {
+  def foo: B =
+    a.b() { c =>
+        val d = c.add(
+          D[E](
+            1, {
+              case conditionNameShouldLongEnough =>
+                foo1(objectNameShouldLongEnough, objectNameShouldLongEnough)
+                0
+              case conditionNameShouldLongEnough =>
+                foo1(objectNameShouldLongEnough, objectNameShouldLongEnough)
+                1
+            }
+          )
+        )
+      }
+      .f("f")
+}
+>>>
+class A {
+  def foo: B =
+    a.b() { c =>
+        val d = c.add(D[E](1, {
+          case conditionNameShouldLongEnough =>
+            foo1(objectNameShouldLongEnough, objectNameShouldLongEnough)
+            0
+          case conditionNameShouldLongEnough =>
+            foo1(objectNameShouldLongEnough, objectNameShouldLongEnough)
+            1
+        }))
+      }
+      .f("f")
+}

scalafmt-tests/src/test/resources/newlines/afterCurlyLambdaNever.stat

@@ -37,3 +37,37 @@ def f = {
 def f = {
   something.call { x => g(x) }
 }
+<<< #1717
+maxColumn = 100
+===
+class Bar {
+  def foo(request: Request): RpcResult[Response] = {
+    repository
+      .flatMap { campaigns =>
+        request match {
+          case Nil => connection.pure(Right(toResponse(campaigns, Map.empty)))
+          case _ =>
+            getProfileEmails(request.securityContext, profileIds).to[ConnectionIO].map {
+              profileEmails =>
+                Right(toResponse(campaigns, profileEmails))
+            }
+        }
+      }.transact(managerTransactor)
+  }
+}
+>>>
+class Bar {
+  def foo(request: Request): RpcResult[Response] = {
+    repository
+      .flatMap { campaigns =>
+        request match {
+          case Nil => connection.pure(Right(toResponse(campaigns, Map.empty)))
+          case _ =>
+            getProfileEmails(request.securityContext, profileIds).to[ConnectionIO].map {
+              profileEmails => Right(toResponse(campaigns, profileEmails))
+            }
+        }
+      }
+      .transact(managerTransactor)
+  }
+}

scalafmt-tests/src/test/resources/newlines/beforeConfigSingleArgParenLambdaParamsF_danglingT.stat

@@ -702,3 +702,92 @@ object a {
         _ => e
       )
 }
+<<< #1599 1: original, with partial func
+object A {
+  def foo() = {
+    x.map(_.y(
+      abcd,
+      { case ((abcdefghij, aswbasdfaw), asfda) => aswdf},
+      { case (abcdefghij, sadfasdass) => (asdfa.sadvfs(abcdefghij).get, asdfasdfasfdasda.asdfas(asdfasdaas).get) },
+      foo
+    ))
+  }
+}
+>>>
+object A {
+  def foo() = {
+    x.map(
+      _.y(
+        abcd,
+        { case ((abcdefghij, aswbasdfaw), asfda) => aswdf }, {
+          case (abcdefghij, sadfasdass) =>
+            (
+              asdfa.sadvfs(abcdefghij).get,
+              asdfasdfasfdasda.asdfas(asdfasdaas).get
+            )
+        },
+        foo
+      )
+    )
+  }
+}
+<<< #1599 2: modified, with partial func inside block
+object A {
+  def foo() = {
+    x.map(_.y(
+      abcd,
+      {{ case ((abcdefghij, aswbasdfaw), asfda) => aswdf}},
+      {{ case (abcdefghij, sadfasdass) => (asdfa.sadvfs(abcdefghij).get, asdfasdfasfdasda.asdfas(asdfasdaas).get) }},
+      foo
+    ))
+  }
+}
+>>>
+object A {
+  def foo() = {
+    x.map(
+      _.y(
+        abcd,
+        { { case ((abcdefghij, aswbasdfaw), asfda) => aswdf } }, {
+          {
+            case (abcdefghij, sadfasdass) =>
+              (
+                asdfa.sadvfs(abcdefghij).get,
+                asdfasdfasfdasda.asdfas(asdfasdaas).get
+              )
+          }
+        },
+        foo
+      )
+    )
+  }
+}
+<<< #1599 3: modified, with lambda
+object A {
+  def foo() = {
+    x.map(_.y(
+      abcd,
+      { case_abcdefghij_aswbasdfaw_asfda => aswdf},
+      { case_abcdefghij_sadfasdass => (asdfa.sadvfs(abcdefghij).get, asdfasdfasfdasda.asdfas(asdfasdaas).get) },
+      foo
+    ))
+  }
+}
+>>>
+object A {
+  def foo() = {
+    x.map(
+      _.y(
+        abcd,
+        { case_abcdefghij_aswbasdfaw_asfda => aswdf }, {
+          case_abcdefghij_sadfasdass =>
+            (
+              asdfa.sadvfs(abcdefghij).get,
+              asdfasdfasfdasda.asdfas(asdfasdaas).get
+            )
+        },
+        foo
+      )
+    )
+  }
+}

scalafmt-tests/src/test/resources/vertical-multiline/VerticalMultilineDefnSite.stat

@@ -14,3 +14,28 @@ def format_![T <: Tree](
 )(implicit ev: Parse[T],
   ev2: EC
 ): String
+<<< #1462
+maxColumn = 91
+===
+object Test {
+
+  def create(
+    somethingInfo: foo.T = Something("my-something", Some("test-something"), "something", Some(Format("something", "something")),
+      Some(something.Something("something", "something", "something", "something")))
+  ): Unit = {
+    w(foo)
+  }
+
+}
+>>>
+object Test {
+
+  def create(
+    somethingInfo: foo.T = Something("my-something", Some("test-something"), "something",
+      Some(Format("something", "something")), Some(something.Something("something",
+          "something", "something", "something")))
+  ): Unit = {
+    w(foo)
+  }
+
+}