[JS IR] handle main function per-file

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrModuleToJsTransformer.kt

@@ -150,7 +150,7 @@ class IrModuleToJsTransformer(
                 declareFreshGlobal = { JsName(sanitizeName(it)) } // TODO: Declare fresh name
             )
 
-        val moduleBody = generateModuleBody(modules, staticContext)
+        val (moduleBody, callToMain) = generateModuleBody(modules, staticContext)
 
         val internalModuleName = JsName("_")
         val globalNames = NameTable<String>(namer.globalNames)
@@ -175,8 +175,8 @@ class IrModuleToJsTransformer(
                     statements.addWithComment("block: imports", importStatements + crossModuleImports)
                     statements += moduleBody
                     statements.addWithComment("block: exports", exportStatements + crossModuleExports)
-                    if (generateMainCall) {
+                    if (generateMainCall && callToMain != null) {
-                        statements += generateCallToMain(modules, staticContext)
+                        statements += callToMain
                     }
                     statements += JsReturn(internalModuleName.makeRef())
                 }
@@ -282,10 +282,12 @@ class IrModuleToJsTransformer(
         }
     }
 
-    private fun generateModuleBody(modules: Iterable<IrModuleFragment>, staticContext: JsStaticContext): List<JsStatement> {
+    private data class ModuleBody(val statements: List<JsStatement>, val mainFunction: JsStatement?)
-        val fragments = modules.flatMap { it.files.map { generateProgramFragment(it, staticContext) } }
 
-        val statements = merge(fragments, staticContext).toMutableList()
+    private fun generateModuleBody(modules: Iterable<IrModuleFragment>, staticContext: JsStaticContext): ModuleBody {
+        val fragments = modules.map { it.files.map { generateProgramFragment(it, staticContext) } }
+
+        val statements = merge(fragments.flatMap { it }, staticContext).toMutableList()
 
         // TODO: handle tests incrementally
         modules.forEach {
@@ -296,7 +298,9 @@ class IrModuleToJsTransformer(
             }
         }
 
-        return statements
+        val callToMain = fragments.last().sortedBy { it.packageFqn }.firstNotNullOfOrNull { it.mainFunction }
+
+        return ModuleBody(statements, callToMain)
     }
 
     private val generateFilePaths = backendContext.configuration.getBoolean(JSConfigurationKeys.GENERATE_COMMENTS_WITH_FILE_PATH)
@@ -306,7 +310,7 @@ class IrModuleToJsTransformer(
         require(staticContext.classModels.isEmpty())
         require(staticContext.initializerBlock.statements.isEmpty())
 
-        val result = JsIrProgramFragment()
+        val result = JsIrProgramFragment(file.fqName.asString())
         val statements = result.declarations.statements
 
         val fileStatements = file.accept(IrFileToJsTransformer(), staticContext).statements
@@ -338,6 +342,15 @@ class IrModuleToJsTransformer(
         result.classes += staticContext.classModels
         result.initializers.statements += staticContext.initializerBlock.statements
 
+        if (mainArguments != null) {
+            JsMainFunctionDetector(backendContext).getMainFunctionOrNull(file)?.let {
+                val jsName = staticContext.getNameForStaticFunction(it)
+                val generateArgv = it.valueParameters.firstOrNull()?.isStringArrayParameter() ?: false
+                val generateContinuation = it.isLoweredSuspendFunction(backendContext)
+                result.mainFunction = JsInvocation(jsName.makeRef(), generateMainArguments(generateArgv, generateContinuation, staticContext)).makeStmt()
+            }
+        }
+
         staticContext.classModels.clear()
         staticContext.initializerBlock.statements.clear()
 
@@ -390,18 +403,6 @@ class IrModuleToJsTransformer(
         return listOfNotNull(mainArgumentsArray, continuation)
     }
 
-    private fun generateCallToMain(modules: Iterable<IrModuleFragment>, staticContext: JsStaticContext): List<JsStatement> {
-        // TODO: Generate calls to main as IR->IR lowering
-        if (mainArguments == null) return emptyList() // in case `NO_MAIN` and `main(..)` exists
-        val mainFunction = JsMainFunctionDetector(backendContext).getMainFunctionOrNull(modules.last())
-        return mainFunction?.let {
-            val jsName = staticContext.getNameForStaticFunction(it)
-            val generateArgv = it.valueParameters.firstOrNull()?.isStringArrayParameter() ?: false
-            val generateContinuation = it.isLoweredSuspendFunction(backendContext)
-            listOf(JsInvocation(jsName.makeRef(), generateMainArguments(generateArgv, generateContinuation, staticContext)).makeStmt())
-        } ?: emptyList()
-    }
-
     private fun generateImportStatements(
         getNameForExternalDeclaration: (IrDeclarationWithName) -> JsName,
         declareFreshGlobal: (String) -> JsName

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/JsIrProgramFragment.kt

@@ -7,9 +7,11 @@ package org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
 
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.js.backend.ast.JsGlobalBlock
+import org.jetbrains.kotlin.js.backend.ast.JsStatement
 
-class JsIrProgramFragment {
+class JsIrProgramFragment(val packageFqn: String) {
     val declarations = JsGlobalBlock()
     val classes = mutableMapOf<IrClassSymbol, JsIrClassModel>()
     val initializers = JsGlobalBlock()
+    var mainFunction: JsStatement? = null
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/utils/JsMainFunctionDetector.kt

@@ -46,14 +46,18 @@ class JsMainFunctionDetector(val context: JsCommonBackendContext) {
         }
     }
 
+    fun getMainFunctionOrNull(file: IrFile): IrSimpleFunction? {
+        // TODO: singleOrNull looks suspicious
+        return file.declarations.filterIsInstance<IrSimpleFunction>().singleOrNull { it.isMain(allowEmptyParameters = true) }
+    }
+
     fun getMainFunctionOrNull(module: IrModuleFragment): IrSimpleFunction? {
 
         var resultPair: Pair<String, IrSimpleFunction>? = null
 
         module.files.forEach { f ->
             val fqn = f.fqName.asString()
-
+            getMainFunctionOrNull(f)?.let {
-            f.declarations.filterIsInstance<IrSimpleFunction>().singleOrNull { it.isMain(allowEmptyParameters = true) }?.let {
                 val result = resultPair
                 if (result == null) {
                     resultPair = Pair(fqn, it)