Introduce DocMiniPhase concept, fused phases

Like dotty, dottydoc can now also enjoy the modularity of having
fused mini-phases. By extending `DocMiniPhase` and overriding the
appropriate methods, one can alter the DocAST easily.

ping: @DarkDimius

8 files changed, 312 insertions, 105 deletions

diff --git a/dottydoc/jvm/src/dotty/tools/dottydoc/DottyDoc.scala b/dottydoc/jvm/src/dotty/tools/dottydoc/DottyDoc.scala
index db2a394be..5eead13c0 100644
--- a/dottydoc/jvm/src/dotty/tools/dottydoc/DottyDoc.scala
+++ b/dottydoc/jvm/src/dotty/tools/dottydoc/DottyDoc.scala
@@ -1,7 +1,8 @@
 package dotty.tools
 package dottydoc
 
-import core.{ AddImplicitsPhase, DocPhase }
+import core._
+import core.transform._
 import dotc.config.CompilerCommand
 import dotc.config.Printers.dottydoc
 import dotc.core.Contexts._
@@ -10,8 +11,6 @@ import dotc.typer.FrontEnd
 import dotc.{ CompilationUnit, Compiler, Driver, Run }
 import io.PlainFile
 
-import scala.util.control.NonFatal
-
 /** Custom Compiler with phases for the documentation tool
  *
  *  The idea here is to structure `dottydoc` around the new infrastructure. As
@@ -26,11 +25,14 @@ import scala.util.control.NonFatal
  *    5. Serve content using Scala.js
  */
 class DottyDocCompiler extends Compiler {
-  override def phases: List[List[Phase]] =
-    List(new DocFrontEnd) ::
-    List(new AddImplicitsPhase) ::
-    List(new DocPhase) ::
-    Nil
+  override def phases: List[List[Phase]] = List(
+    List(new DocFrontEnd),
+    List(new AddImplicitsPhase),
+    List(new DocASTPhase),
+    List(DocMiniTransformations(new ReturnTypeLinker,
+                                new ParamListTypeLinker)),
+    List(new PrintPhase)
+  )
 
   override def newRun(implicit ctx: Context): Run = {
     reset()
@@ -38,20 +40,6 @@ class DottyDocCompiler extends Compiler {
   }
 }
 
-/** TODO: this object is only temporary, should be factored out, just wanted it to demo dottydoc */
-object ImplicitlyAdded {
-  import dotc.core.Symbols.Symbol
-  import dotc.ast.tpd.DefDef
-  import scala.collection.mutable
-
-  private var _defs: Map[Symbol, Set[Symbol]] = Map.empty
-  def defs(sym: Symbol): Set[Symbol] = _defs.get(sym).getOrElse(Set.empty)
-
-  def addDef(s: Symbol, d: Symbol): Unit = _defs = (_defs + {
-    s -> _defs.get(s).map(xs => xs + d).getOrElse(Set(d))
-  })
-}
-
 class DocFrontEnd extends FrontEnd {
   override protected def discardAfterTyper(unit: CompilationUnit)(implicit ctx: Context) =
     unit.isJava
diff --git a/dottydoc/jvm/src/dotty/tools/dottydoc/core/AddImplicitsPhase.scala b/dottydoc/jvm/src/dotty/tools/dottydoc/core/AddImplicitsPhase.scala
index 15994ee28..3b1f2b26a 100644
--- a/dottydoc/jvm/src/dotty/tools/dottydoc/core/AddImplicitsPhase.scala
+++ b/dottydoc/jvm/src/dotty/tools/dottydoc/core/AddImplicitsPhase.scala
@@ -19,7 +19,7 @@ class AddImplicitsPhase extends MiniPhaseTransform { thisTransformer =>
       tree.vparamss(0).length == 1       // should only take one arg, since it has to be a transformation
     ) {
       val convertee = tree.vparamss(0)(0).symbol.info.widenDealias.finalResultType.typeSymbol//denot.typeRef // the pimped type (i.e. `class`)
-      ImplicitlyAdded.addDef(convertee, tree.symbol.info.widenDealias.finalResultType.typeSymbol)
+      ctx.base.addDef(convertee, tree.symbol.info.widenDealias.finalResultType.typeSymbol)
     }
 
     tree
diff --git a/dottydoc/jvm/src/dotty/tools/dottydoc/core/DocPhase.scala b/dottydoc/jvm/src/dotty/tools/dottydoc/core/DocASTPhase.scala
index 522b224d1..72e679565 100644
--- a/dottydoc/jvm/src/dotty/tools/dottydoc/core/DocPhase.scala
+++ b/dottydoc/jvm/src/dotty/tools/dottydoc/core/DocASTPhase.scala
@@ -10,11 +10,11 @@ import dotc.core.Contexts.Context
 import dotc.core.Phases.Phase
 import dotc.core.Symbols.Symbol
 
-class DocPhase extends Phase {
-  import model.parsers.{ WikiParser, ReturnTypeLinker, ParamListLinker }
+class DocASTPhase extends Phase {
   import model._
   import model.factories._
   import model.internal._
+  import model.parsers.WikiParser
   import model.comment.Comment
   import dotty.tools.dotc.core.Flags
   import dotty.tools.dotc.ast.tpd._
@@ -24,10 +24,6 @@ class DocPhase extends Phase {
   def phaseName = "docphase"
 
   private[this] val commentParser = new WikiParser
-  private[this] val linkers =
-    new ReturnTypeLinker ::
-    new ParamListLinker  ::
-    Nil
 
   /** Saves the commentParser function for later evaluation, for when the AST has been filled */
   def track(symbol: Symbol, ctx: Context)(op: => Entity) = {
@@ -41,7 +37,7 @@ class DocPhase extends Phase {
 
   /** Build documentation hierarchy from existing tree */
   def collect(tree: Tree, prev: List[String] = Nil)(implicit ctx: Context): Entity = track(tree.symbol, ctx) {
-    val implicitlyAddedMembers = ImplicitlyAdded.defs(tree.symbol)//denot.info)
+    val implicitlyAddedMembers = ctx.base.defs(tree.symbol)
 
     def collectList(xs: List[Tree], ps: List[String])(implicit ctx: Context): List[Entity] =
       xs.map(collect(_, ps)).filter(_ != NonEntity)
@@ -165,25 +161,16 @@ class DocPhase extends Phase {
       child  <- parent.children
     } setParent(child, to = parent)
 
-
-    // (3) Set returntypes to correct entities
-    for (linker <- linkers) linker.link(packages)
-
     // (3) Create documentation template from docstrings, with internal links
     println("Generating documentation, this might take a while...")
     commentParser.parse(packages)
 
-    // (4) Write the finished model to JSON
-    val outputDir = {
-      val out = ctx.settings.DocOutput.value
-      if (out.last == '/') out.dropRight(1)
-      else out
-    }
-    if (!ctx.settings.YDocNoWrite.value) (new util.OutputWriter).write(packages, outputDir)
-
-    // (5) Clear caches
+    // (4) Clear caches
     commentParser.clear()
 
+    // (5) Update Doc AST in ctx.base
+    for (kv <- packages) ctx.base.packages += kv
+
     // Return super's result
     compUnits
   }
diff --git a/dottydoc/jvm/src/dotty/tools/dottydoc/core/MiniPhaseTransform.scala b/dottydoc/jvm/src/dotty/tools/dottydoc/core/MiniPhaseTransform.scala
new file mode 100644
index 000000000..ba7ea040a
--- /dev/null
+++ b/dottydoc/jvm/src/dotty/tools/dottydoc/core/MiniPhaseTransform.scala
@@ -0,0 +1,189 @@
+package dotty.tools
+package dottydoc
+package core
+
+import dotc.CompilationUnit
+import dotc.core.Contexts.Context
+import dotc.core.Phases.Phase
+import model._
+import model.internal._
+
+object transform {
+  /**
+   * The idea behind DocMiniTransformations is to fuse transformations to the
+   * doc AST, much like `MiniPhaseTransform` in dotty core - but in a much more
+   * simple implementation
+   *
+   * Usage
+   * -----
+   *
+   * Create a `DocMiniPhase` which overrides the relevant method:
+   *
+   * {{{
+   *    override def transformDef(implicit ctx: Context) = {
+   *      case x if shouldTransform(x) => x.copy(newValue = ...)
+   *    }
+   * }}}
+   *
+   * On each node in the AST, the appropriate method in `DocMiniPhase` will be
+   * called in the order that they are supplied in
+   * `DocMiniphaseTransformations`.
+   *
+   * There won't be a match-error as `transformX` is composed with an
+   * `identity` function.
+   *
+   * The transformations in `DocMiniTransformations` will apply transformations
+   * to all nodes - this means that you do _not_ need to transform children in
+   * `transformPackage`, because `transformX` will be called for the relevant
+   * children. If you want to add children to `Package` you need to do that in
+   * `transformPackage`, these additions will be persisted.
+   */
+  abstract class DocMiniTransformations(transformations: List[DocMiniPhase]) extends Phase {
+
+    override def runOn(units: List[CompilationUnit])(implicit ctx: Context): List[CompilationUnit] = {
+      for {
+        rootName    <- rootPackages
+        pack        =  ctx.base.packages[Package](rootName)
+        transformed =  performPackageTransform(pack)
+      } yield ctx.base.packages(rootName) = transformed
+      super.runOn(units)
+    }
+
+    private def rootPackages(implicit ctx: Context): List[String] = {
+      var currentDepth = Int.MaxValue
+      var packs = List.empty[String]
+
+      for (key <- ctx.base.packages.keys) {
+        val keyDepth = key.split("\\.").length
+        packs =
+          if (keyDepth < currentDepth) {
+            currentDepth = keyDepth
+            key :: Nil
+          } else if (keyDepth == currentDepth) {
+            key :: packs
+          } else packs
+      }
+      packs
+    }
+
+    private def performPackageTransform(pack: Package)(implicit ctx: Context): Package = {
+      def transformEntity[E <: Entity](e: E, f: DocMiniPhase => E => E)(createNew: E => E): E = {
+        val transformedEntity = transformations.foldLeft(e) { case (oldE, transf) =>
+          f(transf)(oldE)
+        }
+        createNew(transformedEntity)
+      }
+
+      def traverse(ent: Entity): Entity = ent match {
+        case p: Package => transformEntity(p, _.packageTransformation) { p =>
+          PackageImpl(
+            p.name,
+            p.members.map(traverse),
+            p.path,
+            p.comment
+          )
+        }
+        case c: Class => transformEntity(c, _.classTransformation) { newClass =>
+          ClassImpl(
+            newClass.name,
+            newClass.members.map(traverse),
+            newClass.modifiers,
+            newClass.path,
+            newClass.typeParams,
+            newClass.superTypes,
+            newClass.comment
+          )
+        }
+        case cc: CaseClass => transformEntity(cc, _.cclassTransformation) { cc =>
+          CaseClassImpl(
+            cc.name,
+            cc.members.map(traverse),
+            cc.modifiers,
+            cc.path,
+            cc.typeParams,
+            cc.superTypes,
+            cc.comment
+          )
+        }
+        case trt: Trait => transformEntity(trt, _.traitTransformation) { trt =>
+          TraitImpl(
+            trt.name,
+            trt.members.map(traverse),
+            trt.modifiers,
+            trt.path,
+            trt.typeParams,
+            trt.superTypes,
+            trt.comment
+          )
+        }
+        case obj: Object => transformEntity(obj, _.objectTransformation) { obj =>
+          ObjectImpl(
+            obj.name,
+            obj.members.map(traverse),
+            obj.modifiers,
+            obj.path,
+            obj.superTypes,
+            obj.comment
+          )
+        }
+        case df: Def => transformEntity(df, _.defTransformation) { df =>
+          DefImpl(
+            df.name,
+            df.modifiers,
+            df.path,
+            df.returnValue,
+            df.typeParams,
+            df.paramLists,
+            df.comment
+          )
+        }
+        case vl: Val => transformEntity(vl, _.valTransformation) { vl =>
+          ValImpl(
+            vl.name,
+            vl.modifiers,
+            vl.path,
+            vl.returnValue,
+            vl.comment
+          )
+        }
+      }
+
+      traverse(pack).asInstanceOf[Package]
+    }
+
+    override def run(implicit ctx: Context): Unit = ()
+  }
+
+  object DocMiniTransformations {
+    private var previousPhase = 0
+    def apply(transformations: DocMiniPhase*) =
+      new DocMiniTransformations(transformations.toList) {
+        val packages = Map.empty[String, Package]
+
+        def phaseName = s"MiniTransformation${ previousPhase += 1 }"
+      }
+  }
+
+  trait DocMiniPhase { phase =>
+    private def identity[E]: PartialFunction[E, E] = {
+      case id => id
+    }
+
+    // Partial functions instead????
+    def transformPackage(implicit ctx: Context): PartialFunction[Package, Package] = identity
+    def transformClass(implicit ctx: Context): PartialFunction[Class, Class] = identity
+    def transformCaseClass(implicit ctx: Context): PartialFunction[CaseClass, CaseClass] = identity
+    def transformTrait(implicit ctx: Context): PartialFunction[Trait, Trait] = identity
+    def transformObject(implicit ctx: Context): PartialFunction[Object, Object] = identity
+    def transformDef(implicit ctx: Context): PartialFunction[Def, Def] = identity
+    def transformVal(implicit ctx: Context): PartialFunction[Val, Val] = identity
+
+    private[transform] def packageTransformation(p: Package)(implicit ctx: Context) = (transformPackage orElse identity)(p)
+    private[transform] def classTransformation(cls: Class)(implicit ctx: Context) = (transformClass orElse identity)(cls)
+    private[transform] def cclassTransformation(cc: CaseClass)(implicit ctx: Context) = (transformCaseClass orElse identity)(cc)
+    private[transform] def traitTransformation(trt: Trait)(implicit ctx: Context) = (transformTrait orElse identity)(trt)
+    private[transform] def objectTransformation(obj: Object)(implicit ctx: Context) = (transformObject orElse identity)(obj)
+    private[transform] def defTransformation(df: Def)(implicit ctx: Context) = (transformDef orElse identity)(df)
+    private[transform] def valTransformation(vl: Val)(implicit ctx: Context) = (transformVal orElse identity)(vl)
+  }
+}
diff --git a/dottydoc/jvm/src/dotty/tools/dottydoc/core/PrintPhase.scala b/dottydoc/jvm/src/dotty/tools/dottydoc/core/PrintPhase.scala
new file mode 100644
index 000000000..9c9070d24
--- /dev/null
+++ b/dottydoc/jvm/src/dotty/tools/dottydoc/core/PrintPhase.scala
@@ -0,0 +1,30 @@
+package dotty.tools
+package dottydoc
+package core
+
+import dotc.CompilationUnit
+import dotc.core.Contexts.Context
+import dotc.core.Phases.Phase
+import model.{Package, Entity}
+
+/** TODO: re-write to `DocMiniPhase` */
+class PrintPhase extends Phase {
+  def phaseName = "docPrintPhase"
+
+  var currentRun = 0
+  override def run(implicit ctx: Context): Unit = ()
+
+  override def runOn(units: List[CompilationUnit])(implicit ctx: Context): List[CompilationUnit] = {
+    val compUnits = super.runOn(units)
+    val packages = ctx.base.packages[Package].toMap
+
+    val outputDir = {
+      val out = ctx.settings.DocOutput.value
+      if (out.last == '/') out.dropRight(1)
+      else out
+    }
+    if (!ctx.settings.YDocNoWrite.value) (new util.OutputWriter).write(packages, outputDir)
+
+    compUnits
+  }
+}
diff --git a/dottydoc/jvm/src/dotty/tools/dottydoc/core/TypeLinkerPhases.scala b/dottydoc/jvm/src/dotty/tools/dottydoc/core/TypeLinkerPhases.scala
new file mode 100644
index 000000000..5339844ba
--- /dev/null
+++ b/dottydoc/jvm/src/dotty/tools/dottydoc/core/TypeLinkerPhases.scala
@@ -0,0 +1,69 @@
+package dotty.tools
+package dottydoc
+package core
+
+import dotc.core.Contexts.Context
+import dotc.util.Positions.NoPosition
+
+import transform.DocMiniPhase
+import model._
+import comment._
+import BodyParsers._
+import model.internal._
+import util.MemberLookup
+import util.traversing._
+import util.internal.setters._
+
+class ReturnTypeLinker extends DocMiniPhase with TypeLinker {
+  override def transformDef(implicit ctx: Context) = { case df: DefImpl =>
+    val returnValue = linkReference(df, df.returnValue, ctx.base.packages[Package].toMap)
+    df.copy(returnValue = returnValue)
+  }
+
+  override def transformVal(implicit ctx: Context) = { case vl: ValImpl =>
+    val returnValue = linkReference(vl, vl.returnValue, ctx.base.packages[Package].toMap)
+    vl.copy(returnValue = returnValue)
+  }
+}
+
+class ParamListTypeLinker extends DocMiniPhase with TypeLinker {
+  override def transformDef(implicit ctx: Context) = { case df: DefImpl =>
+    val newParamLists = for {
+      list <- df.paramLists
+      newList = list.map(linkReference(df, _, ctx.base.packages[Package].toMap))
+    } yield newList.asInstanceOf[List[NamedReference]]
+
+    df.copy(paramLists = newParamLists)
+  }
+}
+
+trait TypeLinker extends MemberLookup {
+  def linkReference(ent: Entity, rv: Reference, packs: Map[String, Package]): Reference =
+    rv match {
+      case rv @ TypeReference(_, UnsetLink(t, query), tps) =>
+        val inlineToHtml = InlineToHtml(ent)
+        val title = inlineToHtml(t)
+
+        def handleEntityLink(title: String, lt: LinkTo): MaterializableLink = lt match {
+          case Tooltip(str)           => NoLink(title, str)
+          case LinkToExternal(_, url) => MaterializedLink(title, url)
+          case LinkToEntity(target)   => MaterializedLink(title, util.traversing.relativePath(ent, target))
+        }
+
+        val target = handleEntityLink(title, makeEntityLink(ent, packs, t, NoPosition, query).link)
+
+        val tpTargets = tps.map {
+          case UnsetLink(t, query) =>
+            handleEntityLink(inlineToHtml(t), makeEntityLink(ent, packs, t, NoPosition, query).link)
+          case x => x
+        }
+
+        rv.copy(tpeLink = target, paramLinks = tpTargets)
+      case rv @ OrTypeReference(left, right) =>
+        rv.copy(left = linkReference(ent, left, packs), right = linkReference(ent, right, packs))
+      case rv @ AndTypeReference(left, right) =>
+        rv.copy(left = linkReference(ent, left, packs), right = linkReference(ent, right, packs))
+      case rv @ NamedReference(_, ref) => rv.copy(ref = linkReference(ent, ref, packs))
+      case _ => rv
+    }
+}
diff --git a/dottydoc/jvm/src/dotty/tools/dottydoc/model/parsers.scala b/dottydoc/jvm/src/dotty/tools/dottydoc/model/parsers.scala
index 2c163b1e6..3b2f2df43 100644
--- a/dottydoc/jvm/src/dotty/tools/dottydoc/model/parsers.scala
+++ b/dottydoc/jvm/src/dotty/tools/dottydoc/model/parsers.scala
@@ -83,60 +83,4 @@ object parsers {
 
     def clear(): Unit = commentCache = Map.empty
   }
-
-  sealed trait TypeLinker extends MemberLookup {
-    protected def linkReference(ent: Entity, rv: Reference, packs: Map[String, Package]): Reference =
-      rv match {
-        case rv @ TypeReference(_, UnsetLink(t, query), tps) =>
-          val inlineToHtml = InlineToHtml(ent)
-          val title = inlineToHtml(t)
-
-          def handleEntityLink(title: String, lt: LinkTo): MaterializableLink = lt match {
-            case Tooltip(str)           => NoLink(title, str)
-            case LinkToExternal(_, url) => MaterializedLink(title, url)
-            case LinkToEntity(target)   => MaterializedLink(title, util.traversing.relativePath(ent, target))
-          }
-
-          val target = handleEntityLink(title, makeEntityLink(ent, packs, t, NoPosition, query).link)
-
-          val tpTargets = tps.map {
-            case UnsetLink(t, query) =>
-              handleEntityLink(inlineToHtml(t), makeEntityLink(ent, packs, t, NoPosition, query).link)
-            case x => x
-          }
-
-          rv.copy(tpeLink = target, paramLinks = tpTargets)
-        case rv @ OrTypeReference(left, right) =>
-          rv.copy(left = linkReference(ent, left, packs), right = linkReference(ent, right, packs))
-        case rv @ AndTypeReference(left, right) =>
-          rv.copy(left = linkReference(ent, left, packs), right = linkReference(ent, right, packs))
-        case rv @ NamedReference(_, ref) => rv.copy(ref = linkReference(ent, ref, packs))
-        case _ => rv
-      }
-
-    def link(packs: Map[String, Package]): Unit
-  }
-
-  class ReturnTypeLinker extends TypeLinker {
-    def link(packs: Map[String, Package]): Unit =
-      for (pack <- packs.values) mutateEntities(pack) {
-        case ent: ReturnValue =>
-          setReturnValue(ent, linkReference(ent, ent.returnValue, packs))
-        case _ => ()
-      }
-  }
-
-  class ParamListLinker extends TypeLinker {
-    def link(packs: Map[String, Package]): Unit =
-      for (pack <- packs.values) mutateEntities(pack) {
-        case ent: Def =>
-          val newParamLists = for {
-            list <- ent.paramLists
-            newList = list.map(linkReference(ent, _, packs))
-          } yield newList.asInstanceOf[List[NamedReference]]
-
-          setParamLists(ent, newParamLists)
-        case _ => ()
-      }
-  }
 }
diff --git a/dottydoc/jvm/test/BaseTest.scala b/dottydoc/jvm/test/BaseTest.scala
index 1a2fb25cf..0c6fad1c8 100644
--- a/dottydoc/jvm/test/BaseTest.scala
+++ b/dottydoc/jvm/test/BaseTest.scala
@@ -6,7 +6,7 @@ import Contexts.{ Context, ContextBase, FreshContext }
 import dotc.util.SourceFile
 import dotc.core.Phases.Phase
 import dotc.typer.FrontEnd
-import dottydoc.core.DocPhase
+import dottydoc.core.DocASTPhase
 
 trait DottyTest {
   dotty.tools.dotc.parsing.Scanners // initialize keywords
@@ -22,8 +22,8 @@ trait DottyTest {
     ctx
   }
 
-  private def compilerWithChecker(assertion: DocPhase => Unit) = new DottyDocCompiler {
-    private[this] val docPhase = new DocPhase
+  private def compilerWithChecker(assertion: DocASTPhase => Unit) = new DottyDocCompiler {
+    private[this] val docPhase = new DocASTPhase
 
     override def phases =
       List(new FrontEnd) ::
@@ -35,21 +35,21 @@ trait DottyTest {
       Nil
   }
 
-  def checkSource(source: String)(assertion: DocPhase => Unit): Unit = {
+  def checkSource(source: String)(assertion: DocASTPhase => Unit): Unit = {
     val c = compilerWithChecker(assertion)
     c.rootContext(ctx)
     val run = c.newRun
     run.compile(source)
   }
 
-  def checkFiles(sources: List[String])(assertion: DocPhase => Unit): Unit = {
+  def checkFiles(sources: List[String])(assertion: DocASTPhase => Unit): Unit = {
     val c = compilerWithChecker(assertion)
     c.rootContext(ctx)
     val run = c.newRun
     run.compile(sources)
   }
 
-  def checkSources(sourceFiles: List[SourceFile])(assertion: DocPhase => Unit): Unit = {
+  def checkSources(sourceFiles: List[SourceFile])(assertion: DocASTPhase => Unit): Unit = {
     val c = compilerWithChecker(assertion)
     c.rootContext(ctx)
     val run = c.newRun