TreeOps: define SingleArgInBraces tree matcher

Use to generalize a few cases, to avoid making assumptions about the underlying argument or just make them future-proof.
scalameta · Feb 18, 2024 · 178e15e · 178e15e
1 parent ee5ebdc
commit 178e15e
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 9 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
@@ -626,6 +626,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 _ => false
       }
       def isAloneEnclosed(child: Tree) = child.parent.exists {
@@ -944,7 +945,7 @@ class FormatOps(
 
     def getRToks = dropWS(function.tokens.reverse)
     function.parent match {
-      case Some(p @ Term.Block(_ :: Nil)) =>
+      case Some(p @ SingleArgInBraces.OrBlock(_)) =>
         p.tokens.last -> ExpiresOn.Before
       case Some(Case(_, _, `function`)) =>
         orElse(dropComment(getRToks))
@@ -2677,7 +2678,7 @@ class FormatOps(
       def getBlocks(ft: FormatToken, nft: FormatToken, all: Boolean): Result =
         ft.meta.leftOwner match {
           case t: Term.FunctionTerm =>
-            val skip = t.parent.exists(_.is[Term.Block])
+            val skip = t.parent.exists(TreeOps.isExprWithParentInBraces(t))
             if (skip) None else Some((t.body, seq(all, t.paramClause.values)))
           case _ => None
         }

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
@@ -48,6 +48,29 @@ object TreeOps {
     ret.result()
   }
 
+  object SingleArgInBraces {
+    def unapply(tree: Tree): Option[Term] = tree match {
+      case t: Term.ArgClause => unapply(t)
+      case _ => None
+    }
+    def unapply(tree: Term.ArgClause): Option[Term] = tree.values match {
+      case arg :: Nil if inBraces(tree) => Some(arg)
+      case _ => None
+    }
+    @inline def inBraces(tree: Term.ArgClause): Boolean =
+      tree.tokens.head.is[LeftBrace]
+
+    def orBlock(tree: Tree): Option[Tree] = tree match {
+      case t: Term.ArgClause => unapply(t)
+      case Term.Block(arg :: Nil) => Some(arg)
+      case _ => None
+    }
+
+    object OrBlock {
+      def unapply(tree: Tree): Option[Tree] = orBlock(tree)
+    }
+  }
+
   @tailrec
   def isBlockFunction(fun: Term.FunctionTerm): Boolean =
     fun.parent match {
@@ -59,15 +82,13 @@ object TreeOps {
   def isFunctionWithBraces(fun: Term.FunctionTerm): Boolean =
     fun.parent.exists(isExprWithParentInBraces(fun))
 
-  private def isExprWithParentInBraces(expr: Tree)(parent: Tree): Boolean =
-    parent match {
-      case Term.Block(stat :: Nil) => stat eq expr
-      case _ => false
-    }
+  def isExprWithParentInBraces(expr: Tree)(parent: Tree): Boolean =
+    SingleArgInBraces.orBlock(parent).contains(expr)
 
   def extractStatementsIfAny(tree: Tree): Seq[Tree] =
     tree match {
       case b: Term.Block => b.stats
+      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.
@@ -828,8 +849,8 @@ object TreeOps {
     case Some(_: Term.Select) => true
     case Some(p: Member.Infix) => p.lhs.eq(tree) || followedBySelectOrApply(p)
     case Some(p: Member.Apply) if p.fun eq tree =>
-      p.argClause.values match {
-        case (_: Term.Block | _: Term.PartialFunction) :: Nil => false
+      p.argClause match {
+        case SingleArgInBraces(_) => false
         case _ => true
       }
     case _ => false