TreeOps: get last tree in followedBySelectOrApply

scalameta · May 13, 2024 · 692defc · 692defc
1 parent 302288c
commit 692defc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
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
@@ -1184,7 +1184,7 @@ class Router(formatOps: FormatOps) {
           else firstArg.map(getLast).flatMap(findComma)
         val needOnelinePolicy = oneline &&
           (nextCommaOneline.isDefined ||
-            leftOwner.parent.exists(followedBySelectOrApply))
+            leftOwner.parent.exists(followedBySelectOrApply(_).isDefined))
         val nextCommaOnelinePolicy =
           if (needOnelinePolicy) nextCommaOneline
             .map(splitOneArgPerLineAfterCommaOnBreak)
@@ -1445,7 +1445,8 @@ class Router(formatOps: FormatOps) {
                 else delayedBreakPolicyFor(t)(decideNewlinesOnlyAfterClose)
               case Some(FormatToken(_, t, _))
                   if !callSite ||
-                    leftOwner.parent.exists(followedBySelectOrApply) =>
+                    leftOwner.parent
+                      .exists(followedBySelectOrApply(_).isDefined) =>
                 decideNewlinesOnlyBeforeCloseOnBreak(t)
               case _ => NoPolicy
             }

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
@@ -786,16 +786,17 @@ object TreeOps {
   @tailrec
   final def followedBySelectOrApply(
       tree: Tree,
-  )(implicit ftoks: FormatTokens): Boolean = tree.parent match {
+  )(implicit ftoks: FormatTokens): Option[Tree] = tree.parent match {
     case Some(p: Term.New) => followedBySelectOrApply(p)
-    case Some(_: Term.Select) => true
-    case Some(p: Member.Infix) => p.lhs.eq(tree) || followedBySelectOrApply(p)
+    case Some(p: Term.Select) if p.qual eq tree => Some(tree)
+    case Some(p: Member.Infix) =>
+      if (p.lhs eq tree) Some(tree) else followedBySelectOrApply(p)
     case Some(p: Member.Apply) if p.fun eq tree =>
       p.argClause match {
-        case SingleArgInBraces(_) => false
-        case _ => true
+        case SingleArgInBraces(_) => None
+        case _ => Some(tree)
       }
-    case _ => false
+    case _ => None
   }
 
   // Scala syntax allows commas before right braces in weird places,