From 4ed7ed365a274a33b54d02984e9d8ea1903ed22e Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Fri, 4 Oct 2024 17:24:05 -0700
Subject: [PATCH] FormatTokens: move distance from FormatOps

---
 .../scalafmt/internal/BestFirstSearch.scala   |  2 +-
 .../org/scalafmt/internal/FormatOps.scala     | 16 +----------
 .../org/scalafmt/internal/FormatTokens.scala  | 28 +++++++++++++++++--
 3 files changed, 28 insertions(+), 18 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 cb3d7ffcee..9e05afa16d 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
@@ -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.
-          distance(ft.left, close.left) > style.maxColumn * 3
+          tokens.distance(ft, close) > style.maxColumn * 3
         } => close.left
   }
 
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
index b66263c283..972a896847 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
@@ -40,20 +40,9 @@ class FormatOps(
   implicit val (tokens: FormatTokens, styleMap: StyleMap) =
     FormatTokens(topSourceTree.tokens, owners)(initStyle)
   import tokens._
-  private val usedTokens = tokens.head.left +: tokens.map(_.right)
 
   private[internal] val soft = new SoftKeywordClasses(dialect)
   private[internal] val statementStarts = getStatementStarts(topSourceTree, soft)
-  // Maps token to number of non-whitespace bytes before the token's position.
-  private final val nonWhitespaceOffset: Map[T, Int] = {
-    val resultB = Map.newBuilder[T, Int]
-    var curr = 0
-    usedTokens.foreach { t =>
-      resultB += (t -> curr)
-      curr += (t.end - t.start)
-    }
-    resultB.result()
-  }
 
   val (forceConfigStyle, emptyQueueSpots) = getForceConfigStyle
 
@@ -897,9 +886,6 @@ class FormatOps(
     }
   }
 
-  def distance(left: T, right: T): Int = nonWhitespaceOffset(right) -
-    nonWhitespaceOffset(left)
-
   def typeTemplateSplits(template: Template, indentIfSecond: Int)(implicit
       fileLine: FileLine,
       ft: FormatToken,
@@ -1030,7 +1016,7 @@ class FormatOps(
         val values = clause.values
         if (
           values.lengthCompare(cfg.minCount) >= 0 &&
-          (cfg.minSpan == 0 || cfg.minSpan < distance(ftOpen.left, close.left))
+          (cfg.minSpan == 0 || cfg.minSpan < distance(ftOpen, close))
         ) {
           forces += ftOpen.meta.idx
           values.foreach(x => clearQueues += getHead(x).meta.idx)
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala
index 2e9a777190..22ec161356 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala
@@ -236,8 +236,10 @@ class FormatTokens(leftTok2tok: Map[TokenHash, Int])(val arr: Array[FormatToken]
     else getOnOrAfterOwned(nextFt, tree)
   }
 
+  @inline
+  private def getHeadImpl(tokens: Tokens): FormatToken = after(tokens.head)
   def getHead(tokens: Tokens, tree: Tree): FormatToken =
-    getOnOrBeforeOwned(after(tokens.head), tree)
+    getOnOrBeforeOwned(getHeadImpl(tokens), tree)
   @inline
   def getHead(tree: Tree): FormatToken = getHead(tree.tokens, tree)
 
@@ -246,8 +248,11 @@ class FormatTokens(leftTok2tok: Map[TokenHash, Int])(val arr: Array[FormatToken]
   @inline
   def getHeadOpt(tree: Tree): Option[FormatToken] = getHeadOpt(tree.tokens, tree)
 
+  @inline
+  private def getLastImpl(tokens: Tokens): FormatToken =
+    apply(findLastVisibleToken(tokens))
   def getLast(tokens: Tokens, tree: Tree): FormatToken =
-    getOnOrAfterOwned(apply(findLastVisibleToken(tokens)), tree)
+    getOnOrAfterOwned(getLastImpl(tokens), tree)
   @inline
   def getLast(tree: Tree): FormatToken = getLast(tree.tokens, tree)
 
@@ -367,6 +372,25 @@ class FormatTokens(leftTok2tok: Map[TokenHash, Int])(val arr: Array[FormatToken]
   def isAttachedCommentThenBreak(ft: FormatToken): Boolean = ft.noBreak &&
     isRightCommentThenBreak(ft)
 
+  // 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
+    arr.foreach { t =>
+      result(t.idx) = curr
+      curr += t.left.len
+    }
+    result
+  }
+
+  def distance(left: FormatToken, right: FormatToken): Int =
+    nonWhitespaceOffset(right.idx) - nonWhitespaceOffset(left.idx)
+  def distance(tokens: Tokens): Int =
+    if (tokens.isEmpty) 0
+    else distance(getHeadImpl(tokens), getLastImpl(tokens))
+  @inline
+  def distance(tree: Tree): Int = distance(tree.tokens)
+
 }
 
 object FormatTokens {