improvement: Detect objects with main class in scripts

tgodzik · tgodzik · commit 4793fd8cc8da · 2025-02-12T13:04:14.000+01:00
Prebiously, if user had a legacy script with main method then it would not be picked up at all. Now, when we detect the correct signature we try to run it.

This will work in case of `def main..` and when object extends App

The possibility of false positives is pretty low, since user would have to have their own App, String or Array types. We will also only use that object if there are no toplevel statements
diff --git a/modules/build/src/main/java/scala/build/internal/JavaParserProxyMakerSubst.java b/modules/build/src/main/java/scala/build/internal/JavaParserProxyMakerSubst.java
@@ -20,6 +20,6 @@ public JavaParserProxy get(
     scala.build.Logger logger,
     Supplier<String> javaCommand
   ) {
-    return new JavaParserProxyBinary(archiveCache, logger, javaClassNameVersionOpt, javaCommand);
+    return new JavaParserProxyBinary( archiveCache, logger, javaClassNameVersionOpt, javaCommand);
   }
 }
diff --git a/modules/build/src/main/scala/scala/build/ScopedSources.scala b/modules/build/src/main/scala/scala/build/ScopedSources.scala
@@ -62,14 +62,14 @@ final case class ScopedSources(
   ): Either[BuildException, Sources] = either {
     val combinedOptions = combinedBuildOptions(scope, baseOptions)
 
-    val codeWrapper = ScriptPreprocessor.getScriptWrapper(combinedOptions)
+    val codeWrapper = ScriptPreprocessor.getScriptWrapper(combinedOptions, logger)
 
     val wrappedScripts = unwrappedScripts
       .flatMap(_.valueFor(scope).toSeq)
       .map(_.wrap(codeWrapper))
 
     codeWrapper match {
-      case _: AppCodeWrapper.type if wrappedScripts.size > 1 =>
+      case _: AppCodeWrapper if wrappedScripts.size > 1 =>
         wrappedScripts.find(_.originalPath.exists(_._1.toString == "main.sc"))
           .foreach(_ => logger.diagnostic(WarningMessages.mainScriptNameClashesWithAppWrapper))
       case _ => ()
diff --git a/modules/build/src/main/scala/scala/build/internal/AppCodeWrapper.scala b/modules/build/src/main/scala/scala/build/internal/AppCodeWrapper.scala
@@ -1,6 +1,6 @@
 package scala.build.internal
 
-case object AppCodeWrapper extends CodeWrapper {
+case class AppCodeWrapper(scalaVersion: String, log: String => Unit) extends CodeWrapper {
   override def mainClassObject(className: Name) = className
 
   def apply(
@@ -12,13 +12,19 @@ case object AppCodeWrapper extends CodeWrapper {
   ) = {
     val wrapperObjectName = indexedWrapperName.backticked
 
+    val mainObject = WrapperUtils.mainObjectInScript(scalaVersion, code)
+    val invokeMain = mainObject match
+      case WrapperUtils.ScriptMainMethod.Exists(name) => s"\n$name.main(args)"
+      case otherwise =>
+        otherwise.warningMessage.foreach(log)
+        ""
     val packageDirective =
       if (pkgName.isEmpty) "" else s"package ${AmmUtil.encodeScalaSourcePath(pkgName)}" + "\n"
     val top = AmmUtil.normalizeNewlines(
       s"""$packageDirective
          |
          |object $wrapperObjectName extends App {
-         |val scriptPath = \"\"\"$scriptPath\"\"\"
+         |val scriptPath = \"\"\"$scriptPath\"\"\"$invokeMain
          |""".stripMargin
     )
     val bottom = AmmUtil.normalizeNewlines(
diff --git a/modules/build/src/main/scala/scala/build/internal/ClassCodeWrapper.scala b/modules/build/src/main/scala/scala/build/internal/ClassCodeWrapper.scala
@@ -5,7 +5,7 @@ package scala.build.internal
   * running interconnected scripts using Scala CLI <br> <br> Incompatible with Scala 2 - it uses
   * Scala 3 feature 'export'<br> Incompatible with native JS members - the wrapper is a class
   */
-case object ClassCodeWrapper extends CodeWrapper {
+case class ClassCodeWrapper(scalaVersion: String, log: String => Unit) extends CodeWrapper {
 
   override def mainClassObject(className: Name): Name =
     Name(className.raw ++ "_sc")
@@ -16,8 +16,16 @@ case object ClassCodeWrapper extends CodeWrapper {
     extraCode: String,
     scriptPath: String
   ) = {
+
+    val mainObject = WrapperUtils.mainObjectInScript(scalaVersion, code)
+    val mainInvocation = mainObject match
+      case WrapperUtils.ScriptMainMethod.Exists(name) => s"script.$name.main(args)"
+      case otherwise =>
+        otherwise.warningMessage.foreach(log)
+        s"val _ = script.hashCode()"
+
     val name             = mainClassObject(indexedWrapperName).backticked
-    val wrapperClassName = Name(indexedWrapperName.raw ++ "$_").backticked
+    val wrapperClassName = scala.build.internal.Name(indexedWrapperName.raw ++ "$_").backticked
     val mainObjectCode =
       AmmUtil.normalizeNewlines(s"""|object $name {
                                     |  private var args$$opt0 = Option.empty[Array[String]]
@@ -33,7 +41,7 @@ case object ClassCodeWrapper extends CodeWrapper {
                                     |
                                     |  def main(args: Array[String]): Unit = {
                                     |    args$$set(args)
-                                    |    val _ = script.hashCode() // hashCode to clear scalac warning about pure expression in statement position
+                                    |    $mainInvocation // hashCode to clear scalac warning about pure expression in statement position
                                     |  }
                                     |}
                                     |
diff --git a/modules/build/src/main/scala/scala/build/internal/ObjectCodeWrapper.scala b/modules/build/src/main/scala/scala/build/internal/ObjectCodeWrapper.scala
@@ -4,7 +4,7 @@ package scala.build.internal
   * or/and not using JS native prefer [[ClassCodeWrapper]], since it prevents deadlocks when running
   * threads from script
   */
-case object ObjectCodeWrapper extends CodeWrapper {
+case class ObjectCodeWrapper(scalaVersion: String, log: String => Unit) extends CodeWrapper {
 
   override def mainClassObject(className: Name): Name =
     Name(className.raw ++ "_sc")
@@ -15,12 +15,19 @@ case object ObjectCodeWrapper extends CodeWrapper {
     extraCode: String,
     scriptPath: String
   ) = {
+    val mainObject         = WrapperUtils.mainObjectInScript(scalaVersion, code)
     val name               = mainClassObject(indexedWrapperName).backticked
     val aliasedWrapperName = name + "$$alias"
-    val funHashCodeMethod =
+    val realScript =
       if (name == "main_sc")
-        s"$aliasedWrapperName.alias.hashCode()" // https://github.com/VirtusLab/scala-cli/issues/314
-      else s"${indexedWrapperName.backticked}.hashCode()"
+        s"$aliasedWrapperName.alias" // https://github.com/VirtusLab/scala-cli/issues/314
+      else s"${indexedWrapperName.backticked}"
+
+    val funHashCodeMethod = mainObject match
+      case WrapperUtils.ScriptMainMethod.Exists(name) => s"$realScript.$name.main(args)"
+      case otherwise =>
+        otherwise.warningMessage.foreach(log)
+        s"val _ = $realScript.hashCode()"
     // We need to call hashCode (or any other method so compiler does not report a warning)
     val mainObjectCode =
       AmmUtil.normalizeNewlines(s"""|object $name {
@@ -34,7 +41,7 @@ case object ObjectCodeWrapper extends CodeWrapper {
                                     |  }
                                     |  def main(args: Array[String]): Unit = {
                                     |    args$$set(args)
-                                    |    val _ = $funHashCodeMethod // hashCode to clear scalac warning about pure expression in statement position
+                                    |    $funHashCodeMethod // hashCode to clear scalac warning about pure expression in statement position
                                     |  }
                                     |}
                                     |""".stripMargin)
diff --git a/modules/build/src/main/scala/scala/build/internal/WrapperUtils.scala b/modules/build/src/main/scala/scala/build/internal/WrapperUtils.scala
@@ -0,0 +1,80 @@
+package scala.build.internal
+
+import scala.build.internal.util.WarningMessages
+
+object WrapperUtils {
+
+  enum ScriptMainMethod:
+    case Exists(name: String)
+    case Multiple(names: Seq[String])
+    case ToplevelStatsPresent
+    case NoMain
+
+    def warningMessage: Option[String] =
+      this match
+        case ScriptMainMethod.Multiple(names) =>
+          Some(WarningMessages.multipleMainObjectsInScript(names))
+        case ScriptMainMethod.ToplevelStatsPresent => Some(
+            WarningMessages.mixedToplvelAndObjectInScript
+          )
+        case _ => None
+
+  def mainObjectInScript(scalaVersion: String, code: String): ScriptMainMethod =
+    import scala.meta.*
+
+    val scriptDialect =
+      if scalaVersion.startsWith("3") then dialects.Scala3Future else dialects.Scala213Source3
+
+    given Dialect = scriptDialect.withAllowToplevelStatements(true).withAllowToplevelTerms(true)
+    val parsedCode = code.parse[Source] match
+      case Parsed.Success(Source(stats)) => stats
+      case _                             => Nil
+
+    // Check if there is a main function defined inside an object
+    def checkSignature(defn: Defn.Def) =
+      defn.paramClauseGroups match
+        case List(Member.ParamClauseGroup(
+              Type.ParamClause(Nil),
+              List(Term.ParamClause(
+                List(Term.Param(
+                  Nil,
+                  _: Term.Name,
+                  Some(Type.Apply.After_4_6_0(
+                    Type.Name("Array"),
+                    Type.ArgClause(List(Type.Name("String")))
+                  )),
+                  None
+                )),
+                None
+              ))
+            )) => true
+        case _ => false
+
+    def noToplevelStatements = parsedCode.forall {
+      case _: Term => false
+      case _       => true
+    }
+
+    def hasMainSignature(templ: Template) = templ.body.stats.exists {
+      case defn: Defn.Def =>
+        defn.name.value == "main" && checkSignature(defn)
+      case _ => false
+    }
+    def extendsApp(templ: Template) = templ.inits match
+      case Init.After_4_6_0(Type.Name("App"), _, Nil) :: Nil => true
+      case _                                                 => false
+    val potentialMains = parsedCode.collect {
+      case Defn.Object(_, objName, templ) if extendsApp(templ) || hasMainSignature(templ) =>
+        Seq(objName.value)
+    }.flatten
+
+    potentialMains match
+      case head :: Nil if noToplevelStatements =>
+        ScriptMainMethod.Exists(head)
+      case head :: Nil =>
+        ScriptMainMethod.ToplevelStatsPresent
+      case Nil => ScriptMainMethod.NoMain
+      case seq =>
+        ScriptMainMethod.Multiple(seq)
+
+}
diff --git a/modules/build/src/main/scala/scala/build/internal/util/WarningMessages.scala b/modules/build/src/main/scala/scala/build/internal/util/WarningMessages.scala
@@ -105,6 +105,12 @@ object WarningMessages {
   val offlineModeBloopJvmNotFound =
     "Offline mode is ON and a JVM for Bloop could not be fetched from the local cache, using scalac as fallback"
 
+  def multipleMainObjectsInScript(names: Seq[String]) =
+    s"Only single main is allowed within scripts and multiple main classes were found in the script: ${names.mkString(", ")}"
+
+  def mixedToplvelAndObjectInScript =
+    "Script contains objects with main methods and top-level statements, only the latter will be run."
+
   def directivesInMultipleFilesWarning(
     projectFilePath: String,
     pathsToReport: Iterable[String] = Nil
diff --git a/modules/build/src/main/scala/scala/build/preprocessing/ScriptPreprocessor.scala b/modules/build/src/main/scala/scala/build/preprocessing/ScriptPreprocessor.scala
@@ -134,7 +134,7 @@ case object ScriptPreprocessor extends Preprocessor {
     (codeWrapper: CodeWrapper) =>
       if (containsMainAnnot) logger.diagnostic(
         codeWrapper match {
-          case _: AppCodeWrapper.type =>
+          case _: AppCodeWrapper =>
             WarningMessages.mainAnnotationNotSupported( /* annotationIgnored */ true)
           case _ => WarningMessages.mainAnnotationNotSupported( /* annotationIgnored */ false)
         }
@@ -157,24 +157,27 @@ case object ScriptPreprocessor extends Preprocessor {
     * @return
     *   code wrapper compatible with provided BuildOptions
     */
-  def getScriptWrapper(buildOptions: BuildOptions): CodeWrapper = {
+  def getScriptWrapper(buildOptions: BuildOptions, logger: Logger): CodeWrapper = {
     val effectiveScalaVersion =
       buildOptions.scalaOptions.scalaVersion.flatMap(_.versionOpt)
         .orElse(buildOptions.scalaOptions.defaultScalaVersion)
         .getOrElse(Constants.defaultScalaVersion)
+    def logWarning(msg: String) = logger.diagnostic(msg)
 
     def objectCodeWrapperForScalaVersion =
       // AppObjectWrapper only introduces the 'main.sc' restriction when used in Scala 3, there's no gain in using it with Scala 3
-      if effectiveScalaVersion.startsWith("2") then AppCodeWrapper
-      else ObjectCodeWrapper
+      if effectiveScalaVersion.startsWith("2") then
+        AppCodeWrapper(effectiveScalaVersion, logWarning)
+      else ObjectCodeWrapper(effectiveScalaVersion, logWarning)
 
     buildOptions.scriptOptions.forceObjectWrapper match {
       case Some(true) => objectCodeWrapperForScalaVersion
       case _ =>
         buildOptions.scalaOptions.platform.map(_.value) match {
-          case Some(_: Platform.JS.type)                  => objectCodeWrapperForScalaVersion
-          case _ if effectiveScalaVersion.startsWith("2") => AppCodeWrapper
-          case _                                          => ClassCodeWrapper
+          case Some(_: Platform.JS.type) => objectCodeWrapperForScalaVersion
+          case _ if effectiveScalaVersion.startsWith("2") =>
+            AppCodeWrapper(effectiveScalaVersion, logWarning)
+          case _ => ClassCodeWrapper(effectiveScalaVersion, logWarning)
         }
     }
   }
diff --git a/modules/integration/src/test/scala/scala/cli/integration/RunScriptTestDefinitions.scala b/modules/integration/src/test/scala/scala/cli/integration/RunScriptTestDefinitions.scala
diff --git a/project/deps.sc b/project/deps.sc

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,6 @@ public JavaParserProxy get(`
`20`	`20`	`scala.build.Logger logger,`
`21`	`21`	`Supplier<String> javaCommand`
`22`	`22`	`) {`
`23`		`- return new JavaParserProxyBinary(archiveCache, logger, javaClassNameVersionOpt, javaCommand);`
	`23`	`+ return new JavaParserProxyBinary( archiveCache, logger, javaClassNameVersionOpt, javaCommand);`
`24`	`24`	`}`
`25`	`25`	`}`