[JS IR] handle tests per-file

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/Dce.kt

@@ -87,7 +87,7 @@ private fun buildRoots(modules: Iterable<IrModuleFragment>, context: JsIrBackend
         }
     }
 
-    rootDeclarations += context.testRoots.values
+    rootDeclarations += context.testFunsPerFile.values
 
     val dceRuntimeDiagnostic = context.dceRuntimeDiagnostic
     if (dceRuntimeDiagnostic != null) {

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsCommonBackendContext.kt

@@ -42,7 +42,7 @@ interface JsCommonBackendContext : CommonBackendContext {
     val suiteFun: IrSimpleFunctionSymbol?
     val testFun: IrSimpleFunctionSymbol?
 
-    fun createTestContainerFun(module: IrModuleFragment): IrSimpleFunction
+    fun createTestContainerFun(irFile: IrFile): IrSimpleFunction
 }
 
 // TODO: investigate if it could be removed

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.ir.backend.js
 
-import org.jetbrains.kotlin.backend.common.atMostOne
 import org.jetbrains.kotlin.backend.common.ir.Ir
 import org.jetbrains.kotlin.backend.common.ir.Symbols
 import org.jetbrains.kotlin.builtins.PrimitiveType
@@ -121,12 +120,11 @@ class JsIrBackendContext(
         }
     }
 
-    private val testContainerFuns = mutableMapOf<IrModuleFragment, IrSimpleFunction>()
+    val testFunsPerFile = mutableMapOf<IrFile, IrSimpleFunction>()
 
-    override fun createTestContainerFun(module: IrModuleFragment): IrSimpleFunction {
+    override fun createTestContainerFun(irFile: IrFile): IrSimpleFunction {
-        return testContainerFuns.getOrPut(module) {
+        return testFunsPerFile.getOrPut(irFile) {
-            val file = syntheticFile("tests", module)
+            irFactory.addFunction(irFile) {
-            irFactory.addFunction(file) {
                 name = Name.identifier("test fun")
                 returnType = irBuiltIns.unitType
                 origin = JsIrBuilder.SYNTHESIZED_DECLARATION
@@ -136,9 +134,6 @@ class JsIrBackendContext(
         }
     }
 
-    val testRoots: Map<IrModuleFragment, IrSimpleFunction>
-        get() = testContainerFuns
-
     override val inlineClassesUtils = JsInlineClassesUtils(this)
 
     val innerClassesSupport = JsInnerClassesSupport(mapping, irFactory)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/codegen/IrToJs.kt

@@ -319,9 +319,10 @@ class IrToJs(
             }
         }
 
-        backendContext.testRoots[module]?.let { testContainer ->
+        // TODO: tests
-            statements += invokeFunctionFromEntryJsFileAsStatements(testContainer)
+//        backendContext.testRoots[module]?.let { testContainer ->
-        }
+//            statements += invokeFunctionFromEntryJsFileAsStatements(testContainer)
+//        }
     }
 
     fun generateModule(

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

@@ -8,12 +8,10 @@ package org.jetbrains.kotlin.ir.backend.js
 import org.jetbrains.kotlin.backend.common.lower
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.common.phaser.invokeToplevel
-import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.ir.backend.js.codegen.JsGenerationGranularity
-import org.jetbrains.kotlin.ir.backend.js.ic.SerializedIcData
 import org.jetbrains.kotlin.ir.backend.js.ic.ModuleCache
 import org.jetbrains.kotlin.ir.backend.js.ic.icCompile
-import org.jetbrains.kotlin.ir.backend.js.lower.generateTests
+import org.jetbrains.kotlin.ir.backend.js.lower.generateJsTests
 import org.jetbrains.kotlin.ir.backend.js.lower.moveBodilessDeclarationsToSeparatePlace
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.IrModuleToJsTransformer
 import org.jetbrains.kotlin.ir.backend.js.utils.NameTables
@@ -118,7 +116,7 @@ fun compile(
     }
 
     // TODO should be done incrementally
-    generateTests(context, allModules.last())
+    generateJsTests(context, allModules.last())
 
     if (dceDriven) {
         val controller = MutableController(context, pirLowerings)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/ic/ic.kt

@@ -9,12 +9,10 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.backend.common.lower
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.ir.backend.js.*
-import org.jetbrains.kotlin.ir.backend.js.lower.generateTests
+import org.jetbrains.kotlin.ir.backend.js.lower.generateJsTests
 import org.jetbrains.kotlin.ir.backend.js.lower.moveBodilessDeclarationsToSeparatePlace
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsIrLinker
-import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.IrModuleToJsTransformer
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
-import org.jetbrains.kotlin.ir.declarations.StageController
 import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.declarations.persistent.PersistentIrFactory
 import org.jetbrains.kotlin.ir.util.ExternalDependenciesGenerator
@@ -156,7 +154,7 @@ fun icCompile(
             moveBodilessDeclarationsToSeparatePlace(context, module)
         }
 
-        generateTests(context, modulesToLower.last())
+        generateJsTests(context, modulesToLower.last())
 
         modulesToLower.forEach {
             lowerPreservingIcData(it, context, controller)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/TestGenerator.kt

@@ -27,20 +27,21 @@ import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
-fun generateTests(context: JsCommonBackendContext, moduleFragment: IrModuleFragment) {
+fun generateJsTests(context: JsIrBackendContext, moduleFragment: IrModuleFragment) {
-    val generator = TestGenerator(context) { context.createTestContainerFun(moduleFragment) }
+    val generator = TestGenerator(context, false)
 
     moduleFragment.files.toList().forEach {
         generator.lower(it)
     }
 }
 
-class TestGenerator(val context: JsCommonBackendContext, val testContainerFactory: () -> IrSimpleFunction) : FileLoweringPass {
+class TestGenerator(val context: JsCommonBackendContext, val groupByPackage: Boolean) : FileLoweringPass {
 
     override fun lower(irFile: IrFile) {
-        irFile.declarations.forEach {
+        // Additional copy to prevent ConcurrentModificationException
+        ArrayList(irFile.declarations).forEach {
             if (it is IrClass) {
-                generateTestCalls(it) { suiteForPackage(irFile.fqName) }
+                generateTestCalls(it) { if (groupByPackage) suiteForPackage(irFile) else context.createTestContainerFun(irFile) }
             }
 
             // TODO top-level functions
@@ -49,8 +50,8 @@ class TestGenerator(val context: JsCommonBackendContext, val testContainerFactor
 
     private val packageSuites = mutableMapOf<FqName, IrSimpleFunction>()
 
-    private fun suiteForPackage(fqName: FqName) = packageSuites.getOrPut(fqName) {
+    private fun suiteForPackage(irFile: IrFile) = packageSuites.getOrPut(irFile.fqName) {
-        context.suiteFun!!.createInvocation(fqName.asString(), testContainerFactory())
+        context.suiteFun!!.createInvocation(irFile.fqName.asString(), context.createTestContainerFun(irFile))
     }
 
     private fun IrSimpleFunctionSymbol.createInvocation(

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

@@ -287,20 +287,7 @@ class IrModuleToJsTransformer(
     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()
+        return merge(fragments, staticContext)
-
-        // TODO: handle tests incrementally
-        modules.forEach {
-            backendContext.testRoots[it]?.let { testContainer ->
-                statements.startRegion("block: tests")
-                statements += JsInvocation(staticContext.getNameForStaticFunction(testContainer).makeRef()).makeStmt()
-                statements.endRegion()
-            }
-        }
-
-        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)
@@ -351,13 +338,17 @@ class IrModuleToJsTransformer(
             }
         }
 
+        backendContext.testFunsPerFile[file]?.let {
+            result.testFunInvocation = JsInvocation(staticContext.getNameForStaticFunction(it).makeRef()).makeStmt()
+        }
+
         staticContext.classModels.clear()
         staticContext.initializerBlock.statements.clear()
 
         return result
     }
 
-    private fun merge(fragments: Iterable<JsIrProgramFragment>, staticContext: JsStaticContext): List<JsStatement> {
+    private fun merge(fragments: List<List<JsIrProgramFragment>>, staticContext: JsStaticContext): ModuleBody {
         val statements = mutableListOf<JsStatement>().also {
             if (!generateScriptModule) it += JsStringLiteral("use strict").makeStmt()
         }
@@ -370,9 +361,11 @@ class IrModuleToJsTransformer(
         val classModels = mutableMapOf<IrClassSymbol, JsIrClassModel>()
         val initializerBlock = JsGlobalBlock()
         fragments.forEach {
-            statements += it.declarations.statements
+            it.forEach {
-            classModels += it.classes
+                statements += it.declarations.statements
-            initializerBlock.statements += it.initializers.statements
+                classModels += it.classes
+                initializerBlock.statements += it.initializers.statements
+            }
         }
 
         // sort member forwarding code
@@ -381,7 +374,30 @@ class IrModuleToJsTransformer(
         statements.addWithComment("block: post-declaration", postDeclarationBlock.statements)
         statements.addWithComment("block: init", initializerBlock.statements)
 
-        return statements
+        val lastModuleFragments = fragments.last()
+
+        // Merge test function invocations
+        if (lastModuleFragments.any { it.testFunInvocation != null }) {
+            val testFunBody = JsBlock()
+            val testFun = JsFunction(emptyScope, testFunBody, "root test fun")
+            val suiteFunRef = staticContext.getNameForStaticFunction(backendContext.suiteFun!!.owner).makeRef()
+
+            val tests = lastModuleFragments.filter { it.testFunInvocation != null }.groupBy({ it.packageFqn }) { it.testFunInvocation } // String -> [IrSimpleFunction]
+
+            for ((pkg, testCalls) in tests) {
+                val pkgTestFun = JsFunction(emptyScope, JsBlock(), "test fun for $pkg")
+                pkgTestFun.body.statements += testCalls
+                testFun.body.statements += JsInvocation(suiteFunRef, JsStringLiteral(pkg), JsBooleanLiteral(false), pkgTestFun).makeStmt()
+            }
+
+            statements.startRegion("block: tests")
+            statements += JsInvocation(testFun).makeStmt()
+            statements.endRegion()
+        }
+
+        val callToMain = lastModuleFragments.sortedBy { it.packageFqn }.firstNotNullOfOrNull { it.mainFunction }
+
+        return ModuleBody(statements, callToMain)
     }
 
     private fun generateMainArguments(

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

@@ -14,4 +14,5 @@ class JsIrProgramFragment(val packageFqn: String) {
     val classes = mutableMapOf<IrClassSymbol, JsIrClassModel>()
     val initializers = JsGlobalBlock()
     var mainFunction: JsStatement? = null
+    var testFunInvocation: JsStatement? = null
 }

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmBackendContext.kt

@@ -157,7 +157,8 @@ class WasmBackendContext(
     val testEntryPoints: Collection<IrSimpleFunction>
         get() = testContainerFuns.values
 
-    override fun createTestContainerFun(module: IrModuleFragment): IrSimpleFunction {
+    override fun createTestContainerFun(irFile: IrFile): IrSimpleFunction {
+        val module = irFile.module
         return testContainerFuns.getOrPut(module) {
             val file = syntheticFile("tests", module)
             irFactory.addFunction(file) {

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/TestGenerator.kt

@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
 
 fun generateWasmTests(context: WasmBackendContext, moduleFragment: IrModuleFragment) {
-    val generator = TestGenerator(context) { context.createTestContainerFun(moduleFragment) }
+    val generator = TestGenerator(context, true)
 
     moduleFragment.files.toList().forEach {
         generator.lower(it)