From 48c599f0c7c2508054acfce7c6217451e869784d Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Thu, 14 Mar 2024 10:26:56 -0700
Subject: [PATCH] FormatTokens: add nextNonCommentAfter

---
 .../main/scala/org/scalafmt/internal/FormatOps.scala   |  4 ++--
 .../scala/org/scalafmt/internal/FormatTokens.scala     | 10 +++++++++-
 .../src/main/scala/org/scalafmt/internal/Router.scala  |  8 +++++---
 3 files changed, 16 insertions(+), 6 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
index 50d9fc960e..1526db97d4 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatOps.scala
@@ -1430,7 +1430,7 @@ class FormatOps(
   private def getArrowAfter(ft: FormatToken): FormatToken = {
     val maybeArrow = next(ft)
     if (maybeArrow.left.is[T.RightArrow]) maybeArrow
-    else next(nextNonComment(maybeArrow))
+    else tokens.nextAfterNonComment(maybeArrow)
   }
 
   def getArgs(owner: Tree, orNil: Boolean = false): Seq[Tree] =
@@ -1914,7 +1914,7 @@ class FormatOps(
     } else None
 
   def getKwMatchAfterDot(ft: FormatToken): T.KwMatch =
-    nextNonComment(next(ft)).right.asInstanceOf[T.KwMatch]
+    tokens.nextNonCommentAfter(ft).right.asInstanceOf[T.KwMatch]
 
   object GetSelectLike {
     def unapply(tree: Tree): Option[SelectLike] = tree match {
diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala
index decdc7ba5d..db16a607d6 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/FormatTokens.scala
@@ -171,6 +171,14 @@ class FormatTokens(leftTok2tok: Map[TokenHash, Int])(
   final def nextNonComment(curr: FormatToken.Meta): FormatToken =
     nextNonComment(arr(curr.idx))
 
+  @inline
+  final def nextNonCommentAfter(curr: FormatToken): FormatToken =
+    nextNonComment(next(curr))
+
+  @inline
+  final def nextAfterNonComment(curr: FormatToken): FormatToken =
+    next(nextNonComment(curr))
+
   final def prevNonCommentSameLine(curr: FormatToken): FormatToken =
     findToken(curr, prev)(ft => ft.hasBreak || !ft.left.is[Token.Comment])
 
@@ -194,7 +202,7 @@ class FormatTokens(leftTok2tok: Map[TokenHash, Int])(
 
   @tailrec
   final def getOnOrAfterOwned(ft: FormatToken, tree: Tree): FormatToken = {
-    val nextFt = next(nextNonComment(ft))
+    val nextFt = nextAfterNonComment(ft)
     if (nextFt == ft || nextFt.meta.leftOwner != tree) ft
     else getOnOrAfterOwned(nextFt, tree)
   }
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
index b8c6f1068d..789fef9b93 100644
--- a/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala
+++ b/scalafmt-core/shared/src/main/scala/org/scalafmt/internal/Router.scala
@@ -53,6 +53,8 @@ class Router(formatOps: FormatOps) {
     getLastNonTrivial,
     prevNonComment,
     nextNonComment,
+    nextAfterNonComment,
+    nextNonCommentAfter,
     prevNonCommentSameLine,
     nextNonCommentSameLine
   }
@@ -300,7 +302,7 @@ class Router(formatOps: FormatOps) {
                 val indent = style.indent.main
                 val annoFT = tokens(anno)
                 val arrow = annoFT.left.is[T.RightArrow]
-                val expire = if (arrow) annoFT else next(nextNonComment(annoFT))
+                val expire = if (arrow) annoFT else nextAfterNonComment(annoFT)
                 (expire, Some(expire.left), indent, Some(isSelfAnnotationNL))
               }
               annoOpt.getOrElse { (null, None, 0, None) }
@@ -1436,7 +1438,7 @@ class Router(formatOps: FormatOps) {
         val expire = expireFt.left
         val sameLineSplit = Space(endsWithSymbolIdent(left))
         val bopSplits = style.newlines.getBeforeOpenParenDefnSite.map { x =>
-          val ob = OptionalBraces.get(next(nextNonComment(expireFt))).nonEmpty
+          val ob = OptionalBraces.get(nextAfterNonComment(expireFt)).nonEmpty
           def extraIfBody = style.indent.extraBeforeOpenParenDefnSite
           val indent =
             if (ob) style.indent.getSignificant + extraIfBody
@@ -2442,7 +2444,7 @@ class Router(formatOps: FormatOps) {
         val baseSplit = Split(mod, 0)
         rightOwner match {
           case GetSelectLike(t) =>
-            val expire = nextNonComment(next(formatToken)).left
+            val expire = nextNonCommentAfter(formatToken).left
             val indent = Indent(style.indent.main, expire, ExpiresOn.After)
             val split = baseSplit.withIndent(indent)
             if (findPrevSelect(t, style.encloseSelectChains).isEmpty) Seq(split)