Owner tree work

build.sbt

@@ -17,6 +17,8 @@ lazy val sharedSettings: Def.SettingsDefinition = Def.settings(
       "org.scalameta" %% "scalameta" % "2.1.3" ::
       "org.scalameta" %% "contrib" % "2.1.3" ::
       "org.typelevel" %% "cats-core" % "1.0.1" ::
+      "org.typelevel" %% "cats-free" % "1.0.1" ::
+      "com.github.julien-truffaut" %% "monocle-core" % "1.5.0-cats" ::
       "org.scalactic" %% "scalactic" % "3.0.4" ::
       "org.scalactic" %% "scalactic" % "3.0.4" ::
       "org.scalatest" %% "scalatest" % "3.0.4" % "test" :: Nil,
@@ -29,8 +31,11 @@ lazy val sharedSettings: Def.SettingsDefinition = Def.settings(
 lazy val scalagen =
   project
     .in(file("scalagen"))
-    .settings(sharedSettings)
-    .settings(name := "scalagen")
+    .settings(
+      sharedSettings,
+      name := "scalagen",
+      addCompilerPlugin("org.spire-math" %% "kind-projector" % "0.9.4")
+    )
 
 // JVM sbt plugin
 lazy val sbtScalagen =

scalagen/src/main/scala/org/scalameta/scalagen/GeneratorTree.scala

@@ -0,0 +1,167 @@
+package org.scalameta.scalagen
+
+import cats._
+import cats.free._
+import cats.implicits._
+import monocle._
+
+import scala.meta.{XtensionShow => _, _}
+import scala.meta.gen.TypeclassExtractors.retrieveAnnotExtractInstance
+import scala.meta.gen._
+
+case class GeneratorInputContext(src: Tree, toExpand: List[Generator])
+object GeneratorTree {
+
+  /**
+    * OwnerTree is just simple tree with an arbritraty ,
+    * The issue is we cannot use the name Tree as we do not want
+    * conflicts with scalameta.tree
+    *
+    * TODO: Consider moving this out of scalagen
+    */
+  type GeneratorTree = Cofree[List, GeneratorInputContext]
+
+  /**
+    * Partially applied alias for OwnerTree. Allows use as a Functor/Monad etc.
+    */
+  type GeneratorTreeF[A] = Cofree[List, A]
+
+  /**
+    * Lazily build a Cofree from this tree.
+    *
+    * Cofree objects are only created as the tree is traversed.
+    */
+  def apply(t: Tree): GeneratorTreeF[Tree] =
+    Cofree(t, Eval.later(t.children.map(apply)))
+
+  /**
+    * Produce a product prefix based tree representation.
+    *
+    * Includes all nodes of the tree.
+    *
+    * Defn.Class
+    *  - Defn.Var
+    *    - Defn.Def
+    */
+  implicit def treeShowInstance: Show[GeneratorTree] =
+    Show.show[GeneratorTree](genTraversalString(regularTraversal, _))
+
+  /**
+    * Will print all nodes visited by the given traversal
+    */
+  def genTraversalString(
+      t: Traversal[GeneratorTreeF[GeneratorInputContext], GeneratorInputContext],
+      ot: GeneratorTree): String = {
+    val childString =
+      ot.tailForced
+        .map(genTraversalString(t, _))
+        .filterNot(_.isEmpty)
+        .flatMap(_.lines.toList)
+        .mkString("\n  ")
+
+    val res = t.headOption(ot) match {
+      case None if childString.isEmpty =>
+        ""
+      case None =>
+        error(childString)
+        childString
+      case Some(ctx) if childString.isEmpty =>
+        s" - ${treePrefixAndName(ctx.src)}"
+      case Some(ctx) =>
+        s""" - ${treePrefixAndName(ctx.src)}
+           |  $childString""".stripMargin
+    }
+
+    res
+  }
+
+  /**
+    * Primarily used for debug
+    *
+    * For example
+    * "Defn.Class: Foo"
+    *
+    * TODO: Make an extract/replace instance for names
+    */
+  private def treePrefixAndName(tree: Tree) = {
+    val nameStr =
+      tree match {
+        case o: Pkg => o.ref.syntax + "." + o.name.syntax
+        case o: Defn.Object => o.name.syntax
+        case c: Defn.Class => c.name.syntax
+        case t: Defn.Trait => t.name.syntax
+        case t: Defn.Type => t.name.syntax
+        case t: Decl.Type => t.name.syntax
+        case d: Defn.Def => d.name.syntax
+        case d: Decl.Def => d.name.syntax
+        case v: Defn.Val => genNameSyntax(v.pats)
+        case v: Decl.Val => genNameSyntax(v.pats)
+        case v: Defn.Var => genNameSyntax(v.pats)
+        case v: Decl.Var => genNameSyntax(v.pats)
+        case s: Term.Select => s.syntax
+        case s: Type.Select => s.syntax
+        case n: Term.Name => n.syntax
+        case n: Type.Name => n.syntax
+        case _ => ""
+      }
+    if (nameStr.isEmpty) {
+      tree.productPrefix
+    } else {
+      tree.productPrefix + s": $nameStr"
+    }
+  }
+
+  private def genNameSyntax(pats: List[Pat]): String =
+    pats
+      .collect({ case Pat.Var(name) => name })
+      .mkString(", ")
+
+  def regularTraversal[A]: Traversal[GeneratorTreeF[A], A] =
+    Traversal.fromTraverse[GeneratorTreeF, A](Traverse[GeneratorTreeF])
+
+  val ownerPrism: Prism[GeneratorTreeF[Tree], GeneratorTreeF[Tree]] =
+    Prism[GeneratorTreeF[Tree], GeneratorTreeF[Tree]](t => {
+      if (t.head.isOwner) Some(t)
+      else None
+    })(identity)
+
+  val ownerTraversal: Traversal[GeneratorTreeF[Tree], Tree] =
+    ownerPrism.composeTraversal(regularTraversal[Tree])
+
+  def generatorPrism(gs: Set[Generator]): Prism[GeneratorTreeF[Tree], GeneratorTree] =
+    Prism[GeneratorTreeF[Tree], GeneratorTree](t => {
+      if (hasGenerator(t.head, gs)) {
+        val context = GeneratorInputContext(t.head, getGeneratorsToExpand(t.head, gs))
+        Some(Cofree(context, Eval.now(List.empty[GeneratorTree])))
+      } else {
+        None
+      }
+    })(_.map(_.src))
+
+  private def hasMatchingGenerator(a: Mod.Annot, gs: Set[Generator]): Boolean =
+    gs.exists(isMatching(a, _))
+
+  private def isMatching(a: Mod.Annot, g: Generator) = {
+    a.init.tpe match {
+      case Type.Name(value) => g.name == value
+      case _ => false
+    }
+  }
+
+  private def hasGenerator(tree: Tree, gs: Set[Generator]): Boolean =
+    retrieveAnnotExtractInstance(tree)
+      .map(_.extract(tree))
+      .exists(_.exists(hasMatchingGenerator(_, gs)))
+
+  private def getGeneratorsToExpand(tree: Tree, gs: Set[Generator]): List[Generator] =
+    retrieveAnnotExtractInstance(tree)
+      .fold(List[Mod.Annot]())(_.extract(tree))
+      .flatMap(annot => gs.find(isMatching(annot, _)))
+
+  def generatorTraversal(
+      gs: Set[Generator]): Traversal[GeneratorTreeF[Tree], GeneratorInputContext] =
+    ownerPrism
+      .composePrism(generatorPrism(gs))
+      .composeTraversal(regularTraversal[GeneratorInputContext])
+
+}

scalagen/src/main/scala/org/scalameta/scalagen/Runner2.scala

@@ -0,0 +1,14 @@
+package org.scalameta.scalagen
+
+import org.scalameta.scalagen.GeneratorTree.{GeneratorTree, GeneratorTreeF}
+
+import scala.meta.{XtensionShow => _, _}
+import scala.meta.gen._
+
+object Runner2 {
+
+  def apply(t: Tree, gens: Set[Generator]): Option[Tree] =
+    expandGenerators(GeneratorTree(t), gens)
+
+  private def expandGenerators(t: GeneratorTreeF[Tree], gs: Set[Generator]): Option[Tree] = ???
+}

scalagen/src/main/scala/scala/meta/gen/TypeclassExtractors.scala

@@ -0,0 +1,106 @@
+package scala.meta.gen
+
+import scala.meta._
+import scala.meta.gen._
+import scala.meta.contrib._
+
+/** This is a big hack of a class.
+  * We have to retrieve instances at runtime based in input type
+  */
+// TODO move to scalameta/contrib
+object TypeclassExtractors {
+
+  def retrieveModExtractInstance[A](a: A): Option[Extract[A, Mod]] = {
+    val res =
+      a match {
+        case _: Defn.Class => Option(implicitly[Extract[Defn.Class, Mod]])
+        case _: Defn.Type => Option(implicitly[Extract[Defn.Type, Mod]])
+        case _: Defn.Trait => Option(implicitly[Extract[Defn.Trait, Mod]])
+        case _: Defn.Object => Option(implicitly[Extract[Defn.Object, Mod]])
+        case _: Defn.Def => Option(implicitly[Extract[Defn.Def, Mod]])
+        case _: Defn.Val => Option(implicitly[Extract[Defn.Val, Mod]])
+        case _: Defn.Var => Option(implicitly[Extract[Defn.Var, Mod]])
+        case _: Decl.Val => Option(implicitly[Extract[Decl.Val, Mod]])
+        case _: Decl.Var => Option(implicitly[Extract[Decl.Var, Mod]])
+        case _: Decl.Def => Option(implicitly[Extract[Decl.Def, Mod]])
+        case _: Decl.Type => Option(implicitly[Extract[Decl.Type, Mod]])
+        case _ => None
+      }
+
+    res.asInstanceOf[Option[Extract[A, Mod]]]
+  }
+
+  def retrieveModReplaceInstance[A](a: A): Option[Replace[A, Mod]] = {
+    val res =
+      a match {
+        case _: Defn.Class => Option(implicitly[Replace[Defn.Class, Mod]])
+        case _: Defn.Type => Option(implicitly[Replace[Defn.Type, Mod]])
+        case _: Defn.Trait => Option(implicitly[Replace[Defn.Trait, Mod]])
+        case _: Defn.Object => Option(implicitly[Replace[Defn.Object, Mod]])
+        case _: Defn.Def => Option(implicitly[Replace[Defn.Def, Mod]])
+        case _: Defn.Val => Option(implicitly[Replace[Defn.Val, Mod]])
+        case _: Defn.Var => Option(implicitly[Replace[Defn.Var, Mod]])
+        case _: Decl.Val => Option(implicitly[Replace[Decl.Val, Mod]])
+        case _: Decl.Var => Option(implicitly[Replace[Decl.Var, Mod]])
+        case _: Decl.Def => Option(implicitly[Replace[Decl.Def, Mod]])
+        case _: Decl.Type => Option(implicitly[Replace[Decl.Type, Mod]])
+        case _ => None
+      }
+
+    res.asInstanceOf[Option[Replace[A, Mod]]]
+  }
+
+  def retrieveAnnotExtractInstance[A](a: A): Option[Extract[A, Mod.Annot]] = {
+    val res =
+      a match {
+        case _: Defn.Class => Option(implicitly[Extract[Defn.Class, Mod.Annot]])
+        case _: Defn.Type => Option(implicitly[Extract[Defn.Type, Mod.Annot]])
+        case _: Defn.Trait => Option(implicitly[Extract[Defn.Trait, Mod.Annot]])
+        case _: Defn.Object => Option(implicitly[Extract[Defn.Object, Mod.Annot]])
+        case _: Defn.Def => Option(implicitly[Extract[Defn.Def, Mod.Annot]])
+        case _: Defn.Val => Option(implicitly[Extract[Defn.Val, Mod.Annot]])
+        case _: Defn.Var => Option(implicitly[Extract[Defn.Var, Mod.Annot]])
+        case _: Decl.Val => Option(implicitly[Extract[Decl.Val, Mod.Annot]])
+        case _: Decl.Var => Option(implicitly[Extract[Decl.Var, Mod.Annot]])
+        case _: Decl.Def => Option(implicitly[Extract[Decl.Def, Mod.Annot]])
+        case _: Decl.Type => Option(implicitly[Extract[Decl.Type, Mod.Annot]])
+        case _ => None
+      }
+
+    res.asInstanceOf[Option[Extract[A, Mod.Annot]]]
+  }
+
+  def retrieveStatReplaceInstance[A](a: A): Option[Replace[A, Stat]] = {
+    val res =
+      a match {
+        case _: Defn.Class => Option(implicitly[Replace[Defn.Class, Stat]])
+        case _: Defn.Trait => Option(implicitly[Replace[Defn.Trait, Stat]])
+        case _: Defn.Object => Option(implicitly[Replace[Defn.Object, Stat]])
+        case _: Defn.Def => Option(implicitly[Replace[Defn.Def, Stat]])
+        case _: Defn.Val => Option(implicitly[Replace[Defn.Val, Stat]])
+        case _: Defn.Var => Option(implicitly[Replace[Defn.Var, Stat]])
+        case _: Source => Option(implicitly[Replace[Source, Stat]])
+        case _: Pkg => Option(implicitly[Replace[Pkg, Stat]])
+        case _ => None
+      }
+
+    res.asInstanceOf[Option[Replace[A, Stat]]]
+
+  }
+
+  def retrieveStatExtractInstance[A](a: A): Option[Extract[A, Stat]] = {
+    val res =
+      a match {
+        case _: Defn.Class => Option(implicitly[Extract[Defn.Class, Stat]])
+        case _: Defn.Trait => Option(implicitly[Extract[Defn.Trait, Stat]])
+        case _: Defn.Object => Option(implicitly[Extract[Defn.Object, Stat]])
+        case _: Defn.Def => Option(implicitly[Extract[Defn.Def, Stat]])
+        case _: Defn.Val => Option(implicitly[Extract[Defn.Val, Stat]])
+        case _: Defn.Var => Option(implicitly[Extract[Defn.Var, Stat]])
+        case _: Source => Option(implicitly[Extract[Source, Stat]])
+        case _: Pkg => Option(implicitly[Extract[Pkg, Stat]])
+        case _ => None
+      }
+    res.asInstanceOf[Option[Extract[A, Stat]]]
+  }
+}