FormatTokensRewrite: add getClaimed method

It makes sure that the replacement stored at that index refers to the same format token. The opposite could happen if a rule staked claim for this format token previously.
scalameta · Mar 27, 2024 · 49e8271 · 49e8271
1 parent c18045c
commit 49e8271
Showing 1 changed file with 10 additions and 2 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
@@ -300,6 +300,15 @@ object FormatTokensRewrite {
     private[FormatTokensRewrite] val tokens =
       new mutable.ArrayBuffer[Replacement]()
 
+    private[rewrite] def getClaimed(ftIdx: Int): Option[(Int, Replacement)] =
+      claimed.get(ftIdx) match {
+        case Some(x) =>
+          val repl = tokens(x)
+          val ok = (repl eq null) || repl.ft.meta.idx == ftIdx
+          if (ok) Some((x, repl)) else None
+        case _ => None
+      }
+
     @inline
     def claimedRule(implicit ft: FormatToken): Option[Replacement] =
       claimedRule(ft.meta.idx)
@@ -317,8 +326,7 @@ object FormatTokensRewrite {
         else
           (repl.how match {
             case rt: ReplacementType.RemoveAndResurrect =>
-              claimed.get(rt.idx).flatMap { oldidx =>
-                val orepl = tokens(oldidx)
+              getClaimed(rt.idx).flatMap { case (oldidx, orepl) =>
                 val ok = orepl != null && (orepl.rule eq repl.rule) &&
                   (orepl.how eq ReplacementType.Remove)
                 if (ok) {