Use Transformer instead of Traverser

Author: Duhemm

compiler/src/dotty/tools/dotc/transform/SetRootTree.scala

@@ -3,37 +3,34 @@ package dotty.tools.dotc.transform
 import dotty.tools.dotc.CompilationUnit
 import dotty.tools.dotc.ast.tpd
 import dotty.tools.dotc.core.Contexts.Context
-import dotty.tools.dotc.core.Phases.Phase
 
 /** Set the `rootTreeOrProvider` property of class symbols. */
-class SetRootTree extends Phase {
+class SetRootTree extends MacroTransform {
 
   override val phaseName: String = SetRootTree.name
   override def isRunnable(implicit ctx: Context) =
     super.isRunnable && ctx.settings.YretainTrees.value
 
-  override def run(implicit ctx: Context): Unit = {
-    val tree = ctx.compilationUnit.tpdTree
-    traverser.traverse(tree)
-  }
-
-  private def traverser = new tpd.TreeTraverser {
-    override def traverse(tree: tpd.Tree)(implicit ctx: Context): Unit = {
-      tree match {
-        case pkg: tpd.PackageDef =>
-          traverseChildren(pkg)
-        case td: tpd.TypeDef =>
-          if (td.symbol.isClass) {
-            val sym = td.symbol.asClass
-            tpd.sliceTopLevel(ctx.compilationUnit.tpdTree, sym) match {
-              case (pkg: tpd.PackageDef) :: Nil =>
-                sym.rootTreeOrProvider = pkg
-              case _ =>
-                sym.rootTreeOrProvider = td
+  override protected def newTransformer(implicit ctx: Context): Transformer = {
+    new Transformer {
+      override def transform(tree: tpd.Tree)(implicit ctx: Context): tpd.Tree = {
+        tree match {
+          case pkg: tpd.PackageDef =>
+            super.transform(pkg)
+          case td: tpd.TypeDef =>
+            if (td.symbol.isClass) {
+              val sym = td.symbol.asClass
+              tpd.sliceTopLevel(ctx.compilationUnit.tpdTree, sym) match {
+                case (pkg: tpd.PackageDef) :: Nil =>
+                  sym.rootTreeOrProvider = pkg
+                case _ =>
+                  sym.rootTreeOrProvider = td
+              }
             }
-          }
-        case _ =>
-          ()
+            td
+        case other =>
+          other
+        }
       }
     }
   }