FormatTokensRewrite: pass FormatToken to resurrect

scalameta · Mar 27, 2024 · 4ad80a0 · 4ad80a0
1 parent 937eef7
commit 4ad80a0
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 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
@@ -76,11 +76,12 @@ class FormatTokensRewrite(
         case ReplacementType.Remove => remove(appended)
         case ReplacementType.Replace => append()
         case r: ReplacementType.RemoveAndResurrect =>
-          if (r.idx == idx) { // we moved here
+          val rtidx = r.ft.meta.idx
+          if (rtidx == idx) { // we moved here
             append()
             shiftedIndexMap.put(idx, appended)
           } else { // we moved from here
-            remove(shiftedIndexMap.remove(r.idx).getOrElse(appended))
+            remove(shiftedIndexMap.remove(rtidx).getOrElse(appended))
           }
       }
     }
@@ -324,13 +325,15 @@ object FormatTokensRewrite {
         else
           (repl.how match {
             case rt: ReplacementType.RemoveAndResurrect =>
-              getClaimed(rt.idx).flatMap { case (oldidx, orepl) =>
-                val ok = orepl != null && (orepl.rule eq repl.rule) &&
-                  orepl.isRemove
-                if (ok) {
-                  tokens(oldidx) = repl.copy(ft = repl.ft.withIdx(orepl.idx))
-                  Some(repl.copy(ft = orepl.ft.withIdx(repl.idx)))
-                } else None
+              val rtidx = rt.ft.meta.idx
+              def swapWith(oldidx: Int) = Some {
+                tokens(oldidx) = repl.copy(ft = repl.ft.withIdx(rtidx))
+                repl.copy(ft = rt.ft.withIdx(repl.idx))
+              }
+              getClaimed(rtidx) match {
+                case Some((oidx, x)) if x != null && x.isRemove =>
+                  swapWith(oidx)
+                case _ => None
               }
             case _ => None
           }).getOrElse(repl)
@@ -422,7 +425,7 @@ object FormatTokensRewrite {
   private[rewrite] object ReplacementType {
     object Remove extends ReplacementType
     object Replace extends ReplacementType
-    class RemoveAndResurrect(val idx: Int) extends ReplacementType
+    class RemoveAndResurrect(val ft: FormatToken) extends ReplacementType
   }
 
   private def mergeWhitespaceLeftToRight(

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/RedundantBraces.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/rewrite/RedundantBraces.scala
@@ -153,7 +153,7 @@ class RedundantBraces(implicit val ftoks: FormatTokens)
       val rb = pft.left
       if (rb.is[Token.RightBrace]) {
         // move right to the end of the function
-        val rType = new ReplacementType.RemoveAndResurrect(pft.meta.idx - 1)
+        val rType = new ReplacementType.RemoveAndResurrect(ftoks.prev(pft))
         left -> replaceToken("}", rtype = rType) {
           // create a shifted token so that any child tree wouldn't own it
           new Token.RightBrace(rb.input, rb.dialect, rb.start + 1)