From a5b404082bdc21334f45048e8062e41220344cdc Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Wed, 13 Mar 2024 14:33:06 -0700
Subject: [PATCH] FormatTokensRewrite: "claim" index in `tokens`

---
 .../org/scalafmt/rewrite/FormatTokensRewrite.scala     | 10 ++++++----
 1 file changed, 6 insertions(+), 4 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
index 8d8830ae70..39c3dd768e 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/FormatTokensRewrite.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/FormatTokensRewrite.scala
@@ -273,7 +273,8 @@ object FormatTokensRewrite {
 
   private[rewrite] class Session(rules: Seq[Rule]) {
     private implicit val implicitSession: Session = this
-    private val claimed = new mutable.HashMap[Int, Replacement]()
+    // map FT index to index in tokens below
+    private val claimed = new mutable.HashMap[Int, Int]()
     private[FormatTokensRewrite] val tokens =
       new mutable.ArrayBuffer[Replacement]()
 
@@ -283,14 +284,15 @@ object FormatTokensRewrite {
 
     @inline
     private[rewrite] def claimedRule(ftIdx: Int): Option[Replacement] =
-      claimed.get(ftIdx)
+      claimed.get(ftIdx).map(tokens.apply)
 
     private[FormatTokensRewrite] def applyRule(
         attemptedRule: Rule
     )(implicit ft: FormatToken, style: ScalafmtConfig): Option[Rule] =
       if (attemptedRule.enabled) attemptedRule.onToken.map { repl =>
-        claimed.getOrElseUpdate(ft.meta.idx, repl)
-        repl.claim.foreach { claimed.getOrElseUpdate(_, repl) }
+        val idx = tokens.length
+        claimed.getOrElseUpdate(ft.meta.idx, idx)
+        repl.claim.foreach { claimed.getOrElseUpdate(_, idx) }
         tokens.append(repl)
         repl.rule
       }