Reduce size of generated code in Discover macro (#4463)

lihaoyi · web-flow · commit 31a8b3c45f6f · 2025-02-04T08:57:16.000+08:00
Fixes #3893 - Recurse over all module parent classes, in addition to member return types - Remove `allNames`, since it can now be derived from the `declaredNames` of every module class - Make `entryPoints` only store the entrypoints specifically declared for each class, and not those inherited Reduces the size of `out/mill-build/compile.dest/classes//build_/package_$.class` in the Mill repo from 820K to 73K. Haven't measured but I would expect a decrease in compile times as well
diff --git a/integration/feature/inspect/src/InspectTests.scala b/integration/feature/inspect/src/InspectTests.scala
@@ -63,7 +63,8 @@ object InspectTests extends UtestIntegrationTestSuite {
         doc
       )
 
-      assert(eval(("inspect", "core.run")).isSuccess)
+      val res2 = eval(("inspect", "core.run"))
+      assert(res2.isSuccess)
       val run = out("inspect").json.str
 
       assertGlobMatches(
diff --git a/integration/feature/scala-3-syntax/resources/sub/package.mill b/integration/feature/scala-3-syntax/resources/sub/package.mill
@@ -8,7 +8,7 @@ import mill.*
 assert(1 + 1 == 2)
 
 // modifiers also allowed at top-level
-private def subCommand(): Command[Unit] = Task.Command:
+def subCommand(): Command[Unit] = Task.Command:
   println("Hello, sub-world!")
 
 // top-level object with no extends clause
diff --git a/main/define/src/mill/define/Discover.scala b/main/define/src/mill/define/Discover.scala
@@ -14,10 +14,21 @@ import scala.collection.mutable
  * the `Task.Command` methods we find. This mapping from `Class[_]` to `MainData`
  * can then be used later to look up the `MainData` for any module.
  */
-class Discover(val classInfo: Map[Class[_], Discover.Node], val allNames: Seq[String])
+class Discover(val classInfo: Map[Class[_], Discover.ClassInfo]) {
+  def resolveEntrypoint(cls: Class[_], name: String) = {
+    val res = for {
+      (cls2, node) <- classInfo
+      if cls2.isAssignableFrom(cls)
+      ep <- node.entryPoints
+      if ep.name == name
+    } yield ep
+
+    res.headOption
+  }
+}
 
 object Discover {
-  class Node(
+  class ClassInfo(
       val entryPoints: Seq[mainargs.MainData[_, _]],
       val declaredNames: Seq[String]
   )
@@ -37,14 +48,19 @@ object Discover {
       def rec(tpe: TypeRepr): Unit = {
         if (seen.add(tpe)) {
           val typeSym = tpe.typeSymbol
-          for {
+          val memberTypes: Seq[TypeRepr] = for {
             m <- typeSym.fieldMembers ++ typeSym.methodMembers
             if m != Symbol.noSymbol
-            memberTpe = m.termRef
-            if memberTpe.baseClasses.contains(moduleSym)
+          } yield m.termRef
+
+          val parentTypes: Seq[TypeRepr] = tpe.baseClasses.map(_.typeRef)
+
+          for {
+            tpe <- memberTypes ++ parentTypes
+            if tpe.baseClasses.contains(moduleSym)
           } {
-            rec(memberTpe)
-            memberTpe.asType match {
+            rec(tpe)
+            tpe.asType match {
               case '[mill.define.Cross[m]] => rec(TypeRepr.of[m])
               case _ => () // no cross argument to extract
             }
@@ -100,83 +116,63 @@ object Discover {
       // otherwise the compiler likes to give us stuff in random orders, which
       // causes the code to be generated in random order resulting in code hashes
       // changing unnecessarily
-      val mapping: Seq[(Expr[(Class[_], Node)], Seq[String])] = for {
-        discoveredModuleType <- seen.toSeq.sortBy(_.typeSymbol.fullName)
-        curCls = discoveredModuleType
-        methods = filterDefs(curCls.typeSymbol.methodMembers)
-        declMethods = filterDefs(curCls.typeSymbol.declaredMethods)
-        _ = {
+      val mapping: Seq[(TypeRepr, (Seq[scala.quoted.Expr[mainargs.MainData[?, ?]]], Seq[String]))] =
+        for {
+          curCls <- seen.toSeq.sortBy(_.typeSymbol.fullName)
+        } yield {
+          val declMethods = filterDefs(curCls.typeSymbol.declaredMethods)
           assertParamListCounts(
             curCls,
-            methods,
+            declMethods,
             (TypeRepr.of[mill.define.Command[?]], 1, "`Task.Command`"),
             (TypeRepr.of[mill.define.Target[?]], 0, "Target")
           )
-        }
 
-        names =
-          sortedMethods(curCls, sub = TypeRepr.of[mill.define.NamedTask[?]], methods).map(_.name)
-        entryPoints = for {
-          m <- sortedMethods(curCls, sub = TypeRepr.of[mill.define.Command[?]], methods)
-        } yield curCls.asType match {
-          case '[t] =>
-            val expr =
-              try
-                createMainData[Any, t](
-                  m,
-                  m.annotations.find(_.tpe =:= TypeRepr.of[mainargs.main]).getOrElse('{
-                    new mainargs.main()
-                  }.asTerm),
-                  m.paramSymss
-                ).asExprOf[mainargs.MainData[?, ?]]
-              catch {
-                case NonFatal(e) =>
-                  val (before, Array(after, _*)) = e.getStackTrace().span(e =>
-                    !(e.getClassName() == "mill.define.Discover$Router$" && e.getMethodName() == "applyImpl")
-                  ): @unchecked
-                  val trace =
-                    (before :+ after).map(_.toString).mkString("trace:\n", "\n", "\n...")
-                  report.errorAndAbort(
-                    s"Error generating maindata for ${m.fullName}: ${e}\n$trace",
-                    m.pos.getOrElse(Position.ofMacroExpansion)
-                  )
-              }
-            expr
-        }
-        declaredNames =
-          sortedMethods(
-            curCls,
-            sub = TypeRepr.of[mill.define.NamedTask[?]],
-            declMethods
-          ).map(_.name)
-        if names.nonEmpty || entryPoints.nonEmpty
-      } yield {
-        // by wrapping the `overridesRoutes` in a lambda function we kind of work around
-        // the problem of generating a *huge* macro method body that finally exceeds the
-        // JVM's maximum allowed method size
-        val overridesLambda = '{
-          def triple() =
-            new Node(${ Expr.ofList(entryPoints) }, ${ Expr(declaredNames) })
-          triple()
+          val names =
+            sortedMethods(
+              curCls,
+              sub = TypeRepr.of[mill.define.NamedTask[?]],
+              declMethods
+            ).map(_.name)
+          val entryPoints = for {
+            m <- sortedMethods(curCls, sub = TypeRepr.of[mill.define.Command[?]], declMethods)
+          } yield curCls.asType match {
+            case '[t] =>
+              val expr =
+                try
+                  createMainData[Any, t](
+                    m,
+                    m.annotations
+                      .find(_.tpe =:= TypeRepr.of[mainargs.main])
+                      .getOrElse('{ new mainargs.main() }.asTerm),
+                    m.paramSymss
+                  ).asExprOf[mainargs.MainData[?, ?]]
+                catch {
+                  case NonFatal(e) =>
+                    report.errorAndAbort(
+                      s"Error generating maindata for ${m.fullName}: ${e}\n${e.getStackTrace().mkString("\n")}",
+                      m.pos.getOrElse(Position.ofMacroExpansion)
+                    )
+                }
+              expr
+          }
+
+          (curCls.widen, (entryPoints, names))
         }
-        val lhs =
-          Ref(defn.Predef_classOf).appliedToType(discoveredModuleType.widen).asExprOf[Class[?]]
-        ('{ $lhs -> $overridesLambda }, names)
+
+      val mappingExpr = mapping.collect {
+        case (cls, (entryPoints, names)) if entryPoints.nonEmpty || names.nonEmpty =>
+          // by wrapping the `overridesRoutes` in a lambda function we kind of work around
+          // the problem of generating a *huge* macro method body that finally exceeds the
+          // JVM's maximum allowed method size
+          '{
+            def func() = new ClassInfo(${ Expr.ofList(entryPoints.toList) }, ${ Expr(names) })
+
+            (${ Ref(defn.Predef_classOf).appliedToType(cls).asExprOf[Class[?]] }, func())
+          }
       }
 
-      val expr: Expr[Discover] =
-        '{
-          // TODO: we can not import this here, so we have to import at the use site now, or redesign?
-          // import mill.main.TokenReaders.*
-          // import mill.api.JsonFormatters.*
-          new Discover(
-            Map[Class[_], Node](${ Varargs(mapping.map(_._1)) }*),
-            ${ Expr(mapping.iterator.flatMap(_._2).distinct.toList.sorted) }
-          )
-        }
-      // TODO: if needed for debugging, we can re-enable this
-      // report.warning(s"generated discovery for ${TypeRepr.of[T].show}:\n${expr.asTerm.show}", TypeRepr.of[T].typeSymbol.pos.getOrElse(Position.ofMacroExpansion))
-      expr
+      '{ new Discover(Map[Class[_], ClassInfo](${ Varargs(mappingExpr) }*)) }
     }
   }
 }
diff --git a/main/resolve/src/mill/resolve/Resolve.scala b/main/resolve/src/mill/resolve/Resolve.scala
@@ -142,7 +142,7 @@ object Resolve {
       val invoked = invokeCommand0(
         p,
         r.segments.last.value,
-        rootModule.millDiscover.asInstanceOf[Discover],
+        rootModule.millDiscover,
         args,
         nullCommandDefaults,
         allowPositionalCommandArgs
@@ -159,11 +159,8 @@ object Resolve {
       rest: Seq[String],
       nullCommandDefaults: Boolean,
       allowPositionalCommandArgs: Boolean
-  ): Iterable[Either[String, Command[_]]] = for {
-    (cls, node) <- discover.classInfo
-    if cls.isAssignableFrom(target.getClass)
-    ep <- node.entryPoints
-    if ep.name == name
+  ): Option[Either[String, Command[_]]] = for {
+    ep <- discover.resolveEntrypoint(target.getClass, name)
   } yield {
     def withNullDefault(a: mainargs.ArgSig): mainargs.ArgSig = {
       if (a.default.nonEmpty) a
@@ -303,7 +300,8 @@ trait Resolve[T] {
       ) match {
         case ResolveCore.Success(value) => Right(value)
         case ResolveCore.NotFound(segments, found, next, possibleNexts) =>
-          val allPossibleNames = rootModule.millDiscover.allNames.toSet
+          val allPossibleNames =
+            rootModule.millDiscover.classInfo.values.flatMap(_.declaredNames).toSet
           Left(ResolveNotFoundHandler(
             selector = sel,
             segments = segments,
diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
@@ -265,15 +265,12 @@ trait MainModule extends BaseModule0 {
               val mainDataOpt = evaluator
                 .rootModule
                 .millDiscover
-                .classInfo
-                .get(t.ctx.enclosingCls)
-                .flatMap(_.entryPoints.find(_.name == t.ctx.segments.last.value))
-                .headOption
+                .resolveEntrypoint(t.ctx.enclosingCls, t.ctx.segments.last.value)
 
               mainDataOpt match {
                 case Some(mainData) if mainData.renderedArgSigs.nonEmpty =>
                   val rendered = mainargs.Renderer.formatMainMethodSignature(
-                    mainDataOpt.get,
+                    mainData,
                     leftIndent = 2,
                     totalWidth = 100,
                     leftColWidth = mainargs.Renderer.getLeftColWidth(mainData.renderedArgSigs),

Original file line number	Diff line number	Diff line change
`@@ -63,7 +63,8 @@ object InspectTests extends UtestIntegrationTestSuite {`
`63`	`63`	`doc`
`64`	`64`	`)`
`65`	`65`
`66`		`- assert(eval(("inspect", "core.run")).isSuccess)`
	`66`	`+ val res2 = eval(("inspect", "core.run"))`
	`67`	`+ assert(res2.isSuccess)`
`67`	`68`	`val run = out("inspect").json.str`
`68`	`69`
`69`	`70`	`assertGlobMatches(`