BestFirstSearch: use visual token width to recurse

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

@@ -41,7 +41,7 @@ private class BestFirstSearch private (range: Set[Range])(implicit
   ): Option[Token] = getEndOfBlock(ft, parensToo = true).collect {
     case close if close.left != stop && {
           // Block must span at least 3 lines to be worth recursing.
-          tokens.span(ft, close) > style.maxColumn * 3
+          tokens.width(ft, close) > style.maxColumn * 3
         } => close.left
   }
 

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

@@ -394,21 +394,31 @@ class FormatTokens(leftTok2tok: Map[TokenHash, Int])(val arr: Array[FormatToken]
   }
 
   // Maps token to number of non-whitespace bytes before the token's position.
-  private final lazy val nonWhitespaceOffset: Array[Int] = {
-    val result = new Array[Int](arr.length)
-    var curr = 0
+  private final lazy val nonWhitespaceOffset: Array[(Int, Int)] = {
+    val result = new Array[(Int, Int)](arr.length)
+    var currLength = 0
+    var currWidth = 0
     arr.foreach { t =>
-      result(t.idx) = curr
-      curr += t.left.len
+      result(t.idx) = (currLength, currWidth)
+      val (head, last) = State
+        .getColumns(t.left, t.meta.left, 0)(identity)(identity)
+      currLength += t.left.len
+      currWidth += head.max(last)
     }
     result
   }
 
-  def span(left: FormatToken, right: FormatToken): Int = {
+  private def spanOrWidth(left: FormatToken, right: FormatToken)(
+      f: ((Int, Int)) => Int,
+  ): Int = {
     val lastIdx = nonWhitespaceOffset.length - 1
     val rightIdx = if (right.idx >= lastIdx) lastIdx else right.idx + 1
-    nonWhitespaceOffset(rightIdx) - nonWhitespaceOffset(left.idx)
+    f(nonWhitespaceOffset(rightIdx)) - f(nonWhitespaceOffset(left.idx))
   }
+  def width(left: FormatToken, right: FormatToken): Int =
+    spanOrWidth(left, right)(_._2)
+  def span(left: FormatToken, right: FormatToken): Int =
+    spanOrWidth(left, right)(_._1)
   def span(tokens: Tokens): Int =
     if (tokens.isEmpty) 0 else span(getHeadImpl(tokens), getLastImpl(tokens))
   @inline

scalafmt-tests-community/scala2/src/test/scala/org/scalafmt/community/scala2/CommunityScala2Suite.scala

@@ -9,7 +9,7 @@ abstract class CommunityScala2Suite(name: String)
 
 class CommunityScala2_12Suite extends CommunityScala2Suite("scala-2.12") {
 
-  override protected def totalStatesVisited: Option[Int] = Some(35259833)
+  override protected def totalStatesVisited: Option[Int] = Some(35259819)
 
   override protected def builds =
     Seq(getBuild("v2.12.20", dialects.Scala212, 1277))
@@ -18,7 +18,7 @@ class CommunityScala2_12Suite extends CommunityScala2Suite("scala-2.12") {
 
 class CommunityScala2_13Suite extends CommunityScala2Suite("scala-2.13") {
 
-  override protected def totalStatesVisited: Option[Int] = Some(43931406)
+  override protected def totalStatesVisited: Option[Int] = Some(43931350)
 
   override protected def builds =
     Seq(getBuild("v2.13.14", dialects.Scala213, 1287))

scalafmt-tests-community/scala3/src/test/scala/org/scalafmt/community/scala3/CommunityScala3Suite.scala

@@ -9,15 +9,15 @@ abstract class CommunityScala3Suite(name: String)
 
 class CommunityScala3_2Suite extends CommunityScala3Suite("scala-3.2") {
 
-  override protected def totalStatesVisited: Option[Int] = Some(32916614)
+  override protected def totalStatesVisited: Option[Int] = Some(32918646)
 
   override protected def builds = Seq(getBuild("3.2.2", dialects.Scala32, 791))
 
 }
 
 class CommunityScala3_3Suite extends CommunityScala3Suite("scala-3.3") {
 
-  override protected def totalStatesVisited: Option[Int] = Some(35515588)
+  override protected def totalStatesVisited: Option[Int] = Some(35518295)
 
   override protected def builds = Seq(getBuild("3.3.3", dialects.Scala33, 861))
 

scalafmt-tests-community/spark/src/test/scala/org/scalafmt/community/spark/CommunitySparkSuite.scala

@@ -9,7 +9,7 @@ abstract class CommunitySparkSuite(name: String)
 
 class CommunitySpark3_4Suite extends CommunitySparkSuite("spark-3.4") {
 
-  override protected def totalStatesVisited: Option[Int] = Some(71541585)
+  override protected def totalStatesVisited: Option[Int] = Some(71540803)
 
   override protected def builds =
     Seq(getBuild("v3.4.1", dialects.Scala213, 2585))
@@ -18,7 +18,7 @@ class CommunitySpark3_4Suite extends CommunitySparkSuite("spark-3.4") {
 
 class CommunitySpark3_5Suite extends CommunitySparkSuite("spark-3.5") {
 
-  override protected def totalStatesVisited: Option[Int] = Some(75679293)
+  override protected def totalStatesVisited: Option[Int] = Some(75678511)
 
   override protected def builds =
     Seq(getBuild("v3.5.3", dialects.Scala213, 2756))

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, 1487968, "total explored")
+      assertEquals(explored, 1488002, "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 :'(