Fix filling Main-Class attribute in manifest file of output jar

Before 5b7cee6221 JVM CLI compiler
was calling `KotlinToJVMBytecodeCompiler.compileBunchOfSources`.
`compileBunchOfSources` detected possible main classes,
and filled the Main-Class attribute in output jar
if if there was only one candidate.

After the change JVM CLI began calling
`KotlinToJVMBytecodeCompiler.compileModules`, which was not searching for a main class.

This change adds searching for main classes to `compileModules`.
We search for a main class only when one module is compiled,
and an output is written a jar file (so the change only affects JVM CLI compilation).

    #KT-32272 Fixed

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt

@@ -83,14 +83,14 @@ object KotlinToJVMBytecodeCompiler {
     private fun writeOutput(
         configuration: CompilerConfiguration,
         outputFiles: OutputFileCollection,
-        mainClass: FqName?
+        mainClassProvider: MainClassProvider?
     ) {
         val reportOutputFiles = configuration.getBoolean(CommonConfigurationKeys.REPORT_OUTPUT_FILES)
         val jarPath = configuration.get(JVMConfigurationKeys.OUTPUT_JAR)
         val messageCollector = configuration.get(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, MessageCollector.NONE)
         if (jarPath != null) {
             val includeRuntime = configuration.get(JVMConfigurationKeys.INCLUDE_RUNTIME, false)
-            CompileEnvironmentUtil.writeToJar(jarPath, includeRuntime, mainClass, outputFiles)
+            CompileEnvironmentUtil.writeToJar(jarPath, includeRuntime, mainClassProvider?.mainClassFqName, outputFiles)
             if (reportOutputFiles) {
                 val message = OutputMessageUtil.formatOutputMessage(outputFiles.asList().flatMap { it.sourceFiles }.distinct(), jarPath)
                 messageCollector.report(OUTPUT, message)
@@ -108,7 +108,7 @@ object KotlinToJVMBytecodeCompiler {
         }
         return GenerationStateEventCallback { state ->
             val currentOutput = SimpleOutputFileCollection(state.factory.currentOutput)
-            writeOutput(configuration, currentOutput, mainClass = null)
+            writeOutput(configuration, currentOutput, null)
             if (!configuration.get(JVMConfigurationKeys.RETAIN_OUTPUT_IN_MEMORY, false)) {
                 state.factory.releaseGeneratedOutput()
             }
@@ -207,7 +207,8 @@ object KotlinToJVMBytecodeCompiler {
         try {
             for ((_, state) in outputs) {
                 ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
-                writeOutput(state.configuration, state.factory, null)
+                val mainClassProvider = if (outputs.size == 1) MainClassProvider(state, environment) else null
+                writeOutput(state.configuration, state.factory, mainClassProvider)
             }
 
             if (projectConfiguration.getBoolean(JVMConfigurationKeys.COMPILE_JAVA)) {
@@ -397,16 +398,20 @@ object KotlinToJVMBytecodeCompiler {
             (File(path).takeIf(File::isAbsolute) ?: buildFile.resolveSibling(path)).absolutePath
         }
 
-    private fun findMainClass(generationState: GenerationState, files: List<KtFile>): FqName? {
-        val mainFunctionDetector = MainFunctionDetector(generationState.bindingContext, generationState.languageVersionSettings)
-        return files.asSequence()
-            .map { file ->
-                if (mainFunctionDetector.hasMain(file.declarations))
-                    JvmFileClassUtil.getFileClassInfoNoResolve(file).facadeClassFqName
-                else
-                    null
-            }
-            .singleOrNull { it != null }
+    private class MainClassProvider(generationState: GenerationState, environment: KotlinCoreEnvironment) {
+        val mainClassFqName: FqName? by lazy { findMainClass(generationState, environment.getSourceFiles()) }
+
+        private fun findMainClass(generationState: GenerationState, files: List<KtFile>): FqName? {
+            val mainFunctionDetector = MainFunctionDetector(generationState.bindingContext, generationState.languageVersionSettings)
+            return files.asSequence()
+                .map { file ->
+                    if (mainFunctionDetector.hasMain(file.declarations))
+                        JvmFileClassUtil.getFileClassInfoNoResolve(file).facadeClassFqName
+                    else
+                        null
+                }
+                .singleOrNull { it != null }
+        }
     }
 
     fun compileBunchOfSources(environment: KotlinCoreEnvironment): Boolean {
@@ -421,10 +426,8 @@ object KotlinToJVMBytecodeCompiler {
 
         val generationState = analyzeAndGenerate(environment) ?: return false
 
-        val mainClass = findMainClass(generationState, environment.getSourceFiles())
-
         try {
-            writeOutput(environment.configuration, generationState.factory, mainClass)
+            writeOutput(environment.configuration, generationState.factory, MainClassProvider(generationState, environment))
             return true
         } finally {
             generationState.destroy()

compiler/tests/org/jetbrains/kotlin/cli/CustomCliTest.kt

@@ -5,10 +5,14 @@
 
 package org.jetbrains.kotlin.cli
 
+import junit.framework.Assert
 import org.jetbrains.kotlin.cli.jvm.K2JVMCompiler
 import org.jetbrains.kotlin.test.CompilerTestUtil
 import org.jetbrains.kotlin.test.TestCaseWithTmpdir
 import java.io.File
+import java.util.jar.JarFile
+
+private const val EMPTY_MAIN_FUN = "fun main() {}"
 
 class CustomCliTest : TestCaseWithTmpdir() {
     fun testArgfileWithNonTrivialWhitespaces() {
@@ -16,4 +20,33 @@ class CustomCliTest : TestCaseWithTmpdir() {
         val argfile = File(tmpdir, "argfile").apply { writeText(text, Charsets.UTF_8) }
         CompilerTestUtil.executeCompilerAssertSuccessful(K2JVMCompiler(), listOf("@" + argfile.absolutePath))
     }
+
+    fun testMainClass() {
+        val mainKt = tmpdir.resolve("main.kt").apply {
+            writeText(EMPTY_MAIN_FUN)
+        }
+        compileAndCheckMainClass(listOf(mainKt), expectedMainClass = "MainKt")
+    }
+
+    fun testMultipleMainClasses() {
+        val main1Kt = tmpdir.resolve("main1.kt").apply {
+            writeText(EMPTY_MAIN_FUN)
+        }
+        val main2Kt = tmpdir.resolve("main2.kt").apply {
+            writeText(EMPTY_MAIN_FUN)
+        }
+
+        compileAndCheckMainClass(listOf(main1Kt, main2Kt), expectedMainClass = null)
+    }
+
+    private fun compileAndCheckMainClass(sourceFiles: List<File>, expectedMainClass: String?) {
+        val jarFile = tmpdir.resolve("output.jar")
+        val args = listOf("-include-runtime", "-d", jarFile.absolutePath) + sourceFiles.map { it.absolutePath }
+        CompilerTestUtil.executeCompilerAssertSuccessful(K2JVMCompiler(), args)
+
+        JarFile(jarFile).use {
+            val mainClassAttr = it.manifest.mainAttributes.getValue("Main-Class")
+            Assert.assertEquals(expectedMainClass, mainClassAttr)
+        }
+    }
 }