Simplify handling of MillBuildRootModule in CodeGen (#5057)

* Now that we are on Scala 3, `MillBuildRootModule` can be a `trait`
like most other modules and still take implicit parameters

* That allows us to remove the special casing of `MillBuildRootModule`
in codegen, since it can now stack with the default `MainRootModule`
that non-meta-build root modules inherit from

* We can also remove the synthetic `import
_root_.mill.runner.meta.MillBuildRootModule` that we were previously
injecting into the scope, since now it is just another stackable trait
like any other

* Also tidied up `CodeGen.scala` a bit to reduce duplication

Author: lihaoyi

example/extending/metabuild/4-meta-build/mill-build/build.mill

@@ -1,5 +1,6 @@
 package build
 import mill._, scalalib._
+import mill.runner.meta.MillBuildRootModule
 
 object `package` extends MillBuildRootModule {
   val scalatagsVersion = "0.13.1"

example/extending/metabuild/5-meta-shared-sources/mill-build/build.mill

@@ -1,4 +1,5 @@
 package build
 import mill._
+import mill.runner.meta.MillBuildRootModule
 
 object `package` extends MillBuildRootModule

integration/failure/root-module-compile-error/src/RootModuleCompileErrorTests.scala

@@ -26,13 +26,13 @@ object RootModuleCompileErrorTests extends UtestIntegrationTestSuite {
 
       locally {
         // For now these error messages still show generated/mangled code; not ideal, but it'll do
-        assert(res.err.replace('\\', '/').contains("""foo/package.mill:6:85"""))
+        assert(res.err.replace('\\', '/').contains("""foo/package.mill:6:92"""))
         assert(res.err.contains("""Not found: type UnknownFooModule"""))
         assert(res.err.contains(
-          """abstract class package_  extends mill.main.SubfolderModule(build.millDiscover) with UnknownFooModule {"""
+          """abstract class package_  extends _root_.mill.main.SubfolderModule(build.millDiscover) with UnknownFooModule {"""
         ))
         assert(res.err.contains(
-          """                                                                ^^^^^^^^^^^^^^^^"""
+          """                                                                                           ^^^^^^^^^^^^^^^^"""
         ))
       }
 

integration/feature/scala-3-syntax/resources/mill-build/build.mill

@@ -3,7 +3,7 @@ package build
 import mill.*
 import mill.scalalib.*
 
-object `package` extends MillBuildRootModule:
+object `package` extends mill.runner.meta.MillBuildRootModule:
   def mvnDeps = Seq(
     mvn"org.scala-lang::toolkit:0.5.0"
   )

integration/ide/build-classpath-contents/resources/mill-build/build.mill

@@ -1,5 +1,6 @@
 package build
 import mill.*
 import mill.scalalib.*
+import mill.runner.meta.MillBuildRootModule
 
 object `package` extends MillBuildRootModule

integration/ide/gen-idea/resources/extended/mill-build/build.mill

@@ -1,5 +1,6 @@
 import mill._
 import mill.scalalib._
+import mill.runner.meta.MillBuildRootModule
 
 object `package` extends MillBuildRootModule {
   def mvnDeps = Seq(mvn"org.scalameta::munit:0.7.29")

integration/invalidation/multi-level-editing/resources/mill-build/build.mill

@@ -1,6 +1,7 @@
 package build
 
 import mill._, scalalib._
+import mill.runner.meta.MillBuildRootModule
 
 object `package` extends MillBuildRootModule {
   def mvnDeps = Seq(mvn"com.lihaoyi::scalatags:${constant.MetaConstant.scalatagsVersion}")

integration/invalidation/multi-level-editing/resources/mill-build/mill-build/build.mill

@@ -1,6 +1,7 @@
 package build
 // build.mill
 import mill._, scalalib._
+import mill.runner.meta.MillBuildRootModule
 
 object `package` extends MillBuildRootModule {
   def generatedSources = Task {

runner/meta/src/mill/runner/meta/CodeGen.scala

@@ -26,43 +26,39 @@ object CodeGen {
       val isBuildScript = specialNames(scriptPath.last)
       val scriptFolderPath = scriptPath / os.up
 
-      if (scriptFolderPath == projectRoot && scriptPath.last.split('.').head == "package") {
-        break()
-      }
+      val scriptBaseName = scriptPath.last.split('.').head
 
-      if (scriptFolderPath != projectRoot && scriptPath.last.split('.').head == "build") {
-        break()
-      }
+      if (scriptFolderPath == projectRoot && scriptBaseName == "package") break()
+      if (scriptFolderPath != projectRoot && scriptBaseName == "build") break()
 
       val packageSegments = FileImportGraph.fileImportToSegments(projectRoot, scriptPath)
       val dest = targetDest / packageSegments
 
       val childNames = scriptSources
-        .flatMap { path =>
-          if (path == scriptPath) None
-          else if (nestedBuildFileNames.contains(path.last)) {
-            Option.when(path / os.up / os.up == scriptFolderPath) {
-              (path / os.up).last
-            }
-          } else None
+        .collect {
+          case path
+              if path != scriptPath
+                && nestedBuildFileNames.contains(path.last)
+                && path / os.up / os.up == scriptFolderPath => (path / os.up).last
         }
         .distinct
 
       val pkgSegments = packageSegments.drop(1).dropRight(1)
 
       def pkgSelector0(pre: Option[String], s: Option[String]) =
         (pre ++ pkgSegments ++ s).map(backtickWrap).mkString(".")
+
       def pkgSelector2(s: Option[String]) = s"_root_.${pkgSelector0(Some(globalPackagePrefix), s)}"
-      val (childSels, childAliases0) = childNames
+
+      val childAliases = childNames
         .map { c =>
           // Dummy references to sub-modules. Just used as metadata for the discover and
           // resolve logic to traverse, cannot actually be evaluated and used
-          val comment = "// subfolder module reference"
           val lhs = backtickWrap(c)
           val rhs = s"${pkgSelector2(Some(c))}.package_"
-          (rhs, s"final lazy val $lhs: $rhs.type = $rhs $comment")
-        }.unzip
-      val childAliases = childAliases0.mkString("\n")
+          s"final lazy val $lhs: $rhs.type = $rhs // subfolder module reference"
+        }
+        .mkString("\n")
 
       val pkg = pkgSelector0(Some(globalPackagePrefix), None)
 
@@ -155,16 +151,11 @@ object CodeGen {
         scriptFolderPath,
         compilerWorkerClasspath,
         millTopLevelProjectRoot,
-        output,
-        projectRoot != millTopLevelProjectRoot
+        output
       )
 
     val objectData = parser.parseObjectData(scriptCode)
 
-    val expectedParent =
-      if (projectRoot != millTopLevelProjectRoot) "MillBuildRootModule"
-      else "_root_.mill.main.MainRootModule"
-
     val expectedModuleMsg =
       if (projectRoot != millTopLevelProjectRoot) "MillBuildRootModule" else "mill.Module"
 
@@ -184,10 +175,12 @@ object CodeGen {
          |}
          |""".stripMargin
 
+    val newParent =
+      if (segments.isEmpty) "_root_.mill.main.MainRootModule"
+      else "_root_.mill.main.SubfolderModule(build.millDiscover)"
+
     objectData.find(o => o.name.text == "`package`") match {
       case Some(objectData) =>
-        val newParent =
-          if (segments.isEmpty) expectedParent else s"mill.main.SubfolderModule(build.millDiscover)"
 
         var newScriptCode = scriptCode
         objectData.endMarker match {
@@ -198,18 +191,13 @@ object CodeGen {
         }
         objectData.finalStat match {
           case Some((leading, finalStat)) =>
+            val statLines = finalStat.text.linesWithSeparators.toSeq
             val fenced = Seq(
-              "", {
-                val statLines = finalStat.text.linesWithSeparators.toSeq
-                if statLines.sizeIs > 1 then
-                  statLines.tail.mkString
-                else
-                  finalStat.text
-              }
+              "",
+              if statLines.sizeIs > 1 then statLines.tail.mkString else finalStat.text
             ).mkString(System.lineSeparator())
             newScriptCode = finalStat.applyTo(newScriptCode, fenced)
-          case None =>
-            ()
+          case None => ()
         }
 
         newScriptCode = objectData.parent.applyTo(
@@ -219,12 +207,10 @@ object CodeGen {
               s"object `package` in ${scriptPath.relativeTo(millTopLevelProjectRoot)} " +
                 s"must extend a subclass of `$expectedModuleMsg`"
             )
-          } else if (objectData.parent.text == expectedParent) newParent
-          else newParent + " with " + objectData.parent.text
+          } else newParent + " with " + objectData.parent.text
         )
 
         newScriptCode = objectData.name.applyTo(newScriptCode, wrapperObjectName)
-
         newScriptCode = objectData.obj.applyTo(newScriptCode, "abstract class")
 
         s"""$headerCode
@@ -233,14 +219,8 @@ object CodeGen {
            |""".stripMargin
 
       case None =>
-        val extendsClause =
-          if (segments.nonEmpty) s"extends _root_.mill.main.SubfolderModule(build.millDiscover) "
-          else if (millTopLevelProjectRoot == scriptFolderPath)
-            s"extends _root_.mill.main.MainRootModule "
-          else s"extends _root_.mill.runner.meta.MillBuildRootModule() "
-
         s"""$headerCode
-           |abstract class $wrapperObjectName $extendsClause { this: $wrapperObjectName.type =>
+           |abstract class $wrapperObjectName extends $newParent { this: $wrapperObjectName.type =>
            |$markerComment
            |$scriptCode
            |}""".stripMargin
@@ -252,8 +232,7 @@ object CodeGen {
       scriptFolderPath: os.Path,
       segments: Seq[String]
   ): String = {
-    s"""object MillMiscInfo
-       |extends mill.main.SubfolderModule.Info(
+    s"""object MillMiscInfo extends mill.main.SubfolderModule.Info(
        |  os.Path(${literalize(scriptFolderPath.toString)}),
        |  _root_.scala.Seq(${segments.map(pprint.Util.literalize(_)).mkString(", ")})
        |)
@@ -272,10 +251,9 @@ object CodeGen {
       scriptFolderPath: os.Path,
       compilerWorkerClasspath: Seq[os.Path],
       millTopLevelProjectRoot: os.Path,
-      output: os.Path,
-      isMetaBuild: Boolean
+      output: os.Path
   ): String = {
-    s"""${if (isMetaBuild) "import _root_.mill.runner.meta.MillBuildRootModule" else ""}
+    s"""
        |@_root_.scala.annotation.nowarn
        |object MillMiscInfo extends mill.define.RootModule0.Info(
        |  ${compilerWorkerClasspath.map(p => literalize(p.toString))},

runner/meta/src/mill/runner/meta/MillBuildRootModule.scala

@@ -22,10 +22,10 @@ import scala.jdk.CollectionConverters.ListHasAsScala
  * calls within the scripts.
  */
 @internal
-class MillBuildRootModule()(implicit
+trait MillBuildRootModule()(implicit
     rootModuleInfo: RootModule0.Info,
     scalaCompilerResolver: ScalaCompilerWorker.Resolver
-) extends mill.main.MainRootModule() with ScalaModule {
+) extends ScalaModule {
   override def bspDisplayName0: String = rootModuleInfo
     .projectRoot
     .relativeTo(rootModuleInfo.topLevelProjectRoot)
@@ -296,7 +296,7 @@ object MillBuildRootModule {
   class BootstrapModule()(implicit
       rootModuleInfo: RootModule0.Info,
       scalaCompilerResolver: ScalaCompilerWorker.Resolver
-  ) extends MillBuildRootModule() {
+  ) extends mill.main.MainRootModule() with MillBuildRootModule() {
     override lazy val millDiscover = Discover[this.type]
   }