TreeOps: extract both braces in SingleArgInBraces

scalameta · Mar 16, 2024 · cc8af48 · cc8af48
1 parent 20280c4
commit cc8af48
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 14 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
@@ -627,7 +627,7 @@ class FormatOps(
         case _: Term.Block | _: Term.If | _: Term.While | _: Source => true
         case fun: Term.FunctionTerm => isBlockFunction(fun)
         case t: Case => t.pat.eq(child) || t.body.eq(child)
-        case SingleArgInBraces(arg) => child eq arg
+        case SingleArgInBraces(_, arg, _) => child eq arg
         case _ => false
       }
       def isAloneEnclosed(child: Tree) = child.parent.exists {
@@ -936,8 +936,8 @@ class FormatOps(
 
     def getRToks = dropWS(function.tokens.reverse)
     function.parent match {
-      case Some(p @ SingleArgInBraces.OrBlock(_)) =>
-        tokens.getLast(p).left -> ExpiresOn.Before
+      case Some(SingleArgInBraces.OrBlock(_, _, e)) =>
+        e.left -> ExpiresOn.Before
       case Some(Case(_, _, `function`)) =>
         orElse(dropComment(getRToks))
       case _ =>

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
@@ -165,7 +165,7 @@ class FormatWriter(formatOps: FormatOps) {
               b.parent.exists(tokens.getLast(_) eq tok)
             case f: Term.FunctionTerm =>
               checkApply(f) && RedundantBraces.canRewriteFuncWithParens(f)
-            case t @ TreeOps.SingleArgInBraces(arg) =>
+            case t @ TreeOps.SingleArgInBraces(_, arg, _) =>
               TreeOps.isParentAnApply(t) &&
               RedundantBraces.canRewriteStatWithParens(arg)
             case _ => false

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
@@ -171,7 +171,7 @@ class RedundantBraces(implicit val ftoks: FormatTokens)
         else if (okToReplaceFunctionInSingleArgApply(t)) replaceWithLeftParen
         else removeToken
       case t: Term.PartialFunction if t.parent.exists { p =>
-            SingleArgInBraces.orBlock(p).contains(t) &&
+            SingleArgInBraces.orBlock(p).exists(_._2 eq t) &&
             t.pos.start != p.pos.start
           } =>
         removeToken

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/util/TreeOps.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/util/TreeOps.scala
@@ -52,33 +52,38 @@ object TreeOps {
   object SingleArgInBraces {
     def unapply(tree: Tree)(implicit
         ftoks: FormatTokens
-    ): Option[Term] = tree match {
+    ): Option[(FormatToken, Term, FormatToken)] = tree match {
       case t: Term.ArgClause => unapply(t)
       case _ => None
     }
     def unapply(tree: Term.ArgClause)(implicit
         ftoks: FormatTokens
-    ): Option[Term] = tree.values match {
-      case arg :: Nil if inBraces(tree) => Some(arg)
+    ): Option[(FormatToken, Term, FormatToken)] = tree.values match {
+      case arg :: Nil => getBraces(tree).map { case (b, e) => (b, arg, e) }
       case _ => None
     }
     @inline def inBraces(tree: Tree)(implicit
         ftoks: FormatTokens
     ): Boolean =
-      ftoks.getHeadIfEnclosed(tree).exists(_.left.is[LeftBrace])
+      getBraces(tree).isDefined
+    @inline def getBraces(tree: Tree)(implicit
+        ftoks: FormatTokens
+    ): Option[(FormatToken, FormatToken)] =
+      ftoks.getDelimsIfEnclosed(tree).filter(_._1.left.is[LeftBrace])
 
     def orBlock(tree: Tree)(implicit
         ftoks: FormatTokens
-    ): Option[Tree] = tree match {
+    ): Option[(FormatToken, Stat, FormatToken)] = tree match {
       case t: Term.ArgClause => unapply(t)
-      case Term.Block(arg :: Nil) if inBraces(tree) => Some(arg)
+      case Term.Block(arg :: Nil) =>
+        getBraces(tree).map { case (b, e) => (b, arg, e) }
       case _ => None
     }
 
     object OrBlock {
       def unapply(tree: Tree)(implicit
           ftoks: FormatTokens
-      ): Option[Tree] = orBlock(tree)
+      ): Option[(FormatToken, Stat, FormatToken)] = orBlock(tree)
     }
   }
 
@@ -100,14 +105,14 @@ object TreeOps {
   def isExprWithParentInBraces(expr: Tree)(parent: Tree)(implicit
       ftoks: FormatTokens
   ): Boolean =
-    SingleArgInBraces.orBlock(parent).contains(expr)
+    SingleArgInBraces.orBlock(parent).exists(_._2 eq expr)
 
   def extractStatementsIfAny(tree: Tree)(implicit
       ftoks: FormatTokens
   ): Seq[Tree] =
     tree match {
       case b: Term.Block => b.stats
-      case SingleArgInBraces(fun: Term.FunctionTerm) => fun :: Nil
+      case SingleArgInBraces(_, fun: Term.FunctionTerm, _) => fun :: Nil
       case b: Term.FunctionTerm if isBlockFunction(b) => b.body :: Nil
       case t: Pkg => t.stats
       // TODO(olafur) would be nice to have an abstract "For" superclass.