Router: extract template to handle extends

scalameta · Jul 10, 2023 · 114c883 · 114c883
1 parent e0dafe0
commit 114c883
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 23 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
@@ -381,6 +381,13 @@ class FormatOps(
     }
   )
 
+  val TemplateOnRight = new ExtractFromMeta(ft =>
+    ft.meta.rightOwner match {
+      case ro: Template => Some(ro)
+      case _ => None
+    }
+  )
+
   def templateCurlyFt(template: Template): Option[FormatToken] =
     getStartOfTemplateBody(template).map(tokenBefore)
 

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
@@ -1978,22 +1978,17 @@ class Router(formatOps: FormatOps) {
         )
 
       // Template
-      case FormatToken(_, soft.ExtendsOrDerives(), _) =>
-        val template = defnTemplate(rightOwner)
-        def lastToken = template.fold(getLastNonTrivialToken(rightOwner)) { x =>
-          templateDerivesOrCurlyOrLastNonTrivial(x)
-        }
+      case FormatToken(_, soft.ExtendsOrDerives(), TemplateOnRight(template)) =>
+        def lastToken = templateDerivesOrCurlyOrLastNonTrivial(template)
 
         binPackParentConstructorSplits(
           true,
-          template.toSet,
-          template.flatMap(findTemplateGroupOnRight(_.superType)),
+          Set(template),
+          findTemplateGroupOnRight(_.superType)(template),
           lastToken,
           style.indent.extendSite,
-          template.exists(x =>
-            if (x.early.nonEmpty) x.inits.nonEmpty
-            else x.inits.lengthCompare(1) > 0
-          )
+          if (template.early.nonEmpty) template.inits.nonEmpty
+          else template.inits.lengthCompare(1) > 0
         )
 
       // trait A extends B, C, D, E

diff --git a/scalafmt-core/shared/src/main/scala/org/scalafmt/util/TokenOps.scala b/scalafmt-core/shared/src/main/scala/org/scalafmt/util/TokenOps.scala
@@ -1,6 +1,6 @@
 package org.scalafmt.util
 
-import scala.meta.{Defn, Pkg, Source, Template, Term, Tree}
+import scala.meta._
 import scala.meta.tokens.Token
 import scala.meta.tokens.Token._
 import scala.meta.tokens.Tokens
@@ -112,17 +112,6 @@ object TokenOps {
   @inline
   def getMod(ft: FormatToken): Modification = Space.orNL(ft.newlinesBetween)
 
-  def defnTemplate(tree: Tree): Option[Template] =
-    tree match {
-      case t: Defn.Object => Some(t.templ)
-      case t: Defn.Class => Some(t.templ)
-      case t: Defn.Trait => Some(t.templ)
-      case t: Defn.Enum => Some(t.templ)
-      case t: Pkg.Object => Some(t.templ)
-      case t: Template => Some(t)
-      case _ => None
-    }
-
   val formatOnCode = Set(
     "@formatter:on", // IntelliJ
     "format: on" // scalariform