Fix `dest`-assignment for overriden `T.command`s, and add basic unit tests to validate the `dest` paths of overriden commands and overriden targets

7 files changed, 83 insertions, 26 deletions

diff --git a/core/src/mill/define/Discover.scala b/core/src/mill/define/Discover.scala
index 52f4ab77..5e0e9a05 100644
--- a/core/src/mill/define/Discover.scala
+++ b/core/src/mill/define/Discover.scala
@@ -1,11 +1,14 @@
 package mill.define
 import language.experimental.macros
-import ammonite.main.Router.EntryPoint
+import ammonite.main.Router.{EntryPoint, Overrides}
+import sourcecode.Compat.Context
 
 import scala.collection.mutable
 import scala.reflect.macros.blackbox
 
-case class Discover(value: Map[Class[_], Seq[EntryPoint[_]]])
+
+
+case class Discover(value: Map[Class[_], Seq[(Int, EntryPoint[_])]])
 object Discover {
   def apply[T]: Discover = macro applyImpl[T]
 
@@ -41,14 +44,20 @@ object Discover {
     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
+      val curCls = discoveredModuleType.asInstanceOf[router.c.Type]
+      val methods = router.getValsOrMeths(curCls)
+      val overridesRoutes = {
+        for{
+          m <- methods.toList
+          if m.returnType <:< weakTypeOf[mill.define.Command[_]].asInstanceOf[router.c.Type]
+        } yield (m.overrides.length, router.extractMethod(m, curCls).asInstanceOf[c.Tree])
+      }
+      if overridesRoutes.nonEmpty
+      val (overrides, routes) = overridesRoutes.unzip
+
     } yield {
       val lhs =  q"classOf[${discoveredModuleType.typeSymbol.asClass}]"
-      val rhs = q"scala.Seq[ammonite.main.Router.EntryPoint[${discoveredModuleType.typeSymbol.asClass}]](..$routes)"
+      val rhs = q"scala.Seq[(Int, ammonite.main.Router.EntryPoint[${discoveredModuleType.typeSymbol.asClass}])](..$overridesRoutes)"
       q"$lhs -> $rhs"
     }
 
diff --git a/core/src/mill/define/Task.scala b/core/src/mill/define/Task.scala
index a6d90922..f74171ec 100644
--- a/core/src/mill/define/Task.scala
+++ b/core/src/mill/define/Task.scala
@@ -1,13 +1,26 @@
 package mill.define
 
+import ammonite.main.Router.Overrides
 import mill.define.Applicative.Applyable
 import mill.eval.{PathRef, Result}
-
+import sourcecode.Compat.Context
 import upickle.default.{ReadWriter => RW, Reader => R, Writer => W}
 
 import scala.language.experimental.macros
 import scala.reflect.macros.blackbox.Context
 
+case class EnclosingClass(value: Class[_])
+object EnclosingClass{
+  def apply()(implicit c: EnclosingClass) = c.value
+  implicit def generate: EnclosingClass = macro impl
+  def impl(c: Context): c.Tree = {
+    import c.universe._
+    val cls = c.internal.enclosingOwner.owner.asType.asClass
+//    q"new _root_.mill.define.EnclosingClass(classOf[$cls])"
+    q"new _root_.mill.define.EnclosingClass(this.getClass)"
+  }
+}
+
 /**
   * Models a single node in the Mill build graph, with a list of inputs and a
   * single output of type [[T]].
@@ -164,19 +177,33 @@ object Target extends TargetGenerated with Applicative.Applyer[Task, Task, Resul
 
   def command[T](t: Task[T])
                 (implicit ctx: mill.define.Ctx,
-                 w: W[T]): Command[T] = new Command(t, ctx, w)
+                 w: W[T],
+                 cls: EnclosingClass,
+                 overrides: Overrides): Command[T] = {
+    new Command(t, ctx, w, cls.value, overrides.value)
+  }
 
   def command[T](t: Result[T])
                 (implicit w: W[T],
-                 ctx: mill.define.Ctx): Command[T] = macro commandImpl[T]
+                 ctx: mill.define.Ctx,
+                 cls: EnclosingClass,
+                 overrides: Overrides): Command[T] = macro commandImpl[T]
 
   def commandImpl[T: c.WeakTypeTag](c: Context)
                                    (t: c.Expr[T])
                                    (w: c.Expr[W[T]],
-                                    ctx: c.Expr[mill.define.Ctx]): c.Expr[Command[T]] = {
+                                    ctx: c.Expr[mill.define.Ctx],
+                                    cls: c.Expr[EnclosingClass],
+                                    overrides: c.Expr[Overrides]): c.Expr[Command[T]] = {
     import c.universe._
     reify(
-      new Command[T](Applicative.impl[Task, T, mill.util.Ctx](c)(t).splice, ctx.splice, w.splice)
+      new Command[T](
+        Applicative.impl[Task, T, mill.util.Ctx](c)(t).splice,
+        ctx.splice,
+        w.splice,
+        cls.splice.value,
+        overrides.splice.value
+      )
     )
   }
 
@@ -255,7 +282,9 @@ class TargetImpl[+T](t: Task[T],
 
 class Command[+T](t: Task[T],
                   ctx0: mill.define.Ctx,
-                  val writer: W[_]) extends NamedTaskImpl[T](ctx0, t) {
+                  val writer: W[_],
+                  val cls: Class[_],
+                  val overrides: Int) extends NamedTaskImpl[T](ctx0, t) {
   override def asCommand = Some(this)
 }
 
diff --git a/core/src/mill/eval/Evaluator.scala b/core/src/mill/eval/Evaluator.scala
index f1a34e69..d5c273d2 100644
--- a/core/src/mill/eval/Evaluator.scala
+++ b/core/src/mill/eval/Evaluator.scala
@@ -2,6 +2,7 @@ package mill.eval
 
 import java.net.URLClassLoader
 
+import ammonite.main.Router.EntryPoint
 import ammonite.ops._
 import ammonite.runtime.SpecialClassLoader
 import mill.define.{Ctx => _, _}
@@ -26,6 +27,7 @@ case class Labelled[T](target: NamedTask[T],
 class Evaluator[T](val workspacePath: Path,
                    val basePath: Path,
                    val rootModule: mill.Module,
+                   val discover: Discover,
                    log: Logger,
                    val classLoaderSig: Seq[(Path, Long)] = Evaluator.classLoaderSig){
 
@@ -38,16 +40,30 @@ class Evaluator[T](val workspacePath: Path,
     val topoSorted = Graph.topoSorted(transitive)
     val sortedGroups = Graph.groupAroundImportantTargets(topoSorted){
       case t: NamedTask[Any]   =>
+
         val segments = t.ctx.segments
-        val (finalTaskOverrides, enclosing) = t match{
+        val finalTaskOverrides = t match{
           case t: Target[_] =>
-            rootModule.millInternal.segmentsToTargets.get(segments).fold(0)(_.ctx.overrides) -> t.ctx.enclosing
-          case c: mill.define.Command[_] => 0 -> c.ctx.enclosing
-          case c: mill.define.Worker[_] => 0 -> c.ctx.enclosing
+            rootModule.millInternal.segmentsToTargets.get(segments).fold(0)(_.ctx.overrides)
+          case c: mill.define.Command[_] =>
+            def findMatching(cls: Class[_]): Option[Seq[(Int, EntryPoint[_])]] = {
+              discover.value.get(cls) match{
+                case Some(v) => Some(v)
+                case None =>
+                  cls.getSuperclass match{
+                    case null => None
+                    case superCls => findMatching(superCls)
+                  }
+              }
+            }
+            val public = findMatching(c.cls).get.find(_._2.name == c.ctx.segment.pathSegments.head).get._1
+            public
+          case c: mill.define.Worker[_] => 0
         }
+
         val additional =
           if (finalTaskOverrides == t.ctx.overrides) Nil
-          else Seq(Segment.Label("overriden"), Segment.Label(enclosing))
+          else Seq(Segment.Label("overriden"), Segment.Label(t.ctx.enclosing))
 
         Right(Labelled(t, segments ++ additional))
       case t if goals.contains(t) => Left(t)
diff --git a/core/src/mill/main/MainRunner.scala b/core/src/mill/main/MainRunner.scala
index d3053d7a..81450ba9 100644
--- a/core/src/mill/main/MainRunner.scala
+++ b/core/src/mill/main/MainRunner.scala
@@ -95,13 +95,15 @@ class MainRunner(config: ammonite.main.Cli.Config,
          |  // even if it does nothing...
          |  def $$main() = Iterator[String]()
          |
-         |  val millDiscover: mill.define.Discover = mill.define.Discover[this.type]
+         |  implicit def millDiscover: mill.define.Discover = mill.define.Discover[this.type]
          |  // Need to wrap the returned Module in Some(...) to make sure it
          |  // doesn't get picked up during reflective child-module discovery
          |  val millSelf = Some(this)
          |}
          |
          |sealed trait $wrapName extends mill.Module{this: mill.define.BaseModule =>
+         |
+         |  implicit def millDiscover: mill.define.Discover
          |""".stripMargin
     }
 
diff --git a/core/src/mill/main/ReplApplyHandler.scala b/core/src/mill/main/ReplApplyHandler.scala
index cc4d3c64..9b3a29bb 100644
--- a/core/src/mill/main/ReplApplyHandler.scala
+++ b/core/src/mill/main/ReplApplyHandler.scala
@@ -19,6 +19,7 @@ object ReplApplyHandler{
         ammonite.ops.pwd / 'out,
         ammonite.ops.pwd,
         rootModule,
+        discover,
         new mill.util.PrintLogger(
           colors != ammonite.util.Colors.BlackWhite,
           colors,
@@ -86,8 +87,8 @@ class ReplApplyHandler(pprinter0: pprint.PPrinter,
           case None => Nil
           case Some(commands) =>
             ctx.applyPrefixColor("\nCommands:").toString +: commands.map{c =>
-              "\n    ." + c.name + "(" +
-              c.argSignatures.map(s => s.name + ": " + s.typeString).mkString(", ") +
+              "\n    ." + c._2.name + "(" +
+              c._2.argSignatures.map(s => s.name + ": " + s.typeString).mkString(", ") +
                 ")()"
             }
         }) ++
diff --git a/core/src/mill/main/Resolve.scala b/core/src/mill/main/Resolve.scala
index b01391d0..d8c0ede0 100644
--- a/core/src/mill/main/Resolve.scala
+++ b/core/src/mill/main/Resolve.scala
@@ -28,8 +28,8 @@ object Resolve {
           for{
             (cls, entryPoints) <- discover.value.filterKeys(_.isAssignableFrom(target.getClass))
             ep <- entryPoints
-            if ep.name == name
-          } yield ep.asInstanceOf[EntryPoint[mill.Module]].invoke(target, ammonite.main.Scripts.groupArgs(rest.toList)) match {
+            if ep._2.name == name
+          } yield ep._2.asInstanceOf[EntryPoint[mill.Module]].invoke(target, ammonite.main.Scripts.groupArgs(rest.toList)) match {
             case Router.Result.Success(v) => Right(v)
             case _ => Left(s"Command failed $last")
           }
diff --git a/core/src/mill/main/RunScript.scala b/core/src/mill/main/RunScript.scala
index a69f87be..b4de2966 100644
--- a/core/src/mill/main/RunScript.scala
+++ b/core/src/mill/main/RunScript.scala
@@ -43,7 +43,7 @@ object RunScript{
               for((mapping, discover) <- evaluateMapping(wd, path, interp))
               yield (
                 new Evaluator(
-                  wd / 'out, wd, mapping, log,
+                  wd / 'out, wd, mapping, discover, log,
                   mapping.getClass.getClassLoader.asInstanceOf[SpecialClassLoader].classpathSignature
                 ),
                 discover
@@ -120,7 +120,7 @@ object RunScript{
         Util.withContextClassloader(interp.evalClassloader) {
           Res.Success(
             buildCls.getMethod("millDiscover")
-                    .invoke(null)
+                    .invoke(module)
                     .asInstanceOf[Discover]
           )
         }