2 files changed, 15 insertions, 1 deletions
diff --git a/scalalib/src/mill/scalalib/ScalaModule.scala b/scalalib/src/mill/scalalib/ScalaModule.scala
index 390672c6..05c34249 100644
--- a/scalalib/src/mill/scalalib/ScalaModule.scala
+++ b/scalalib/src/mill/scalalib/ScalaModule.scala
@@ -23,7 +23,19 @@ trait ScalaModule extends mill.Module with TaskModule { outer =>
     override def moduleDeps = Seq(outer)
   }
   def scalaVersion: T[String]
-  def mainClass: T[Option[String]] = None
+
+  def mainClass: T[Option[String]] = T{
+    discoverMainClasses() match {
+      case Seq(main) => Some(main)
+      case _         => None
+    }
+  }
+
+  def discoverMainClasses: T[Seq[String]] = T{
+    Task.traverse(transitiveModuleDeps){ module => T.task {
+      mill.scalalib.ScalaWorkerApi.scalaWorker().discoverMainClasses(module.compile())
+    }}().flatten.distinct
+  }
 
   def ivyDeps = T{ Agg.empty[Dep] }
   def compileIvyDeps = T{ Agg.empty[Dep] }
diff --git a/scalalib/src/mill/scalalib/ScalaWorkerApi.scala b/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
index 03cd7d9d..ec2489af 100644
--- a/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
+++ b/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
@@ -62,4 +62,6 @@ trait ScalaWorkerApi {
                testClassfilePath: Agg[Path],
                args: Seq[String])
               (implicit ctx: mill.util.Ctx.Log): (String, Seq[Result])
+
+  def discoverMainClasses(compilationResult: CompilationResult)(implicit ctx: mill.util.Ctx): Seq[String]
 }