Move finding Main classes to external sbt run

Author: BarkingBad

compiler/src/dotty/tools/backend/jvm/GenBCode.scala

@@ -61,35 +61,26 @@ class GenBCode extends Phase {
 
 
   override def runOn(units: List[CompilationUnit])(using Context): List[CompilationUnit] = {
+    outputDir match
+      case jar: JarArchive =>
+        updateJarManifestWithMainClass(jar)
+      case _ =>
     try super.runOn(units)
-    finally myOutput match {
+    finally outputDir match {
       case jar: JarArchive =>
         if (ctx.run.suspendedUnits.nonEmpty)
           // If we close the jar the next run will not be able to write on the jar.
           // But if we do not close it we cannot use it as part of the macro classpath of the suspended files.
           report.error("Can not suspend and output to a jar at the same time. See suspension with -Xprint-suspension.")
 
-        updateJarManifestWithMainClass(units, jar)
         jar.close()
       case _ =>
     }
   }
 
-  private def updateJarManifestWithMainClass(units: List[CompilationUnit], jarArchive: JarArchive)(using Context): Unit =
+  private def updateJarManifestWithMainClass(jarArchive: JarArchive)(using Context): Unit =
     val mainClass = Option.when(!ctx.settings.XmainClass.isDefault)(ctx.settings.XmainClass.value).orElse {
-      val _mainClassesBuffer = new mutable.HashSet[String]
-      units.map { unit =>
-        unit.tpdTree.foreachSubTree { tree =>
-          val sym = tree.symbol
-          import dotty.tools.dotc.core.NameOps.stripModuleClassSuffix
-          val name = sym.fullName.stripModuleClassSuffix.toString
-          if (sym.isStatic && !sym.is(Flags.Trait) && ctx.platform.hasMainMethod(sym)) {
-            // If sym is an object, all main methods count, otherwise only @static ones count.
-            _mainClassesBuffer += name
-          }
-        }
-      }
-      _mainClassesBuffer.toList.match
+      ctx.entryPoints.toList.match
         case List(mainClass) =>
           Some(mainClass)
         case Nil =>
@@ -99,7 +90,6 @@ class GenBCode extends Phase {
           report.warning(s"No Main-Class due to multiple entry points:\n  ${mcs.mkString("\n  ")}")
           None
     }
-
     mainClass.map { mc =>
       val manifest = Jar.WManifest()
       manifest.mainClass = mc

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -42,6 +42,7 @@ class Compiler {
     List(new semanticdb.ExtractSemanticDB) :: // Extract info into .semanticdb files
     List(new PostTyper) ::          // Additional checks and cleanups after type checking
     List(new sjs.PrepJSInterop) ::  // Additional checks and transformations for Scala.js (Scala.js only)
+    List(new CollectEntryPoints) :: // Collect all entry points and save them in the context
     List(new sbt.ExtractAPI) ::     // Sends a representation of the API of classes to sbt via callbacks
     List(new SetRootTree) ::        // Set the `rootTreeOrProvider` on class symbols
     Nil

compiler/src/dotty/tools/dotc/core/Contexts.scala

@@ -52,8 +52,9 @@ object Contexts {
   private val (notNullInfosLoc,     store8) = store7.newLocation[List[NotNullInfo]]()
   private val (importInfoLoc,       store9) = store8.newLocation[ImportInfo]()
   private val (typeAssignerLoc,    store10) = store9.newLocation[TypeAssigner](TypeAssigner)
+  private val (entryPointsLoc,     store11) = store10.newLocation[EntryPoints](new EntryPoints)
 
-  private val initialStore = store10
+  private val initialStore = store11
 
   /** The current context */
   inline def ctx(using ctx: Context): Context = ctx
@@ -239,6 +240,9 @@ object Contexts {
     /** The current type assigner or typer */
     def typeAssigner: TypeAssigner = store(typeAssignerLoc)
 
+    /** The current entry points */
+    def entryPoints: EntryPoints = store(entryPointsLoc)
+
     /** The new implicit references that are introduced by this scope */
     protected var implicitsCache: ContextualImplicits = null
     def implicits: ContextualImplicits = {

compiler/src/dotty/tools/dotc/core/EntryPoints.scala

@@ -0,0 +1,3 @@
+package dotty.tools.dotc.core
+
+type EntryPoints = scala.collection.mutable.HashSet[String]

compiler/src/dotty/tools/dotc/sbt/ExtractAPI.scala

@@ -66,7 +66,7 @@ class ExtractAPI extends Phase {
 
     val apiTraverser = new ExtractAPICollector
     val classes = apiTraverser.apiSource(unit.tpdTree)
-    val mainClasses = apiTraverser.mainClasses
+    val mainClasses = ctx.entryPoints.toSet
 
     if (ctx.settings.YdumpSbtInc.value) {
       // Append to existing file that should have been created by ExtractDependencies
@@ -144,7 +144,6 @@ private class ExtractAPICollector(using Context) extends ThunkHolder {
   private val refinedTypeCache = new mutable.HashMap[(api.Type, api.Definition), api.Structure]
 
   private val allNonLocalClassesInSrc = new mutable.HashSet[xsbti.api.ClassLike]
-  private val _mainClasses = new mutable.HashSet[String]
 
   private object Constants {
     val emptyStringArray = Array[String]()
@@ -195,11 +194,6 @@ private class ExtractAPICollector(using Context) extends ThunkHolder {
   def apiClass(sym: ClassSymbol): api.ClassLikeDef =
     classLikeCache.getOrElseUpdate(sym, computeClass(sym))
 
-  def mainClasses: Set[String] = {
-    forceThunks()
-    _mainClasses.toSet
-  }
-
   private def computeClass(sym: ClassSymbol): api.ClassLikeDef = {
     import xsbti.api.{DefinitionType => dt}
     val defType =
@@ -234,11 +228,6 @@ private class ExtractAPICollector(using Context) extends ThunkHolder {
 
     allNonLocalClassesInSrc += cl
 
-    if (sym.isStatic && !sym.is(Trait) && ctx.platform.hasMainMethod(sym)) {
-       // If sym is an object, all main methods count, otherwise only @static ones count.
-      _mainClasses += name
-    }
-
     api.ClassLikeDef.of(name, acc, modifiers, anns, tparams, defType)
   }
 

compiler/src/dotty/tools/dotc/transform/CollectEntryPoints.scala

@@ -0,0 +1,50 @@
+package dotty.tools.dotc.transform
+
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.Types
+import dotty.tools.dotc.transform.MegaPhase._
+import dotty.tools.dotc.ast.tpd
+import java.io.{File => _}
+
+import dotty.tools.dotc.core._
+import SymDenotations._
+import Contexts._
+import Types._
+import Symbols._
+import dotty.tools.dotc.util.SourcePosition
+import Decorators._
+import StdNames.nme
+import dotty.tools.io.JarArchive
+
+/**
+ * Small phase to be run to collect main classes and store them in the context.
+ * The general rule to run this phase is either:
+ * - The output of compilation is JarArchive and there is no `-Xmain-class` defined
+ * - The compiler is run from sbt and is forced by flags forcing `ExtractorAPI`
+ *
+ * The following flags affect this phase:
+ *   -d path.jar
+ *   -Xmain-class
+ *   -Yforce-sbt-phases
+ *   -Ydump-sbt-inc
+ */
+class CollectEntryPoints extends MiniPhase:
+  def phaseName: String = "Collect entry points"
+
+  override def isRunnable(using Context): Boolean =
+    def forceRun = (ctx.settings.XmainClass.isDefault && ctx.settings.outputDir.value.isInstanceOf[JarArchive]) ||
+                    ctx.settings.YdumpSbtInc.value ||
+                    ctx.settings.YforceSbtPhases.value
+    super.isRunnable && (ctx.sbtCallback != null || forceRun)
+
+
+  override def transformTypeDef(tree: tpd.TypeDef)(using Context): tpd.Tree =
+    ctx.entryPoints ++= getEntryPoint(tree)
+    tree
+
+  private def getEntryPoint(tree: tpd.TypeDef)(using Context): Option[String] =
+    val sym = tree.symbol
+    import dotty.tools.dotc.core.NameOps.stripModuleClassSuffix
+    val name = sym.fullName.stripModuleClassSuffix.toString
+    Option.when(sym.isStatic && !sym.is(Flags.Trait) && ctx.platform.hasMainMethod(sym))(name)

compiler/src/dotty/tools/io/Jar.scala

@@ -44,7 +44,7 @@ class Jar(file: File) {
   import Jar._
 
   lazy val jarFile: JarFile  = new JarFile(file.jpath.toFile)
-  lazy val manifest: Option[Manifest] = withJarInput(s => Option(s.getManifest).orElse(findManifest(s)))
+  lazy val manifest: Option[Manifest] = withJarInput(s => Option(s.getManifest))
 
   def mainClass: Option[String]     = manifest.map(_(Name.MAIN_CLASS))
   /** The manifest-defined classpath String if available. */
@@ -73,20 +73,6 @@ class Jar(file: File) {
     case x      => x
   }
 
-  /**
-   * Hack for Java reading MANIFEST.MF only if it is at the first entry of the JAR file and otherwise returns null.
-   * https://github.com/openjdk/jdk/blob/master/src/java.base/share/classes/java/util/jar/JarInputStream.java#L74
-   * Suprisingly, such jars still can be successfully runned jars via `java -jar path.jar` so it is only problem for Jar readers.
-   */
-  private def findManifest(s: JarInputStream): Option[Manifest] =
-    val entry = s.getNextEntry
-    if entry != null && entry.getName != JarFile.MANIFEST_NAME then
-      findManifest(s)
-    else if entry == null then
-      None
-    else
-      Some(Manifest(s))
-
   override def toString: String = "" + file
 }