ScalafmtOptimizer: allow to do only no-slow-states

docs/configuration.md

@@ -5566,11 +5566,12 @@ for a state to be considered. A state need to qualify under at least one of
 the checks controlled by these parameters:
 
 - `runner.optimizer.pruneSlowStates`:
-  - if this flag is disabled, any state qualifies
+  - accepts `true` or `false`; also, since v3.8.4, `No`, `Yes` or `Only`
+  - if this flag is disabled (`false` or `No`), any state qualifies
   - if it is enabled:
     - if the search algorithm is ultimately unable to find a completed solution
-      (because some states might be discarded), then it will disable the flag
-      and make another attempt
+      (because some states might be discarded) and if the value is not `Only`,
+      then it will disable the flag and make another attempt
     - during the first run, when the flag is enabled, a state qualifies if it
       is not "worse" than a previously recorded "best" state for this token
     - the "best" state is the first state with a newline split at the given

scalafmt-core/shared/src/main/scala/org/scalafmt/config/ScalafmtOptimizer.scala

@@ -49,9 +49,11 @@ case class ScalafmtOptimizer(
     maxDepth: Int = 100,
     acceptOptimalAtHints: Boolean = true,
     disableOptimizationsInsideSensitiveAreas: Boolean = true,
-    pruneSlowStates: Boolean = true,
+    pruneSlowStates: ScalafmtOptimizer.PruneSlowStates =
+      ScalafmtOptimizer.PruneSlowStates.Yes,
     recurseOnBlocks: Boolean = true,
-    @annotation.ExtraName("forceConfigStyleOnOffset") @annotation.DeprecatedName(
+    @annotation.ExtraName("forceConfigStyleOnOffset")
+    @annotation.DeprecatedName(
       "forceConfigStyleMinSpan",
       "Use `callSite.minSpan` instead",
       "3.8.2",
@@ -111,4 +113,16 @@ object ScalafmtOptimizer {
       .deriveCodecEx(default).noTypos
   }
 
+  sealed abstract class PruneSlowStates
+  object PruneSlowStates {
+    case object No extends PruneSlowStates
+    case object Yes extends PruneSlowStates
+    case object Only extends PruneSlowStates
+
+    implicit val reader: ConfCodecEx[PruneSlowStates] = ReaderUtil
+      .oneOfCustom[PruneSlowStates](No, Yes, Only) { case Conf.Bool(flag) =>
+        Configured.Ok(if (flag) Yes else No)
+      }
+  }
+
 }

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

@@ -38,12 +38,12 @@ private class BestFirstSearch private (range: Set[Range])(implicit
   var deepestYet = State.start
   val best = mutable.Map.empty[Int, State]
   val visits = new Array[Int](tokens.length)
-  var keepSlowStates = !initStyle.runner.optimizer.pruneSlowStates
+  var pruneSlowStates = initStyle.runner.optimizer.pruneSlowStates
 
   /** Returns true if it's OK to skip over state.
     */
-  def shouldEnterState(curr: State): Boolean = keepSlowStates ||
-    curr.policy.noDequeue ||
+  def shouldEnterState(curr: State): Boolean = curr.policy.noDequeue ||
+    (pruneSlowStates eq ScalafmtOptimizer.PruneSlowStates.No) ||
     // TODO(olafur) document why/how this optimization works.
     best.get(curr.depth).forall(curr.possiblyBetter)
 
@@ -187,8 +187,10 @@ private class BestFirstSearch private (range: Set[Range])(implicit
               case _ => nextState
             }
             if (null ne stateToQueue) {
-              if (!keepSlowStates && depth == 0 && split.isNL) best
-                .getOrElseUpdate(curr.depth, nextState)
+              if (
+                (pruneSlowStates ne ScalafmtOptimizer.PruneSlowStates.No) &&
+                depth == 0 && split.isNL
+              ) best.getOrElseUpdate(curr.depth, nextState)
               enqueue(stateToQueue)
             }
           }
@@ -300,14 +302,15 @@ private class BestFirstSearch private (range: Set[Range])(implicit
   def getBestPath: SearchResult = {
     initStyle.runner.event(FormatEvent.Routes(routes))
     val state = {
-      def run = shortestPath(State.start, topSourceTree.tokens.last)
+      val endToken = topSourceTree.tokens.last
+      def run = shortestPath(State.start, endToken)
       val state = run
-      if (null != state || keepSlowStates) state
-      else {
-        best.clear()
-        keepSlowStates = true
+      val retry = (null eq state) &&
+        (pruneSlowStates eq ScalafmtOptimizer.PruneSlowStates.Yes)
+      if (retry) {
+        pruneSlowStates = ScalafmtOptimizer.PruneSlowStates.No
         run
-      }
+      } else state
     }
     if (null != state) {
       complete(state)(initStyle)