TreeOps: get owners via token range, not position

This allows us to handle tokens which occupy no space, such as Xml.End
or Xml.Start.

scalafmt-core/shared/src/main/scala/org/scalafmt/util/TreeOps.scala

@@ -860,15 +860,13 @@ object TreeOps {
     val allTokens = topSourceTree.tokens
     var prevParens: List[Token] = Nil
 
-    def treeAt(elemIdx: Int, elem: Tree, outerPrevLPs: Int): Int = {
+    def treeAt(elemIdx: Int, elem: Tree, elemEnd: Token, outerPrevLPs: Int): Int = {
       if (TreeOps.isInfixApp(elem)) infixCount += 1
 
-      val endPos = elem.pos.end
-      val allChildren: List[(Tree, Int)] = elem.children.flatMap { x =>
-        val pos = x.pos
-        val startPos = pos.start
-        if (startPos == pos.end) None else Some((x, startPos))
-      }.sortBy(_._2)
+      val allChildren: List[(Tree, Token, Token)] = elem.children.flatMap { x =>
+        val tokens = x.tokens
+        if (tokens.isEmpty) None else Some((x, tokens.head, tokens.last))
+      }.sortBy(_._2.start)
 
       allChildren match {
         case Nil =>
@@ -877,41 +875,42 @@ object TreeOps {
             if (idx == allTokens.length) idx
             else {
               val tok = allTokens(idx)
-              if (elem != topSourceTree && tok.start >= endPos) idx
-              else {
-                setOwner(tok, elem)
-                tokenAt(idx + 1)
-              }
+              setOwner(tok, elem)
+              val nextIdx = idx + 1
+              if (tok eq elemEnd) nextIdx else tokenAt(nextIdx)
             }
           tokenAt(elemIdx)
 
-        case (firstChild, firstChildStart) :: rest =>
+        case (firstChild, firstChildBeg, firstChildEnd) :: rest =>
           var nextChild = firstChild
-          var nextChildStart = firstChildStart
+          var nextChildBeg = firstChildBeg
+          var nextChildEnd = firstChildEnd
           var children = rest
           var prevChild: Tree = null
           var prevLPs = outerPrevLPs
           var prevComma: Token = null
 
           @tailrec
           def tokenAt(idx: Int): Int =
-            if (idx == allTokens.length) idx
+            if (
+              idx == allTokens.length ||
+              elem != topSourceTree && idx > 0 && (elemEnd eq allTokens(idx - 1))
+            ) idx
             else {
               val tok = allTokens(idx)
-              val tokStart = tok.start
-              if (elem != topSourceTree && tokStart >= endPos) idx
-              else if (nextChild != null && tokStart >= nextChildStart) {
+              if (nextChild != null && (tok eq nextChildBeg)) {
                 if (prevChild != null) prevLPs = 0
                 prevChild = nextChild
-                val nextIdx = treeAt(idx, nextChild, prevLPs)
+                val nextIdx = treeAt(idx, nextChild, nextChildEnd, prevLPs)
                 children match {
                   case Nil =>
                     nextChild = null
-                    nextChildStart = endPos
-                  case (head, start) :: rest =>
+                    nextChildEnd = elemEnd
+                  case (head, beg, end) :: rest =>
                     nextChild = head
                     children = rest
-                    nextChildStart = start
+                    nextChildBeg = beg
+                    nextChildEnd = end
                 }
                 prevComma = null
                 tokenAt(nextIdx)
@@ -921,8 +920,7 @@ object TreeOps {
                   case _: Term.While | _: Term.ForClause => prevLPs == 1 &&
                     prevChild == firstChild // `expr` is first
                   case _: Member.SyntaxValuesClause | _: Member.Tuple |
-                      _: Term.Do | _: Term.AnonymousFunction => endPos ==
-                      tok.end
+                      _: Term.Do | _: Term.AnonymousFunction => elemEnd eq tok
                   case t: Init => prevChild ne t.tpe // include tpe
                   case _: Ctor.Primary | _: Term.EnumeratorsBlock => true
                   case _ => false
@@ -944,7 +942,7 @@ object TreeOps {
                   setOwner(tok, elem)
                 } else {
                   setOwner(tok, elem)
-                  if (!tok.is[Trivia] && tokStart != tok.end) {
+                  if (!tok.is[Trivia] && !tok.isEmpty) {
                     prevComma = null
                     prevChild = null
                     if (tok.is[LeftParen]) {
@@ -960,7 +958,7 @@ object TreeOps {
       }
     }
 
-    treeAt(0, topSourceTree, 0)
+    treeAt(0, topSourceTree, allTokens.last, 0)
 
     val checkedNewlines = baseStyle.newlines.checkInfixConfig(infixCount)
     val initStyle =

scalafmt-tests/shared/src/test/resources/newlines/source_unfold.stat

@@ -9726,25 +9726,16 @@ object a {
   }
 }
 >>>
-Idempotency violated
-=> Diff (- obtained, + expected)
-   private def waitingBatchRow(batch: BatchUIData): Seq[Node] = {
--    baseRow(batch) ++ createOutputOperationProgressBar(batch) ++ <td>queued</td>
--    ++ {
--      if (firstFailureReason.nonEmpty) {
--        // Waiting batches have not run yet, so must have no failure reasons.
--        <td>-</td>
--      } else {
--        Nil
-+    baseRow(batch) ++ createOutputOperationProgressBar(batch) ++
-+      <td>queued</td>
-+      ++ {
-+        if (firstFailureReason.nonEmpty) {
-+          // Waiting batches have not run yet, so must have no failure reasons.
-+          <td>-</td>
-+        } else {
-+          Nil
-+        }
-       }
--    }
-   }
+object a {
+  private def waitingBatchRow(batch: BatchUIData): Seq[Node] = {
+    baseRow(batch) ++ createOutputOperationProgressBar(batch) ++
+      <td>queued</td> ++ {
+        if (firstFailureReason.nonEmpty) {
+          // Waiting batches have not run yet, so must have no failure reasons.
+          <td>-</td>
+        } else {
+          Nil
+        }
+      }
+  }
+}