BestFirstSearch: memorize only if full cost

scalameta · Oct 27, 2024 · 39c8d19 · 39c8d19
1 parent 7c4dfe4
commit 39c8d19
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 10 deletions.
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/BestFirstSearch.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/BestFirstSearch.scala
@@ -45,7 +45,7 @@ private class BestFirstSearch private (range: Set[Range])(implicit
         } => close.left
   }
 
-  private val memo = mutable.Map.empty[Long, State]
+  private val memo = mutable.Map.empty[Long, Option[State]]
 
   def shortestPathMemo(
       start: State,
@@ -55,12 +55,14 @@ private class BestFirstSearch private (range: Set[Range])(implicit
   ): Option[State] = {
     val key = (start.indentation & 0xffL) | (start.column & 0xffffffL) << 8 |
       (start.depth & 0xffffffffL) << 32
-    memo.get(key).orElse {
-      // Only update state if it reached stop.
-      val nextState = shortestPath(start, stop, depth, isOpt).toOption
-      nextState.foreach(memo.update(key, _))
-      nextState
-    }
+    def orElse =
+      if (isOpt) None
+      else {
+        val nextState = shortestPath(start, stop, depth).toOption
+        memo.update(key, nextState)
+        nextState
+      }
+    memo.getOrElse(key, orElse)
   }
 
   /** Runs best first search to find lowest penalty split.

diff --git a/scalafmt-tests/shared/src/test/resources/default/Advanced.stat b/scalafmt-tests/shared/src/test/resources/default/Advanced.stat
@@ -377,7 +377,7 @@ private def withNewLocalDefs = {
             }))
         }))
   }
->>> { stateVisits = 3779 }
+>>> { stateVisits = 3750, stateVisits2 = 3750 }
 val createIsArrayOfStat = {
   envFieldDef(
       "isArrayOf",

diff --git a/scalafmt-tests/shared/src/test/resources/default/Apply.stat b/scalafmt-tests/shared/src/test/resources/default/Apply.stat
@@ -273,7 +273,7 @@ class ResolutionCopier(x: Int) {
             toNode.rightBracket));
   }
 }
->>> { stateVisits = 57714 }
+>>> { stateVisits = 21703, stateVisits2 = 21703 }
 class ResolutionCopier(x: Int) {
 
   def visitClassDeclaration(node: ClassDeclaration): Boolean = {

diff --git 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 {
     val explored = Debug.explored.get()
     logger.debug(s"Total explored: $explored")
     if (!onlyUnit && !onlyManual)
-      assertEquals(explored, 1488077, "total explored")
+      assertEquals(explored, 1413371, "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 :'(