FormatTokens: use getHead/Last, not tree tokens

Also, rename or refactor some methods.

scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala

@@ -308,9 +308,9 @@ class FormatOps(
     def alt = getLastToken(t)
     t match {
       case t: Defn.Def =>
-        t.body.tokens.headOption.fold(alt) {
-          case head: T.LeftBrace if t.body.is[Term.Block] => head
-          case head => tokenBefore(head).left
+        tokens.getHeadOpt(t.body).fold(alt) { ft =>
+          if (ft.left.is[T.LeftBrace] && t.body.is[Term.Block]) ft.left
+          else prevNonCommentBefore(ft).left
         }
       case _: Ctor.Primary =>
         close match {
@@ -839,7 +839,7 @@ class FormatOps(
     if (opFollowsComment) getLastNonTrivialToken(app.lhs) else opToken
   }
 
-  private def getLastEnclosedToken(tree: Tree): T = {
+  def getLastEnclosedToken(tree: Tree): T = {
     tokens.getLastExceptParen(tree.tokens).left
   }
 
@@ -1796,8 +1796,8 @@ class FormatOps(
           val callPolicy = CallSite.getFoldedPolicy(lia.lhs)
           if (callPolicy.nonEmpty) getPolicySplits(0, callPolicy)
           else {
-            val lp = body.tokens.headOption.filter(_.is[T.LeftParen])
-            val ok = lp.flatMap(matchingOpt).exists(_.end >= lia.op.pos.end)
+            // lia is enclosed in parens if and only if lia == ia (== body)
+            val ok = isEnclosedInParens(body)
             getSplits(getSlbSplit(getLastToken(if (ok) lia.lhs else lia.op)))
           }
         case _ =>
@@ -2674,7 +2674,7 @@ class FormatOps(
       def getBlocks(ft: FormatToken, nft: FormatToken, all: Boolean): Result =
         ft.meta.leftOwner match {
           case t @ Term.For(s, b)
-              if !nft.right.is[T.KwDo] && !isLastToken(ft.left, t) =>
+              if !nft.right.is[T.KwDo] && !isTokenLastOrAfter(ft.left, t) =>
             Some((b, seq(all, s)))
           case _ => None
         }

scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala

@@ -7,10 +7,7 @@ import scala.meta.tokens.Tokens
 
 import org.scalafmt.config.ScalafmtConfig
 import org.scalafmt.rewrite.FormatTokensRewrite
-import org.scalafmt.util.StyleMap
-import org.scalafmt.util.TokenOps
-import org.scalafmt.util.TreeOps
-import org.scalafmt.util.Whitespace
+import org.scalafmt.util._
 
 class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
     val arr: Array[FormatToken]
@@ -99,10 +96,8 @@ class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
   @inline private def areMatchingParens(close: Token)(open: => Token): Boolean =
     close.is[Token.RightParen] && areMatching(close)(open)
 
-  def isEnclosedInParens(tokens: Tokens): Boolean =
-    getClosingIfInParens(tokens).isDefined
   def isEnclosedInParens(tree: Tree): Boolean =
-    isEnclosedInParens(tree.tokens)
+    getClosingIfInParens(tree).isDefined
 
   def getClosingIfInParens(
       last: FormatToken
@@ -114,8 +109,10 @@ class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
         Some(afterLast)
       else None
     }
-  def getClosingIfInParens(tokens: Tokens): Option[FormatToken] =
+  def getClosingIfInParens(tree: Tree): Option[FormatToken] = {
+    val tokens = tree.tokens
     getHeadOpt(tokens).flatMap(getClosingIfInParens(getLastNonTrivial(tokens)))
+  }
 
   def getLastExceptParen(tokens: Tokens): FormatToken = {
     val last = getLastNonTrivial(tokens)
@@ -166,6 +163,7 @@ class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
     after(tokens.head)
   def getHead(tree: Tree): FormatToken =
     getHead(tree.tokens)
+
   def getHeadOpt(tokens: Tokens): Option[FormatToken] =
     tokens.headOption.map(after)
   def getHeadOpt(tree: Tree): Option[FormatToken] =
@@ -175,6 +173,7 @@ class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
     apply(TokenOps.findLastVisibleToken(tokens))
   def getLast(tree: Tree): FormatToken =
     getLast(tree.tokens)
+
   def getLastOpt(tokens: Tokens): Option[FormatToken] =
     TokenOps.findLastVisibleTokenOpt(tokens).map(apply)
   def getLastOpt(tree: Tree): Option[FormatToken] =
@@ -184,6 +183,7 @@ class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
     apply(TokenOps.findLastNonTrivialToken(tokens))
   def getLastNonTrivial(tree: Tree): FormatToken =
     getLastNonTrivial(tree.tokens)
+
   def getLastNonTrivialOpt(tokens: Tokens): Option[FormatToken] =
     TokenOps.findLastNonTrivialTokenOpt(tokens).map(apply)
   def getLastNonTrivialOpt(tree: Tree): Option[FormatToken] =
@@ -208,10 +208,10 @@ class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
   @inline
   def justBefore(token: Token): FormatToken = apply(token, -1)
   @inline
-  def tokenJustBefore(tree: Tree): FormatToken = justBefore(tree.tokens.head)
+  def tokenJustBefore(tree: Tree): FormatToken = prev(getHead(tree))
 
   def tokenJustBeforeOpt(tree: Tree): Option[FormatToken] =
-    tree.tokens.headOption.map(justBefore)
+    getHeadOpt(tree).map(prev)
   def tokenJustBeforeOpt(trees: Seq[Tree]): Option[FormatToken] =
     trees.headOption.flatMap(tokenJustBeforeOpt)
 
@@ -220,7 +220,8 @@ class FormatTokens(leftTok2tok: Map[TokenOps.TokenHash, Int])(
   @inline
   def tokenBefore(token: Token): FormatToken = prevNonComment(justBefore(token))
   @inline
-  def tokenBefore(tree: Tree): FormatToken = tokenBefore(tree.tokens.head)
+  def tokenBefore(tree: Tree): FormatToken =
+    prevNonComment(tokenJustBefore(tree))
   @inline
   def tokenBefore(trees: Seq[Tree]): FormatToken = tokenBefore(trees.head)
 

scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala

@@ -1488,13 +1488,13 @@ class Router(formatOps: FormatOps) {
                 case Some(p: Term.ApplyInfix) => startsInfix(p, ai)
                 case _ => true
               })
-            val roTokens = rightOwner.tokens
-            isFirstToken(lb, roTokens) && rightOwner.parent.exists {
+            val roPos = rightOwner.pos
+            isTokenHeadOrBefore(lb, roPos) && rightOwner.parent.exists {
               case p: Term.ApplyInfix => // exclude start of infix
                 startsInfix(p, rightOwner)
               case p =>
-                isFirstToken(lb, p) &&
-                matchingOpt(lb).exists(isLastToken(_, roTokens))
+                isTokenHeadOrBefore(lb, p) &&
+                matchingOpt(lb).exists(isTokenLastOrAfter(_, roPos))
             }
           } =>
         Seq(Split(Space, 0))
@@ -1671,7 +1671,7 @@ class Router(formatOps: FormatOps) {
           case x: Term.Match => SelectLike(x, getKwMatchAfterDot(t))
         }
         val prevSelect = findPrevSelect(thisSelect, enclosed)
-        val expire = tokens.getLastExceptParen(expireTree.tokens).left
+        val expire = getLastEnclosedToken(expireTree)
         val indentLen = style.indent.main
 
         def breakOnNextDot: Policy =
@@ -1939,7 +1939,7 @@ class Router(formatOps: FormatOps) {
       // If/For/While/For with (
       case FormatToken(open: T.LeftParen, _, _) if (leftOwner match {
             case _: Term.If | _: Term.While | _: Term.For | _: Term.ForYield =>
-              !isFirstToken(open, leftOwner)
+              !isTokenHeadOrBefore(open, leftOwner)
             case _ => false
           }) =>
         val close = matching(open)
@@ -2009,7 +2009,7 @@ class Router(formatOps: FormatOps) {
             case _: Term.While | _: Term.For =>
               !nextNonComment(formatToken).right.is[T.KwDo]
             case _ => false
-          }) && !isLastToken(close, leftOwner) =>
+          }) && !isTokenLastOrAfter(close, leftOwner) =>
         val body = leftOwner match {
           case t: Term.If => t.thenp
           case t: Term.For => t.body
@@ -2306,7 +2306,7 @@ class Router(formatOps: FormatOps) {
 
       // Term.For
       case ft @ FormatToken(rb: T.RightBrace, _, _)
-          if leftOwner.is[Term.For] && !isLastToken(rb, leftOwner) &&
+          if leftOwner.is[Term.For] && !isTokenLastOrAfter(rb, leftOwner) &&
             !nextNonComment(formatToken).right.is[T.KwDo] =>
         val body = leftOwner.asInstanceOf[Term.For].body
         def nlSplit(cost: Int) = Split(Newline, cost)
@@ -2516,7 +2516,7 @@ class Router(formatOps: FormatOps) {
         def modNoNL = {
           def allowSpace = rightOwner match {
             case _: Term.If | _: Term.While | _: Term.For | _: Term.ForYield =>
-              isLastToken(close, rightOwner)
+              isTokenLastOrAfter(close, rightOwner)
             case _ => true
           }
           Space(style.spaces.inParentheses && allowSpace)

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

@@ -87,10 +87,9 @@ object TreeOps {
     val ret = Map.newBuilder[TokenHash, Tree]
     ret.sizeHint(tree.tokens.length)
 
-    def addTok(token: Token, tree: Tree) =
-      ret += hash(ftoks.after(token).left) -> tree
-    def addTree(t: Tree, tree: Tree) =
-      t.tokens.find(!_.is[Trivia]).foreach(addTok(_, tree))
+    def addFT(ft: FormatToken, tree: Tree): Unit = ret += hash(ft.left) -> tree
+    def addTok(token: Token, tree: Tree) = addFT(ftoks.after(token), tree)
+    def addTree(t: Tree, o: Tree) = ftoks.getHeadOpt(t).foreach(addFT(_, o))
     def addAll(trees: Seq[Tree]) = trees.foreach(x => addTree(x, x))
 
     def addDefnTokens(
@@ -411,23 +410,23 @@ object TreeOps {
       isSuperfluousParenthesis(open.asInstanceOf[LeftParen], owner)
 
   def isSuperfluousParenthesis(open: LeftParen, owner: Tree): Boolean =
-    !isTuple(owner) && isFirstToken(open, owner)
+    !isTuple(owner) && isTokenHeadOrBefore(open, owner)
 
   @inline
-  def isFirstToken(token: Token, owner: Tree): Boolean =
-    isFirstToken(token, owner.tokens)
+  def isTokenHeadOrBefore(token: Token, owner: Tree): Boolean =
+    isTokenHeadOrBefore(token, owner.pos)
 
   @inline
-  def isFirstToken(token: Token, owner: Tokens): Boolean =
-    owner.headOption.contains(token)
+  def isTokenHeadOrBefore(token: Token, pos: Position): Boolean =
+    pos.start >= token.start
 
   @inline
-  def isLastToken(token: Token, owner: Tree): Boolean =
-    isLastToken(token, owner.tokens)
+  def isTokenLastOrAfter(token: Token, owner: Tree): Boolean =
+    isTokenLastOrAfter(token, owner.pos)
 
   @inline
-  def isLastToken(token: Token, owner: Tokens): Boolean =
-    owner.lastOption.contains(token)
+  def isTokenLastOrAfter(token: Token, pos: Position): Boolean =
+    pos.end <= token.end
 
   def isCallSite(tree: Tree)(implicit style: ScalafmtConfig): Boolean =
     tree match {
@@ -564,9 +563,7 @@ object TreeOps {
   final def lastLambda(first: Term.FunctionTerm): Term.FunctionTerm =
     first.body match {
       case child: Term.FunctionTerm => lastLambda(child)
-      case block: Term.Block
-          if block.stats.headOption.exists(_.is[Term.FunctionTerm]) =>
-        lastLambda(block.stats.head.asInstanceOf[Term.FunctionTerm])
+      case Term.Block((child: Term.FunctionTerm) :: _) => lastLambda(child)
       case _ => first
     }