FormatOps: get rewritten tree head and last safer

scalameta · Dec 30, 2021 · bd78d8b · bd78d8b
1 parent 1044994
commit bd78d8b
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 39 deletions.
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
@@ -77,6 +77,7 @@ class FormatOps(
   import tokens.{
     matching,
     matchingOpt,
+    isEnclosedInMatching,
     findTokenWith,
     tokenBefore,
     tokenAfter,
@@ -662,7 +663,7 @@ class FormatOps(
               infixes.takeWhile { x =>
                 val close = getLastNonTrivialToken(x.lhs)
                 !close.is[T.RightParen] ||
-                !isEnclosedInMatching(owners(close), matching(close), close)
+                !tokens.isCloseMatchingHead(close)(owners(close))
               }
           if (filtered.isEmpty) None
           else {
@@ -1425,15 +1426,6 @@ class FormatOps(
       case _ => true
     } || ft.right.is[soft.KwUsing]
 
-  def isEnclosedInMatching(tree: Tree, open: T, close: T): Boolean =
-    tree.tokens.headOption.contains(open) && (tree.tokens.last eq close)
-
-  def getClosingIfEnclosedInMatching(tree: Tree): Option[T] =
-    tree.tokens.lastOption.filter(tokens.areMatching(tree.tokens.head))
-
-  def isEnclosedInMatching(tree: Tree): Boolean =
-    getClosingIfEnclosedInMatching(tree).isDefined
-
   @tailrec
   final def findPrevSelect(
       tree: Tree,
@@ -2011,8 +2003,11 @@ class FormatOps(
       val slbExpire = nextNonCommentSameLine(end).left
       val closeOpt =
         if (isTuple(tree)) None
-        else
-          getClosingIfEnclosedInMatching(tree).map(tokens.tokenBefore(_).left)
+        else {
+          val maybeClose = prevNonComment(end)
+          if (!tokens.isCloseMatchingHead(maybeClose.left)(tree)) None
+          else Some(prevNonComment(prev(maybeClose)).left)
+        }
       def nlPolicy(implicit fileLine: FileLine) =
         if (danglingKeyword)
           decideNewlinesOnlyAfterClose(closeOpt.getOrElse(slbExpire))

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
@@ -89,6 +89,20 @@ class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
   @inline def areMatching(t1: Token)(t2: Token): Boolean =
     matchingOpt(t1).contains(t2)
 
+  def getClosingIfEnclosedInMatching(tree: Tree): Option[FormatToken] =
+    getHeadOpt(tree).flatMap { head =>
+      matchingOpt(head.left).flatMap { last =>
+        val ft = getLastNonTrivial(tree)
+        if (ft.left eq last) Some(ft) else None
+      }
+    }
+
+  def isEnclosedInMatching(tree: Tree): Boolean =
+    getClosingIfEnclosedInMatching(tree).isDefined
+
+  def isCloseMatchingHead(close: Token)(tree: => Tree): Boolean =
+    matchingOpt(close).exists(_ eq getHead(tree).left)
+
   @tailrec
   final def findTokenWith[A](
       ft: FormatToken,

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatWriter.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatWriter.scala
@@ -357,7 +357,7 @@ class FormatWriter(formatOps: FormatOps) {
   private def checkInsertBraces(locations: Array[FormatLocation]): Unit = {
     def checkInfix(tree: Tree): Boolean = tree match {
       case ai @ Term.ApplyInfix(lhs, op, _, rhs) => {
-        isEnclosedInMatching(ai) ||
+        tokens.isEnclosedInMatching(ai) ||
         tokens.prevNonCommentSameLine(tokens.tokenJustBefore(op)).noBreak &&
         checkInfix(lhs) && (rhs.lengthCompare(1) != 0 || checkInfix(rhs.head))
       }

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala
@@ -66,6 +66,7 @@ class Router(formatOps: FormatOps) {
   import tokens.{
     matching,
     matchingOpt,
+    isEnclosedInMatching,
     prev,
     next,
     tokenBefore,

diff --git a/scalafmt-tests/src/test/resources/newlines/source_fold.stat b/scalafmt-tests/src/test/resources/newlines/source_fold.stat
@@ -429,20 +429,18 @@ val attributes = (
  }
 )
 >>>
-Idempotency violated
-val attributes =
-  for (_ ← 1 to count) yield {
-    val tag = d.readUnsignedShort()
-    val length = d.readInt()
-    if (debug) println(s"LNB:   tag ${c(tag)} ($length bytes)")
-    if (tag != s) {
-      skip(d, length)
-      None
-    } else {
-      val name = d.readUnsignedShort()
-      Some(c(name))
-    }
+val attributes = for (_ ← 1 to count) yield {
+  val tag = d.readUnsignedShort()
+  val length = d.readInt()
+  if (debug) println(s"LNB:   tag ${c(tag)} ($length bytes)")
+  if (tag != s) {
+    skip(d, length)
+    None
+  } else {
+    val name = d.readUnsignedShort()
+    Some(c(name))
   }
+}
 <<< 2.13.4 def with for/yield
 maxColumn = 80
 ===
@@ -525,20 +523,18 @@ def attributes = (
  }
 )
 >>>
-Idempotency violated
-def attributes =
-  for (_ ← 1 to count) yield {
-    val tag = d.readUnsignedShort()
-    val length = d.readInt()
-    if (debug) println(s"LNB:   tag ${c(tag)} ($length bytes)")
-    if (tag != s) {
-      skip(d, length)
-      None
-    } else {
-      val name = d.readUnsignedShort()
-      Some(c(name))
-    }
+def attributes = for (_ ← 1 to count) yield {
+  val tag = d.readUnsignedShort()
+  val length = d.readInt()
+  if (debug) println(s"LNB:   tag ${c(tag)} ($length bytes)")
+  if (tag != s) {
+    skip(d, length)
+    None
+  } else {
+    val name = d.readUnsignedShort()
+    Some(c(name))
   }
+}
 <<< 2.14 val with short for/yield
 maxColumn = 80
 ===