CodeGen - package_ not always in discover map

com-lihaoyi · Sep 20, 2024 · 6366cc7 · 6366cc7
1 parent 4f27142
commit 6366cc7
Showing 1 changed file with 16 additions and 15 deletions.
diff --git a/runner/src/mill/runner/CodeGen.scala b/runner/src/mill/runner/CodeGen.scala
@@ -231,16 +231,10 @@ object CodeGen {
     def addChildren(initial: String) =
       if childSels.nonEmpty then
         s"""{
-        |      val childModules: Seq[_root_.mill.define.Module] = Seq(
-        |        ${childSels.mkString(",\n      ")}
+        |      val childDiscovers: Seq[_root_.mill.define.Discover] = Seq(
+        |        ${childSels.map(child => s"$child.millDiscover").mkString(",\n      ")}
         |      )
-        |      childModules.foldLeft($initial.value)({ (acc, m) =>
-        |        m match {
-        |          case m: _root_.mill.main.RootModule => acc ++ m.millDiscover.value
-        |          case m: _root_.mill.main.RootModule.Subfolder => acc ++ m.millDiscover.value
-        |          case _ => acc
-        |        }
-        |      })
+        |      childDiscovers.foldLeft($initial.value)(_ ++ _.value)
         |    }""".stripMargin
       else
         s"""$initial.value""".stripMargin
@@ -254,15 +248,22 @@ object CodeGen {
     s"""object ${wrapperObjectName} extends $wrapperObjectName {
        |  ${childAliases.linesWithSeparators.mkString("  ")}
        |  override lazy val millDiscover: _root_.mill.define.Discover = {
-       |    val initial = ${addChildren("this.innerMillDiscover_")}
-       |    val inner = initial(classOf[$wrapperObjectName])
-       |    _root_.mill.define.Discover.apply2(
-       |      value = initial.updated(classOf[$wrapperObjectName.type], inner),
-       |    )
+       |    val base = this.__innerMillDiscover
+       |    val initial = ${addChildren("base")}
+       |    val subbed = {
+       |      initial.get(classOf[$wrapperObjectName]) match {
+       |        case Some(inner) => initial.updated(classOf[$wrapperObjectName.type], inner)
+       |        case None => initial
+       |      }
+       |    }
+       |    if subbed ne base.value then
+       |      _root_.mill.define.Discover.apply2(value = subbed)
+       |    else
+       |      base
        |  }
        |}
        |abstract class $wrapperObjectName $extendsClause {
-       |protected def innerMillDiscover_: _root_.mill.define.Discover = _root_.mill.define.Discover[this.type]""".stripMargin
+       |protected def __innerMillDiscover: _root_.mill.define.Discover = _root_.mill.define.Discover[this.type]""".stripMargin
 
   }