From 89af3603cddf377e6f956907ef5b58dfacc3336e Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Mon, 25 Nov 2024 09:47:10 -0800
Subject: [PATCH] FormatOps: support ContextArrow in `getFuncArrow`

---
 .../main/scala/org/scalafmt/internal/FormatOps.scala | 12 ++++++------
 1 file changed, 6 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 f4c1e359e..a8f8f02d3 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
@@ -1241,19 +1241,19 @@ class FormatOps(
   // look for arrow before body, if any, else after params
   def getFuncArrow(term: Term.FunctionTerm): Option[FT] = tokens
     .tokenBeforeOpt(term.body)
-    .orElse(tokenAfterOpt(term.paramClause).map(getArrowAfter))
+    .orElse(tokenAfterOpt(term.paramClause).map(getArrowAfter[T.FunctionArrow]))
 
   // look for arrow before body, if any, else after cond/pat
-  def getCaseArrow(term: Case): FT = tokenBeforeOpt(term.body)
-    .getOrElse(getArrowAfter(tokenAfter(term.cond.getOrElse(term.pat))))
+  def getCaseArrow(term: Case): FT = tokenBeforeOpt(term.body).getOrElse {
+    getArrowAfter[T.RightArrow](tokenAfter(term.cond.getOrElse(term.pat)))
+  }
 
   // look for arrow before body, if any, else after cond/pat
   def getCaseArrow(term: TypeCase): FT = next(tokenAfter(term.pat))
 
-  private def getArrowAfter(ft: FT): FT = {
+  private def getArrowAfter[A](ft: FT)(implicit f: Classifier[T, A]): FT = {
     val maybeArrow = next(ft)
-    if (maybeArrow.left.is[T.RightArrow]) maybeArrow
-    else nextAfterNonComment(maybeArrow)
+    if (f(maybeArrow.left)) maybeArrow else nextAfterNonComment(maybeArrow)
   }
 
   @tailrec