From a0a013d23820be2e07d5b3f4aa882e427df83ab7 Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Sun, 27 Oct 2024 07:52:58 -0700
Subject: [PATCH] BestFirstSearch: try memorize partial functions

---
 .../main/scala/org/scalafmt/internal/BestFirstSearch.scala   | 4 ++++
 .../shared/src/test/resources/newlines/source_fold.stat      | 5 ++---
 .../shared/src/test/scala/org/scalafmt/FormatTests.scala     | 2 +-
 3 files changed, 7 insertions(+), 4 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
index 51e77980a..66c0bdd93 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/BestFirstSearch.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/BestFirstSearch.scala
@@ -319,6 +319,8 @@ object BestFirstSearch {
     var expire: FormatToken = null
     @inline
     def addRange(t: Token): Unit = expire = tokens.matching(t)
+    @inline
+    def addBlock(t: Token): Unit = result.getOrElseUpdate(t, false)
     tokens.foreach {
       case ft if expire ne null =>
         if (ft eq expire) expire = null else result.update(ft.left, true)
@@ -332,6 +334,8 @@ object BestFirstSearch {
           // Type compounds can be inside defn.defs
           case lo: meta.Stat.Block if lo.parent.is[Type.Refine] => addRange(t)
           case _: Type.Refine => addRange(t)
+          case _: Term.PartialFunction if styleMap.at(t).newlines.fold =>
+            addBlock(t)
           case _ =>
         }
       case _ =>
diff --git a/scalafmt-tests/shared/src/test/resources/newlines/source_fold.stat b/scalafmt-tests/shared/src/test/resources/newlines/source_fold.stat
index 21aa084f2..e4278c385 100644
--- a/scalafmt-tests/shared/src/test/resources/newlines/source_fold.stat
+++ b/scalafmt-tests/shared/src/test/resources/newlines/source_fold.stat
@@ -9460,7 +9460,7 @@ object a {
         }
   }
 }
->>> { stateVisits = 2428, stateVisits2 = 2428 }
+>>> { stateVisits = 8595, stateVisits2 = 8595 }
 object a {
   private object MemoMap {
     def make(implicit trace: Trace): UIO[MemoMap] = Ref.Synchronized
@@ -10059,7 +10059,6 @@ object a {
 preset = default
 maxColumn = 80
 newlines.source = fold
-runner.maxStateVisits = 80000
 #newlines.avoidForSimpleOverflow = all
 ===
 system.dynamicAccess.createInstanceFor[Serializer](fqn, Nil).recoverWith {
@@ -10091,7 +10090,7 @@ system.dynamicAccess.createInstanceFor[Serializer](fqn, Nil).recoverWith {
             }
       }
 }
->>> { stateVisits = 72590, stateVisits2 = 72590 }
+>>> { stateVisits = 5817, stateVisits2 = 5817 }
 system.dynamicAccess.createInstanceFor[Serializer](fqn, Nil).recoverWith {
   case _: NoSuchMethodException => system.dynamicAccess
       .createInstanceFor[Serializer](
diff --git a/scalafmt-tests/shared/src/test/scala/org/scalafmt/FormatTests.scala b/scalafmt-tests/shared/src/test/scala/org/scalafmt/FormatTests.scala
index 7ede4e1c4..98359ac80 100644
--- 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, 1620763, "total explored")
+      assertEquals(explored, 1499555, "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 :'(