Synthetize main classes for @mainarg.main methods in Scala scripts for Intellij/@scala.main compatibility (#6084)

* `@scala.main` synthesizes a main class for each main method, which
IntelliJ expects when you click the "run" button next to the main method
`def`.
* `@mainargs.main` does not, as it expects a single `def main` method to
call `mainargs.Parser(this).runOrExit` to dispatch to the various
annotated entrypoints
* This PR bridges the gap by:
* Hackily parsing out the name of the annotated main methods from the
`mainargs.Parser(this).` macro expansion in the synthesized `def main`
method (Ideally we could fish out the annotations with ASM, but despite
`@mainargs.main` being a `ClassFileAnnotation` it does not appear in the
bytecode),
* Generates synthetic main classes corresponding to each one that
forward to the respective annotated method
* Filters them out during `allLocalMainClasses` discovery, so that
normal `run` only runs the primary synthetic `main`, while `runMain` or
similar (e.g. what IntelliJ does) can run the others

Covered by an integration test and tested manually in IntelliJ to make
sure the run button works in single-main and multi-main scenarios

<img width="976" height="652" alt="Screenshot 2025-11-01 at 11 15 03 PM"
src="https://github.com/user-attachments/assets/a9a4a02e-f122-4977-9be4-c3e2bd55e960"
/>

Author: lihaoyi

integration/feature/auxiliary-class-files/resources/build.mill

@@ -9,6 +9,8 @@ object app extends Module {
   trait Common extends ScalaModule {
     def moduleDir = super.moduleDir / os.up
     def scalaVersion = "3.4.0"
+    // force this on for testing even though there's only one file
+    def zincIncrementalCompilation = true
   }
 
   object jvm extends Common

integration/feature/script-main-forwarder-classes/resources/Multi.scala

@@ -0,0 +1,8 @@
+@main
+def main1(text: String) = println(text + "123")
+
+@main
+def main2(text: String) = println(text + "456")
+
+@main
+def main3(text: String) = println(text + "789")

integration/feature/script-main-forwarder-classes/resources/RawMainSignature.scala

@@ -0,0 +1,2 @@
+@scala.main
+def main1(text: String) = println(text + "OMG")

integration/feature/script-main-forwarder-classes/resources/ScalaMain.scala

@@ -0,0 +1,2 @@
+@scala.main
+def main1(text: String) = println(text + "ABC")

integration/feature/script-main-forwarder-classes/resources/Single.scala

@@ -0,0 +1,2 @@
+@main
+def main1(text: String) = println(text + "XYZ")

integration/feature/script-main-forwarder-classes/src/ScriptMainForwarderClassesTests.scala

@@ -0,0 +1,45 @@
+package mill.integration
+
+import mill.testkit.UtestIntegrationTestSuite
+import utest._
+
+// Make sure that for scripts with multiple `@main` methods (aliased to `@mainargs.main`),
+// we generate synthetic main classes for each one following the name of the class that we
+// can run using `runMain`. This mimics the behavior of `@scala.main`, and allows interop
+// with tools that expect that behavior such as the IntelliJ `run` button
+object ScriptMainForwarderClassesTests extends UtestIntegrationTestSuite {
+  val tests: Tests = Tests {
+    test("test") - integrationTest { tester =>
+      import tester._
+      // When using `run`, we ignore the synthetic main classes when picking a main method,
+      // so in this case we run the default `_MillScriptMain` method which delegates
+      // to the relevant method internally based on the first token
+      val res0 = eval(("Multi.scala:run", "main1", "--text", "HELLO"))
+      assert(res0.out == "HELLO123")
+
+      // Multi-main script have forwarder classes synthesized for each @mainargs.main method,
+      // which forward to `_MillScriptMain` but pass the name as the first param to disambiguate
+      val res1 = eval(("Multi.scala:runMain", "main1", "--text", "hello"))
+      assert(res1.out == "hello123")
+
+      val res2 = eval(("Multi.scala:runMain", "main2", "--text", "world"))
+      assert(res2.out == "world456")
+
+      val res3 = eval(("Multi.scala:runMain", "main3", "--text", "moooo"))
+      assert(res3.out == "moooo789")
+
+      // Single-main script, forwarder class must *not* pass the method name as
+      // the first parameter
+      val res4 = eval(("Single.scala:runMain", "main1", "--text", "iamcow"))
+      assert(res4.out == "iamcowXYZ")
+
+      // scala.main method takes priority over synthetic _MillScriptMain method
+      val res5 = eval(("ScalaMain.scala:run", "hearmemoo"))
+      assert(res5.out == "hearmemooABC")
+
+      // `def main(args: Array[String]): Unit` method takes priority over synthetic _MillScriptMain method
+      val res6 = eval(("RawMainSignature.scala:run", "iweightwiceasmuchasyou"))
+      assert(res6.out == "iweightwiceasmuchasyouOMG")
+    }
+  }
+}

libs/javalib/src/mill/javalib/JavaModule.scala

@@ -836,9 +836,12 @@ trait JavaModule
     ).equalsIgnoreCase("true")
   }
 
-  def zincIncrementalCompilation: T[Boolean] = Task {
-    true
-  }
+  /**
+   * Whether to turn on zinc incremental compilation or not, as it can speed things up
+   * by skipping some source files but also adds some performance overhead. Defaults
+   * to turning it on if there is more than one source file being compiled
+   */
+  def zincIncrementalCompilation: T[Boolean] = Task { allSourceFiles().length > 1 }
 
   /**
    * Compiles the current module to generate compiled classfiles/bytecode.

libs/script/asm-worker/src/mill/script/asm/AsmWorkerImpl.scala

@@ -0,0 +1,134 @@
+package mill.script.asm
+
+import org.objectweb.asm
+
+object AsmWorkerImpl {
+
+  def generateSyntheticClasses(classesDir: java.nio.file.Path, mainMethods: Array[String]): Unit = {
+    mainMethods.foreach { methodName =>
+      generateSyntheticMainClass(os.Path(classesDir), methodName, mainMethods.size > 1)
+    }
+  }
+
+  def findMainArgsMethods(classesDir: java.nio.file.Path): Array[String] = {
+    val mainMethods = collection.mutable.ArrayBuffer[String]()
+
+    // Look for _MillScriptMain$ class which contains the mainargs.Parser code
+    val millScriptMainClass = os.Path(classesDir) / "_MillScriptMain$.class"
+
+    if (os.exists(millScriptMainClass)) {
+      val reader = new asm.ClassReader(os.read.bytes(millScriptMainClass))
+
+      val visitor = new asm.ClassVisitor(asm.Opcodes.ASM9) {
+        override def visitMethod(
+            access: Int,
+            name: String,
+            descriptor: String,
+            signature: String,
+            exceptions: Array[String]
+        ): asm.MethodVisitor = {
+          new asm.MethodVisitor(asm.Opcodes.ASM9) {
+            private val stringsSinceLastCreate = collection.mutable.ArrayBuffer[String]()
+
+            override def visitLdcInsn(value: Any): Unit = {
+              value match {
+                case s: String => stringsSinceLastCreate += s
+                case _ =>
+              }
+              super.visitLdcInsn(value)
+            }
+
+            override def visitMethodInsn(
+                opcode: Int,
+                owner: String,
+                methodName: String,
+                descriptor: String,
+                isInterface: Boolean
+            ): Unit = {
+              // Look for MainData.create calls which include the method name as first parameter
+              if (owner.contains("MainData") && methodName == "create") {
+                // The first string constant before MainData.create is the method name
+                if (stringsSinceLastCreate.nonEmpty) {
+                  val potentialMethodName = stringsSinceLastCreate.head
+                  mainMethods += potentialMethodName
+                  stringsSinceLastCreate.clear()
+                }
+              }
+
+              super.visitMethodInsn(opcode, owner, methodName, descriptor, isInterface)
+            }
+          }
+        }
+      }
+
+      reader.accept(visitor, 0)
+    }
+
+    mainMethods.toArray.distinct
+  }
+
+  def generateSyntheticMainClass(
+      classesDir: os.Path,
+      methodName: String,
+      multiMain: Boolean
+  ): Unit = {
+    val templateClassName = if (multiMain) "TemplateMultiMainClass" else "TemplateSingleMainClass"
+    val templateBytes = os.read.bytes(
+      os.resource(getClass.getClassLoader) / os.SubPath(s"mill/script/asm/$templateClassName.class")
+    )
+    val reader = new asm.ClassReader(templateBytes)
+    val writer = new asm.ClassWriter(reader, 0)
+
+    val visitor = new asm.ClassVisitor(asm.Opcodes.ASM9, writer) {
+      override def visit(
+          version: Int,
+          access: Int,
+          name: String,
+          signature: String,
+          superName: String,
+          interfaces: Array[String]
+      ): Unit = {
+        super.visit(version, access, methodName, signature, superName, interfaces)
+      }
+
+      override def visitMethod(
+          access: Int,
+          name: String,
+          descriptor: String,
+          signature: String,
+          exceptions: Array[String]
+      ): asm.MethodVisitor = {
+        val mv = super.visitMethod(access, name, descriptor, signature, exceptions)
+
+        new asm.MethodVisitor(asm.Opcodes.ASM9, mv) {
+          override def visitLdcInsn(value: Any): Unit = {
+            // Replace "TEMPLATE_METHOD_NAME" with actual method name
+            if (value == "TEMPLATE_METHOD_NAME") super.visitLdcInsn(methodName)
+            else super.visitLdcInsn(value)
+          }
+
+          override def visitMethodInsn(
+              opcode: Int,
+              owner: String,
+              name: String,
+              descriptor: String,
+              isInterface: Boolean
+          ): Unit = {
+            // Replace TemplateMainClass.main call with _MillScriptMain$.main
+            if (owner == s"mill/script/asm/$templateClassName" && name == "main") {
+              super.visitMethodInsn(opcode, "_MillScriptMain", name, descriptor, isInterface)
+            } else {
+              super.visitMethodInsn(opcode, owner, name, descriptor, isInterface)
+            }
+          }
+        }
+      }
+    }
+
+    reader.accept(visitor, 0)
+
+    // Write the modified class file
+    val classBytes = writer.toByteArray
+    os.write.over(classesDir / s"$methodName.class", classBytes)
+  }
+}

libs/script/asm-worker/src/mill/script/asm/TemplateMultiMainClass.java

@@ -0,0 +1,10 @@
+package mill.script.asm;
+
+public class TemplateMultiMainClass {
+  public static void main(String[] args) {
+    String[] newArgs = new String[args.length + 1];
+    newArgs[0] = "TEMPLATE_METHOD_NAME";
+    System.arraycopy(args, 0, newArgs, 1, args.length);
+    TemplateMultiMainClass.main(newArgs);
+  }
+}

libs/script/asm-worker/src/mill/script/asm/TemplateSingleMainClass.java

@@ -0,0 +1,7 @@
+package mill.script.asm;
+
+public class TemplateSingleMainClass {
+  public static void main(String[] args) {
+    TemplateSingleMainClass.main(args);
+  }
+}