FormatTokensRewrite: reuse index claimed earlier

We just need to check that this index is not occupied by another format token, which could happen if a rule staked claim for this format token previously.
scalameta · Mar 27, 2024 · d749bac · d749bac
1 parent cfebea9
commit d749bac
Showing 1 changed file with 22 additions and 6 deletions.
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/FormatTokensRewrite.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/FormatTokensRewrite.scala
@@ -318,11 +318,8 @@ object FormatTokensRewrite {
     private[rewrite] def claimedRule(ftIdx: Int): Option[Replacement] =
       claimed.get(ftIdx).map(tokens.apply).filter(_ ne null)
 
-    @inline
     private[rewrite] def claim(ftIdx: Int, repl: Replacement): Int = {
-      val idx = tokens.length
-      claimed.update(ftIdx, idx)
-      tokens.append(
+      justClaim(ftIdx) {
         if (repl eq null) null
         else
           (repl.how match {
@@ -337,8 +334,27 @@ object FormatTokensRewrite {
               }
             case _ => None
           }).getOrElse(repl)
-      )
-      idx
+      }
+    }
+
+    private def justClaim(ftIdx: Int)(repl: Replacement): Int = {
+      val idx = tokens.length
+      val claimedIdx = claimed.getOrElseUpdate(ftIdx, idx)
+      val preClaimed = claimedIdx < idx
+      if (
+        preClaimed && {
+          val oldrepl = tokens(claimedIdx)
+          oldrepl != null && oldrepl.idx == ftIdx
+        }
+      ) {
+        tokens(claimedIdx) = repl
+        claimedIdx
+      } else {
+        if (preClaimed)
+          claimed.update(ftIdx, idx)
+        tokens.append(repl)
+        idx
+      }
     }
 
     private[FormatTokensRewrite] def applyRule(