FormatOps: fix functionExpire when rewritten

Also, now it can be used in Router after `=>`.
scalameta · Feb 18, 2024 · 6c99a7a · 6c99a7a
1 parent 1d6ca94
commit 6c99a7a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 12 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
@@ -944,8 +944,8 @@ class FormatOps(
 
     def getRToks = dropWS(function.tokens.reverse)
     function.parent match {
-      case Some(b @ Term.Block(_ :: Nil)) =>
-        getLastToken(b) -> ExpiresOn.Before
+      case Some(p @ Term.Block(_ :: Nil)) =>
+        p.tokens.last -> ExpiresOn.Before
       case Some(Case(_, _, `function`)) =>
         orElse(dropComment(getRToks))
       case _ =>

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
@@ -407,30 +407,25 @@ class Router(formatOps: FormatOps) {
             _,
             StartsStatementRight(stmt)
           ) if leftOwner.isInstanceOf[Term.FunctionTerm] =>
-        val leftFuncBody = leftOwner.asInstanceOf[Term.FunctionTerm].body
-        val endOfFunction = getLastNonTrivial(leftFuncBody)
-        val endIndent = leftOwner.parent match {
-          case Some(Term.Block(_ :: Nil)) => nextNonComment(endOfFunction).left
-          case _ => endOfFunction.left
-        }
-        val canBeSpace = stmt.isInstanceOf[Term.FunctionTerm]
+        val leftFunc = leftOwner.asInstanceOf[Term.FunctionTerm]
         val (afterCurlySpace, afterCurlyNewlines) =
           getSpaceAndNewlineAfterCurlyLambda(newlines)
         val spaceSplit =
-          if (canBeSpace) Split(Space, 0)
+          if (stmt.isInstanceOf[Term.FunctionTerm]) Split(Space, 0)
           else if (
             afterCurlySpace &&
             (!rightOwner.is[Defn] || style.newlines.source.eq(Newlines.fold))
           )
             Split(Space, 0).withSingleLineNoOptimal(
-              getOptimalTokenFor(endOfFunction.left),
+              getOptimalTokenFor(getLastNonTrivial(leftFunc.body).left),
               noSyntaxNL = true
             )
           else Split.ignored
+        val (endIndent, expiresOn) = functionExpire(leftFunc)
         Seq(
           spaceSplit,
           Split(afterCurlyNewlines, 1)
-            .withIndent(style.indent.main, endIndent, After)
+            .withIndent(style.indent.main, endIndent, expiresOn)
         )
 
       case FormatToken(_: T.RightArrow | _: T.ContextArrow, right, _)