- DRY up `Resolve.scala`

- Add a version of `Resolve` that resolves names only, but works on entrypoints/etc. without needing arguments
- Fix tests to compile with new `multiSelect` parsing flag

3 files changed, 197 insertions, 136 deletions

diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
index 07ba37a9..56e43f6f 100644
--- a/main/src/mill/main/MainModule.scala
+++ b/main/src/mill/main/MainModule.scala
@@ -2,7 +2,7 @@ package mill.main
 
 import mill.define.NamedTask
 import mill.eval.{Evaluator, Result}
-import mill.util.{PrintLogger, Watched}
+import mill.util.{EitherOps, ParseArgs, PrintLogger, Watched}
 import pprint.{Renderer, Truncated}
 import upickle.Js
 object MainModule{
@@ -30,8 +30,30 @@ trait MainModule extends mill.Module{
   implicit def millScoptEvaluatorReads[T] = new mill.main.EvaluatorScopt[T]()
 
   def resolve(evaluator: Evaluator[Any], targets: String*) = mill.T.command{
-    MainModule.resolveTasks(evaluator, targets, multiSelect = true){ tasks =>
-      tasks.foreach(println)
+    val resolved = for {
+      parsed <- ParseArgs(targets, multiSelect = true)
+      (selectors, args) = parsed
+      taskss <- {
+        val selected = selectors.map { case (scopedSel, sel) =>
+          val (rootModule, discover, crossSelectors) =
+            mill.main.RunScript.prepareResolve(evaluator, scopedSel, sel)
+
+          try {
+            mill.eval.Evaluator.currentEvaluator.set(evaluator)
+            mill.main.ResolveMetadata.resolve(
+              sel.value.toList, rootModule, discover,
+              args, crossSelectors.toList, Nil
+            )
+          } finally{
+            mill.eval.Evaluator.currentEvaluator.set(null)
+          }
+        }
+        EitherOps.sequence(selected)
+      }
+    } yield taskss.flatten
+    resolved match{
+      case Left(err) => throw new Exception(err)
+      case Right(r) => r.foreach(println)
     }
   }
 
diff --git a/main/src/mill/main/Resolve.scala b/main/src/mill/main/Resolve.scala
index c134d70d..507b06a3 100644
--- a/main/src/mill/main/Resolve.scala
+++ b/main/src/mill/main/Resolve.scala
@@ -4,137 +4,171 @@ import mill.define._
 import mill.define.TaskModule
 import ammonite.util.Res
 import mill.util.Router.EntryPoint
-import mill.util.Scripts
+import mill.util.{Router, Scripts}
+
+import scala.reflect.ClassTag
+
+object ResolveMetadata extends Resolve[String]{
+  def singleModuleMeta(obj: Module, discover: Discover[_], isRootModule: Boolean) = {
+    val modules = obj.millModuleDirectChildren.map(_.toString)
+    val targets =
+      obj
+        .millInternal
+        .reflect[Target[_]]
+        .map(_.toString)
+    val commands = for{
+      (cls, entryPoints) <- discover.value
+      if cls.isAssignableFrom(obj.getClass)
+      ep <- entryPoints
+    } yield {
+      if (isRootModule) ep._2.name
+      else obj + "." + ep._2.name
+    }
+
+    modules ++ targets ++ commands
+  }
+  def endResolve(obj: Module,
+                 revSelectorsSoFar: List[Segment],
+                 last: String,
+                 discover: Discover[_],
+                 rest: Seq[String]): Either[String, List[String]] = {
+
+    val direct = singleModuleMeta(obj, discover, revSelectorsSoFar.isEmpty)
+    if (last == "__") {
+      Right(direct.toList ++ obj.millInternal.modules.flatMap(singleModuleMeta(_, discover, false)))
+    } else if (last == "_") Right(direct.toList)
+    else direct.find(_.split('.').last == last) match{
+      case None =>
+        Left(
+          "Unable to resolve " +
+            Segments((Segment.Label(last) :: revSelectorsSoFar).reverse: _*).render
+        )
+      case Some(s) => Right(List(s))
+    }
+  }
+}
+object Resolve extends Resolve[NamedTask[Any]]{
+  def endResolve(obj: Module,
+                 revSelectorsSoFar: List[Segment],
+                 last: String,
+                 discover: Discover[_],
+                 rest: Seq[String]) = {
+    val target =
+      obj
+        .millInternal
+        .reflect[Target[_]]
+        .find(_.label == last)
+        .map(Right(_))
+
+    def shimArgsig[T](a: Router.ArgSig[T, _]) = {
+      ammonite.main.Router.ArgSig[T](
+        a.name,
+        a.typeString,
+        a.doc,
+        a.default
+      )
+    }
+
+    def invokeCommand(target: Module, name: String) = for {
+      (cls, entryPoints) <- discover.value
+      if cls.isAssignableFrom(target.getClass)
+      ep <- entryPoints
+      if ep._2.name == name
+    } yield Scripts.runMainMethod(
+      target,
+      ep._2.asInstanceOf[EntryPoint[Module]],
+      ammonite.main.Scripts.groupArgs(rest.toList)
+    ) match {
+      case Res.Success(v: Command[_]) => Right(v)
+      case Res.Failure(msg) => Left(msg)
+      case Res.Exception(ex, msg) =>
+        val sw = new java.io.StringWriter()
+        ex.printStackTrace(new java.io.PrintWriter(sw))
+        val prefix = if (msg.nonEmpty) msg + "\n" else msg
+        Left(prefix + sw.toString)
+
+    }
+
+    val runDefault = for {
+      child <- obj.millInternal.reflectNestedObjects[Module]
+      if child.millOuterCtx.segment == Segment.Label(last)
+      res <- child match {
+        case taskMod: TaskModule => Some(invokeCommand(child, taskMod.defaultCommandName()).headOption)
+        case _ => None
+      }
+    } yield res
+
+    val command = invokeCommand(obj, last).headOption
+
+    command orElse target orElse runDefault.flatten.headOption match {
+      case None => Left("Cannot resolve task " +
+        Segments((Segment.Label(last) :: revSelectorsSoFar).reverse: _*).render
+      )
+      // Contents of `either` *must* be a `Task`, because we only select
+      // methods returning `Task` in the discovery process
+      case Some(either) => either.right.map(Seq(_))
+    }
+  }
 
-object Resolve {
-  def resolve[T, V](remainingSelector: List[Segment],
-                    obj: mill.Module,
-                    discover: Discover[_],
-                    rest: Seq[String],
-                    remainingCrossSelectors: List[List[String]],
-                    revSelectorsSoFar: List[Segment]): Either[String, Seq[NamedTask[Any]]] = {
+}
+abstract class Resolve[R: ClassTag] {
+  def endResolve(obj: Module,
+                 revSelectorsSoFar: List[Segment],
+                 last: String,
+                 discover: Discover[_],
+                 rest: Seq[String]): Either[String, Seq[R]]
+
+  def resolve(remainingSelector: List[Segment],
+              obj: mill.Module,
+              discover: Discover[_],
+              rest: Seq[String],
+              remainingCrossSelectors: List[List[String]],
+              revSelectorsSoFar: List[Segment]): Either[String, Seq[R]] = {
 
     remainingSelector match{
       case Segment.Cross(_) :: Nil => Left("Selector cannot start with a [cross] segment")
       case Segment.Label(last) :: Nil =>
-        val target =
-          obj
-            .millInternal
-            .reflect[Target[_]]
-            .find(_.label == last)
-            .map(Right(_))
-
-        def shimArgsig[T](a: mill.util.Router.ArgSig[T, _]) = {
-          ammonite.main.Router.ArgSig[T](
-            a.name,
-            a.typeString,
-            a.doc,
-            a.default
-          )
-        }
-        def invokeCommand(target: mill.Module, name: String) = for{
-          (cls, entryPoints) <- discover.value
-          if cls.isAssignableFrom(target.getClass)
-          ep <- entryPoints
-          if ep._2.name == name
-        } yield Scripts.runMainMethod(
-          target,
-          ep._2.asInstanceOf[EntryPoint[mill.Module]],
-          ammonite.main.Scripts.groupArgs(rest.toList)
-        ) match{
-          case Res.Success(v: Command[_]) => Right(v)
-          case Res.Failure(msg) => Left(msg)
-          case Res.Exception(ex, msg) =>
-            val sw = new java.io.StringWriter()
-            ex.printStackTrace(new java.io.PrintWriter(sw))
-            val prefix = if (msg.nonEmpty) msg + "\n" else msg
-            Left(prefix + sw.toString)
-
-        }
-
-        val runDefault = for{
-          child <- obj.millInternal.reflectNestedObjects[mill.Module]
-          if child.millOuterCtx.segment == Segment.Label(last)
-          res <- child match{
-            case taskMod: TaskModule => Some(invokeCommand(child, taskMod.defaultCommandName()).headOption)
-            case _ => None
-          }
-        } yield res
-
-        val command = invokeCommand(obj, last).headOption
-
-        command orElse target orElse runDefault.flatten.headOption match{
-          case None =>  Left("Cannot resolve task " +
-            Segments((Segment.Label(last) :: revSelectorsSoFar).reverse:_*).render
-          )
-          // Contents of `either` *must* be a `Task`, because we only select
-          // methods returning `Task` in the discovery process
-          case Some(either) => either.right.map(Seq(_))
-        }
+        endResolve(obj, revSelectorsSoFar, last, discover, rest)
 
 
       case head :: tail =>
         val newRevSelectorsSoFar = head :: revSelectorsSoFar
+        def resolveFailureMsg = Left(
+          "Cannot resolve module " + Segments(newRevSelectorsSoFar.reverse:_*).render
+        )
+        def recurse(searchModules: Seq[Module]) = {
+          val matching = searchModules
+            .map(resolve(tail, _, discover, rest, remainingCrossSelectors, newRevSelectorsSoFar))
+            .collect{case Right(vs) => vs}.flatten
+
+          if (matching.nonEmpty)Right(matching)
+          else resolveFailureMsg
+        }
         head match{
           case Segment.Label(singleLabel) =>
-            if (singleLabel == "__") {
-
-              val matching =
-                obj.millInternal.modules
-                  .map(resolve(tail, _, discover, rest, remainingCrossSelectors, newRevSelectorsSoFar))
-                  .collect{case Right(vs) => vs}.flatten
-
-              if (matching.nonEmpty)Right(matching)
-              else Left("Cannot resolve module " + Segments(newRevSelectorsSoFar.reverse:_*).render)
-            } else if (singleLabel == "_") {
-
-              val matching =
-                obj.millModuleDirectChildren
-                  .map(resolve(tail, _, discover, rest, remainingCrossSelectors, newRevSelectorsSoFar))
-                  .collect{case Right(vs) => vs}.flatten
-
-              if (matching.nonEmpty)Right(matching)
-              else Left("Cannot resolve module " + Segments(newRevSelectorsSoFar.reverse:_*).render)
-            }else{
-
-              obj.millInternal.reflectNestedObjects[mill.Module].find{
-                _.millOuterCtx.segment == Segment.Label(singleLabel)
-              } match{
-                case Some(child: mill.Module) => resolve(tail, child, discover, rest, remainingCrossSelectors, newRevSelectorsSoFar)
-                case None => Left("Cannot resolve module " + Segments(newRevSelectorsSoFar.reverse:_*).render)
+            recurse(
+              if (singleLabel == "__") obj.millInternal.modules
+              else if (singleLabel == "_") obj.millModuleDirectChildren.toSeq
+              else{
+                obj.millInternal.reflectNestedObjects[mill.Module]
+                  .find(_.millOuterCtx.segment == Segment.Label(singleLabel))
+                  .toSeq
               }
-            }
-
+            )
           case Segment.Cross(cross) =>
             obj match{
-              case c: Cross[_] =>
-                if(cross == Seq("__")){
-                  val matching =
-                    for ((k, v) <- c.items)
-                    yield resolve(tail, v.asInstanceOf[mill.Module], discover, rest, remainingCrossSelectors, newRevSelectorsSoFar)
-
-                  val results = matching.collect{case Right(res) => res}.flatten
-
-                  if (results.isEmpty) Left("Cannot resolve cross " + Segments(newRevSelectorsSoFar.reverse:_*).render)
-                  else Right(results)
-                } else if (cross.contains("_")){
-                  val matching = for {
-                    (k, v) <- c.items
-                    if k.length == cross.length
-                    if k.zip(cross).forall { case (l, r) => l == r || r == "_" }
-                  } yield resolve(tail, v.asInstanceOf[mill.Module], discover, rest, remainingCrossSelectors, newRevSelectorsSoFar)
-
-                  val results = matching.collect{case Right(res) => res}.flatten
-
-                  if (results.isEmpty) Left("Cannot resolve cross " + Segments(newRevSelectorsSoFar.reverse:_*).render)
-                  else Right(results)
-                }else{
-                  c.itemMap.get(cross.toList) match{
-                    case Some(m: mill.Module) => resolve(tail, m, discover, rest, remainingCrossSelectors, newRevSelectorsSoFar)
-                    case None => Left("Cannot resolve cross " + Segments(newRevSelectorsSoFar.reverse:_*).render)
-                  }
-                }
-              case _ => Left("Cannot resolve cross " + Segments(newRevSelectorsSoFar.reverse:_*).render)
+              case c: Cross[Module] =>
+                recurse(
+                  if(cross == Seq("__")) for ((k, v) <- c.items) yield v
+                  else if (cross.contains("_")){
+                      for {
+                        (k, v) <- c.items
+                        if k.length == cross.length
+                        if k.zip(cross).forall { case (l, r) => l == r || r == "_" }
+                      } yield v
+                  }else c.itemMap.get(cross.toList).toSeq
+                )
+              case _ => resolveFailureMsg
             }
         }
 
diff --git a/main/src/mill/main/RunScript.scala b/main/src/mill/main/RunScript.scala
index 4ee16ac3..c9155551 100644
--- a/main/src/mill/main/RunScript.scala
+++ b/main/src/mill/main/RunScript.scala
@@ -8,7 +8,7 @@ import ammonite.runtime.SpecialClassLoader
 import ammonite.util.Util.CodeSource
 import ammonite.util.{Name, Res, Util}
 import mill.define
-import mill.define.{Discover, ExternalModule, Segment, Task}
+import mill.define._
 import mill.eval.{Evaluator, PathRef, Result}
 import mill.util.{EitherOps, Logger, ParseArgs, Watched}
 import mill.util.Strict.Agg
@@ -137,19 +137,8 @@ object RunScript{
       (selectors, args) = parsed
       taskss <- {
         val selected = selectors.map { case (scopedSel, sel) =>
-          val (rootModule, discover) = scopedSel match{
-            case None => (evaluator.rootModule, evaluator.discover)
-            case Some(scoping) =>
-              val moduleCls =
-                evaluator.rootModule.getClass.getClassLoader.loadClass(scoping.render + "$")
-
-              val rootModule = moduleCls.getField("MODULE$").get(moduleCls).asInstanceOf[ExternalModule]
-              (rootModule, rootModule.millDiscover)
-          }
-          val crossSelectors = sel.value.map {
-            case Segment.Cross(x) => x.toList.map(_.toString)
-            case _ => Nil
-          }
+          val (rootModule, discover, crossSelectors) =
+            prepareResolve(evaluator, scopedSel, sel)
 
           try {
             // We inject the `evaluator.rootModule` into the TargetScopt, rather
@@ -159,8 +148,7 @@ object RunScript{
             // is not currently supported
             mill.eval.Evaluator.currentEvaluator.set(evaluator)
             mill.main.Resolve.resolve(
-              sel.value.toList, rootModule,
-              discover,
+              sel.value.toList, rootModule, discover,
               args, crossSelectors.toList, Nil
             )
           } finally{
@@ -172,6 +160,23 @@ object RunScript{
     } yield taskss.flatten
   }
 
+  def prepareResolve[T](evaluator: Evaluator[T], scopedSel: Option[Segments], sel: Segments) = {
+    val (rootModule, discover) = scopedSel match {
+      case None => (evaluator.rootModule, evaluator.discover)
+      case Some(scoping) =>
+        val moduleCls =
+          evaluator.rootModule.getClass.getClassLoader.loadClass(scoping.render + "$")
+
+        val rootModule = moduleCls.getField("MODULE$").get(moduleCls).asInstanceOf[ExternalModule]
+        (rootModule, rootModule.millDiscover)
+    }
+    val crossSelectors = sel.value.map {
+      case Segment.Cross(x) => x.toList.map(_.toString)
+      case _ => Nil
+    }
+    (rootModule, discover, crossSelectors)
+  }
+
   def evaluateTasks[T](evaluator: Evaluator[T],
                        scriptArgs: Seq[String],
                        multiSelect: Boolean) = {