From 114c88313116cf99da5022afb9bfed149db6be0e Mon Sep 17 00:00:00 2001
From: Albert Meltzer <7529386+kitbellew@users.noreply.github.com>
Date: Fri, 7 Jul 2023 07:18:51 -0700
Subject: [PATCH] Router: extract template to handle `extends`

---
 .../scala/org/scalafmt/internal/FormatOps.scala |  7 +++++++
 .../scala/org/scalafmt/internal/Router.scala    | 17 ++++++-----------
 .../main/scala/org/scalafmt/util/TokenOps.scala | 13 +------------
 3 files changed, 14 insertions(+), 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
index c32e9bbecb..a30e65e419 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
@@ -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
index ef00764c04..c58902c6f4 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
@@ -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
index 04fcd53d50..31bc0301dc 100644
--- 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