TreeOps: fix ownership of enclosed trees

- exclude only last right paren (control expressions like if/while/for) - ignore empty tokens (Xml.End, Interpolation.End) - allow procedure-syntax Decl.Def (no type) - detect top-level invocation directly, not via elemIdx
scalameta · Dec 11, 2022 · 74ba074 · 74ba074
1 parent 43adb34
commit 74ba074
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
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
@@ -971,7 +971,7 @@ object TreeOps {
             if (idx == allTokens.length) idx
             else {
               val tok = allTokens(idx)
-              if (tok.start >= endPos && elemIdx != 0) idx
+              if (elem != topSourceTree && tok.start >= endPos) idx
               else {
                 setOwner(tok, elem)
                 tokenAt(idx + 1)
@@ -991,14 +991,16 @@ object TreeOps {
             if (idx == allTokens.length) idx
             else {
               val tok = allTokens(idx)
-              if (nextChild == null && elemIdx != 0 && tok.start >= endPos) idx
-              else if (nextChild != null && tok.end > nextChildStart) {
+              val tokStart = tok.start
+              if (elem != topSourceTree && tokStart >= endPos) idx
+              else if (nextChild != null && tokStart >= nextChildStart) {
                 if (prevChild != null) prevLPs = 0
                 prevChild = nextChild
                 val nextIdx = treeAt(idx, nextChild, prevLPs)
                 children match {
                   case Nil =>
                     nextChild = null
+                    nextChildStart = endPos
                   case (head, start) :: rest =>
                     nextChild = head
                     children = rest
@@ -1007,33 +1009,38 @@ object TreeOps {
                 tokenAt(nextIdx)
               } else {
                 def excludeRightParen: Boolean = elem match {
-                  case t: Term.If => prevChild eq t.cond // `expr` after `mods`
+                  case t: Term.If =>
+                    prevLPs == 1 && prevChild == t.cond // `expr` after `mods`
                   case _: Term.While | _: Term.For | _: Term.ForYield =>
-                    prevChild eq firstChild // `expr` is first
+                    prevLPs == 1 && prevChild == firstChild // `expr` is first
                   case _: Term.Tuple | _: Type.Tuple | _: Pat.Tuple |
                       _: Term.ApplyInfix | _: Pat.ExtractInfix | _: Term.Apply |
                       _: Pat.Extract | _: Term.Do | _: Term.AnonymousFunction =>
                     nextChild == null
-                  case _: Decl.Def | _: Decl.Given | _: Type.FunctionType |
-                      _: Defn.Def | _: Defn.Macro | _: Defn.Given |
-                      _: Defn.GivenAlias | _: Defn.ExtensionGroup =>
-                    nextChild != null // body is next
-                  case _: Init | _: Ctor => true
+                  case t: Decl.Def => // only body can have owners modified
+                    nextChild != null || t.decltpe != prevChild
+                  case t: Decl.Given => // only body can have owners modified
+                    nextChild != null || t.decltpe != prevChild
+                  case _: Type.FunctionType | _: Defn.Def | _: Defn.GivenAlias |
+                      _: Defn.Given | _: Defn.Macro | _: Defn.ExtensionGroup =>
+                    nextChild != null // only body can have owners modified
+                  case _: Init => prevChild ne firstChild // include tpe
+                  case _: Ctor => true
                   case _ => false
                 }
 
                 if (prevParens.nonEmpty && tok.is[RightParen]) {
-                  if (prevChild == null || prevLPs == 0 || excludeRightParen)
+                  if (prevChild == null || prevLPs <= 0 || excludeRightParen)
                     setOwner(tok, elem)
                   else {
-                    prevLPs -= 1
                     setOwner(tok, prevChild)
                     setOwner(prevParens.head, prevChild)
                   }
+                  prevLPs -= 1
                   prevParens = prevParens.tail
                 } else {
                   setOwner(tok, elem)
-                  if (!tok.is[Trivia] && nextChild != null) {
+                  if (!tok.is[Trivia] && tokStart != tok.end) {
                     prevChild = null
                     if (tok.is[LeftParen]) {
                       prevLPs += 1

diff --git a/scalafmt-tests/src/test/resources/newlines/source_classic.stat b/scalafmt-tests/src/test/resources/newlines/source_classic.stat
@@ -6484,7 +6484,7 @@ object a {
 }
 >>>
 object a {
-  info (
+  info(
     "ZKConsumerConnector shutdown completed in " +
       (System
         .nanoTime() - startTime) / 1000000 + " ms"
-Original file line number
+Diff line change
@@ Expand Up / @@ -6484,7 +6484,7 @@ object a { @@
     }
     >>>
     object a {
-      info (
+      info(
         "ZKConsumerConnector shutdown completed in " +
           (System
             .nanoTime() - startTime) / 1000000 + " ms"
@@ Expand Down @@