core/src/main/scala/mill/define/Discover.scala


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

package mill.define
import language.experimental.macros
import ammonite.main.Router.EntryPoint

import scala.collection.mutable
import scala.reflect.macros.blackbox

case class Discover(value: Map[Class[_], Seq[EntryPoint[_]]])
object Discover {
  def apply[T]: Discover = macro applyImpl[T]

  def applyImpl[T: c.WeakTypeTag](c: blackbox.Context): c.Expr[Discover] = {
    import c.universe._
    import compat._
    val seen = mutable.Set.empty[Type]
    def rec(tpe: Type): Unit = {
      if (!seen(tpe)){
        seen.add(tpe)
        for{
          m <- tpe.members
          memberTpe = m.typeSignature
          if memberTpe.resultType <:< typeOf[mill.define.Module] && memberTpe.paramLists.isEmpty
        } rec(memberTpe.resultType)

        if (tpe <:< typeOf[mill.define.Cross[_]]){
          val inner = typeOf[Cross[_]]
            .typeSymbol
            .asClass
            .typeParams
            .head
            .asType
            .toType
            .asSeenFrom(tpe, typeOf[Cross[_]].typeSymbol)

          rec(inner)
        }
      }
    }
    rec(weakTypeOf[T])

    val router = new ammonite.main.Router(c)
    val mapping = for{
      discoveredModuleType <- seen
      val routes = router.getAllRoutesForClass(
        discoveredModuleType.asInstanceOf[router.c.Type],
        _.returnType <:< weakTypeOf[mill.define.Command[_]].asInstanceOf[router.c.Type]
      ).map(_.asInstanceOf[c.Tree])
      if routes.nonEmpty
    } yield {
      val lhs =  q"classOf[${discoveredModuleType.typeSymbol.asClass}]"
      val rhs = q"scala.Seq[ammonite.main.Router.EntryPoint[${discoveredModuleType.typeSymbol.asClass}]](..$routes)"
      q"$lhs -> $rhs"
    }

    c.Expr[Discover](q"mill.define.Discover(scala.collection.immutable.Map(..$mapping))")
  }
}