1 files changed, 22 insertions, 12 deletions
diff --git a/scalalib/src/mill/scalalib/JavaModule.scala b/scalalib/src/mill/scalalib/JavaModule.scala
index 259e5cd3..ab4f8795 100644
--- a/scalalib/src/mill/scalalib/JavaModule.scala
+++ b/scalalib/src/mill/scalalib/JavaModule.scala
@@ -17,12 +17,13 @@ import mill.util.Loose.Agg
   * Core configuration required to compile a single Scala compilation target
   */
 trait JavaModule extends mill.Module with TaskModule { outer =>
-  def scalaWorker: ScalaWorkerModule = mill.scalalib.ScalaWorkerModule
+  def zincWorker: ZincWorkerModule = mill.scalalib.ZincWorkerModule
 
   trait Tests extends TestModule{
     override def moduleDeps = Seq(outer)
     override def repositories = outer.repositories
     override def javacOptions = outer.javacOptions
+    override def zincWorker = outer.zincWorker
   }
   def defaultCommandName() = "run"
 
@@ -38,7 +39,16 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
   def finalMainClassOpt: T[Either[String, String]] = T{
     mainClass() match{
       case Some(m) => Right(m)
-      case None => Left("No main class specified or found")
+      case None =>
+        zincWorker.worker().discoverMainClasses(compile())match {
+          case Seq() => Left("No main class specified or found")
+          case Seq(main) => Right(main)
+          case mains =>
+            Left(
+              s"Multiple main classes found (${mains.mkString(",")}) " +
+                "please explicitly specify which one to use by overriding mainClass"
+            )
+        }
     }
   }
 
@@ -98,7 +108,7 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
   }
 
 
-  def repositories: Seq[Repository] = scalaWorker.repositories
+  def repositories: Seq[Repository] = zincWorker.repositories
 
   def platformSuffix = T{ "" }
 
@@ -133,12 +143,12 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
     } yield PathRef(path)
   }
 
-  def compile: T[CompilationResult] = T{
-    Lib.compileJava(
-      allSourceFiles().map(_.path.toIO).toArray,
-      compileClasspath().map(_.path.toIO).toArray,
-      javacOptions(),
-      upstreamCompileOutput()
+  def compile: T[CompilationResult] = T.persistent{
+    zincWorker.worker().compileJava(
+      upstreamCompileOutput(),
+      allSourceFiles().map(_.path),
+      compileClasspath().map(_.path),
+      javacOptions()
     )
   }
 
@@ -317,7 +327,7 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
     val (procId, procTombstone, token) = backgroundSetup(T.ctx().dest)
     try Result.Success(Jvm.interactiveSubprocess(
       "mill.scalalib.backgroundwrapper.BackgroundWrapper",
-      (runClasspath() ++ scalaWorker.backgroundWrapperClasspath()).map(_.path),
+      (runClasspath() ++ zincWorker.backgroundWrapperClasspath()).map(_.path),
       forkArgs(),
       forkEnv(),
       Seq(procId.toString, procTombstone.toString, token, finalMainClass()) ++ args,
@@ -332,7 +342,7 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
     val (procId, procTombstone, token) = backgroundSetup(T.ctx().dest)
     try Result.Success(Jvm.interactiveSubprocess(
       "mill.scalalib.backgroundwrapper.BackgroundWrapper",
-      (runClasspath() ++ scalaWorker.backgroundWrapperClasspath()).map(_.path),
+      (runClasspath() ++ zincWorker.backgroundWrapperClasspath()).map(_.path),
       forkArgs(),
       forkEnv(),
       Seq(procId.toString, procTombstone.toString, token, mainClass) ++ args,
@@ -384,7 +394,7 @@ trait TestModule extends JavaModule with TaskModule {
 
     Jvm.subprocess(
       mainClass = "mill.scalalib.TestRunner",
-      classPath = scalaWorker.scalalibClasspath().map(_.path),
+      classPath = zincWorker.scalalibClasspath().map(_.path),
       jvmArgs = forkArgs(),
       envArgs = forkEnv(),
       mainArgs =