5 files changed, 19 insertions, 34 deletions

diff --git a/core/src/main/scala/mill/define/Cross.scala b/core/src/main/scala/mill/define/Cross.scala
index 6120d5df..2975f97b 100644
--- a/core/src/main/scala/mill/define/Cross.scala
+++ b/core/src/main/scala/mill/define/Cross.scala
@@ -44,15 +44,9 @@ object Cross{
   */
 class Cross[T](cases: Any*)
               (implicit ci: Cross.Factory[T],
-               ctx: mill.define.Ctx,
-               cmds: Module.Cmds) extends mill.define.Module()(ctx, cmds) {
-
-  override def traverse[T](f: Module => Seq[T]): Seq[T] = {
-    f(this) ++
-    this.reflect[Module].flatMap(f) ++
-    items.map(_._2).flatMap{case t: Module => f(t)}
-  }
+               ctx: mill.define.Ctx) extends mill.define.Module()(ctx) {
 
+  override lazy val modules = this.reflectNestedObjects[Module] ++ items.collect{case (k, v: mill.define.Module) => v}
   val items = for(c0 <- cases.toList) yield{
     val c = c0 match{
       case p: Product => p
diff --git a/core/src/main/scala/mill/define/Discover.scala b/core/src/main/scala/mill/define/Discover.scala
index a5bcd8c6..0db2d617 100644
--- a/core/src/main/scala/mill/define/Discover.scala
+++ b/core/src/main/scala/mill/define/Discover.scala
@@ -13,7 +13,6 @@ object Discover {
     import c.universe._
     val seen = mutable.Set.empty[Type]
     def rec(tpe: Type): Unit = {
-      println("Rec! " + tpe)
       if (!seen(tpe)){
         seen.add(tpe)
         for{
diff --git a/core/src/main/scala/mill/define/Module.scala b/core/src/main/scala/mill/define/Module.scala
index c9ce523c..4634b9a3 100644
--- a/core/src/main/scala/mill/define/Module.scala
+++ b/core/src/main/scala/mill/define/Module.scala
@@ -8,36 +8,23 @@ import ammonite.ops.Path
 import scala.language.experimental.macros
 import scala.reflect.ClassTag
 import scala.reflect.macros.blackbox
-object Module{
-  case class Cmds(value: Seq[EntryPoint[Module]])
-  object Cmds{
-    implicit def make: Cmds = macro makeImpl
-    implicit def makeImpl(c: blackbox.Context): c.Expr[Cmds] = {
-      import c.universe._
-      reify(Cmds(Nil))
-    }
-  }
-}
 /**
   * `Module` is a class meant to be extended by `trait`s *only*, in order to
   * propagate the implicit parameters forward to the final concrete
   * instantiation site so they can capture the enclosing/line information of
   * the concrete instance.
   */
-class Module(implicit ctx0: mill.define.Ctx, cmds: Module.Cmds) extends mill.moduledefs.Cacher{
+class Module(implicit ctx0: mill.define.Ctx) extends mill.moduledefs.Cacher{
 
   def traverse[T](f: Module => Seq[T]): Seq[T] = {
-    def rec(m: Module): Seq[T] = {
-      f(m) ++
-      this.reflect[Module].flatMap(f)
-    }
+    def rec(m: Module): Seq[T] = f(m) ++ m.modules.flatMap(rec)
     rec(this)
   }
 
   lazy val segmentsToModules = traverse{m => Seq(m.ctx.segments -> m)}
     .toMap
 
-  lazy val modules = segmentsToModules.valuesIterator.toSet
+  lazy val modules = this.reflectNestedObjects[Module]
   lazy val segmentsToTargets = traverse{_.reflect[Target[_]]}
     .map(t => (t.ctx.segments, t))
     .toMap
@@ -78,6 +65,9 @@ class Module(implicit ctx0: mill.define.Ctx, cmds: Module.Cmds) extends mill.mod
       .filter(implicitly[ClassTag[T]].runtimeClass isAssignableFrom _.getReturnType)
       .map(_.getName)
   }
+  // For some reason, this fails to pick up concrete `object`s nested directly within
+  // another top-level concrete `object`. This is fine for now, since Mill's Ammonite
+  // script/REPL runner always wraps user code in a wrapper object/trait
   def reflectNestedObjects[T: ClassTag] = {
     reflect[T] ++
     this
@@ -95,9 +85,7 @@ class BaseModule(basePath: Path)
                 (implicit millModuleEnclosing0: sourcecode.Enclosing,
                  millModuleLine0: sourcecode.Line,
                  millName0: sourcecode.Name,
-                 overrides0: Overrides,
-                 cmds: Module.Cmds)
+                 overrides0: Overrides)
   extends Module()(
-    mill.define.Ctx.make(implicitly, implicitly, implicitly, BasePath(basePath), Segments(), implicitly),
-    cmds
+    mill.define.Ctx.make(implicitly, implicitly, implicitly, BasePath(basePath), Segments(), implicitly)
   )
\ No newline at end of file
diff --git a/core/src/main/scala/mill/main/MainRunner.scala b/core/src/main/scala/mill/main/MainRunner.scala
index fd16161f..5dd451c2 100644
--- a/core/src/main/scala/mill/main/MainRunner.scala
+++ b/core/src/main/scala/mill/main/MainRunner.scala
@@ -99,6 +99,9 @@ class MainRunner(config: ammonite.main.Cli.Config,
          |  def $$main() = Iterator[String]()
          |
          |  val millDiscover = mill.define.Discover[$wrapName]
+         |  // 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{
diff --git a/core/src/main/scala/mill/main/RunScript.scala b/core/src/main/scala/mill/main/RunScript.scala
index 8ad144d3..010f134b 100644
--- a/core/src/main/scala/mill/main/RunScript.scala
+++ b/core/src/main/scala/mill/main/RunScript.scala
@@ -100,9 +100,10 @@ object RunScript{
       module <- try {
         Util.withContextClassloader(interp.evalClassloader) {
           Res.Success(
-            buildCls.getField("MODULE$")
-              .get(buildCls)
-              .asInstanceOf[mill.Module]
+            buildCls.getMethod("millSelf")
+                    .invoke(null)
+                    .asInstanceOf[Some[mill.Module]]
+                    .get
           )
         }
       } catch {
@@ -112,8 +113,8 @@ object RunScript{
         Util.withContextClassloader(interp.evalClassloader) {
           Res.Success(
             buildCls.getMethod("millDiscover")
-                .invoke(null)
-                .asInstanceOf[Discover]
+                    .invoke(null)
+                    .asInstanceOf[Discover]
           )
         }
       } catch {