[JS IR] split resulting code into fragments

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

@@ -217,10 +217,10 @@ class IrModuleToJsTransformer(
                     sourceMapContentEmbedding != SourceMapSourceEmbedding.NEVER
                 )
             } else {
-                null
+                NoOpSourceLocationConsumer
             }
 
-        program.accept(JsToStringGenerationVisitor(jsCode, sourceMapBuilderConsumer ?: NoOpSourceLocationConsumer))
+        program.accept(JsToStringGenerationVisitor(jsCode, sourceMapBuilderConsumer))
 
         return CompilationOutputs(
             jsCode.toString(),
@@ -283,21 +283,33 @@ class IrModuleToJsTransformer(
     }
 
     private fun generateModuleBody(modules: Iterable<IrModuleFragment>, staticContext: JsStaticContext): List<JsStatement> {
-        val statements = mutableListOf<JsStatement>().also {
+        val fragments = modules.flatMap { it.files.map { generateProgramFragment(it, staticContext) } }
-            if (!generateScriptModule) it += JsStringLiteral("use strict").makeStmt()
+
+        val statements = merge(fragments, staticContext).toMutableList()
+
+        // 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 preDeclarationBlock = JsGlobalBlock()
+        return statements
-        val postDeclarationBlock = JsGlobalBlock()
+    }
 
-        statements.addWithComment("block: pre-declaration", preDeclarationBlock)
+    private val generateFilePaths = backendContext.configuration.getBoolean(JSConfigurationKeys.GENERATE_COMMENTS_WITH_FILE_PATH)
+    private val pathPrefixMap = backendContext.configuration.getMap(JSConfigurationKeys.FILE_PATHS_PREFIX_MAP)
 
-        val generateFilePaths = backendContext.configuration.getBoolean(JSConfigurationKeys.GENERATE_COMMENTS_WITH_FILE_PATH)
+    private fun generateProgramFragment(file: IrFile, staticContext: JsStaticContext): JsIrProgramFragment {
-        val pathPrefixMap = backendContext.configuration.getMap(JSConfigurationKeys.FILE_PATHS_PREFIX_MAP)
+        require(staticContext.classModels.isEmpty())
+        require(staticContext.initializerBlock.statements.isEmpty())
 
-        modules.forEach { module ->
+        val result = JsIrProgramFragment()
-            module.files.forEach {
+        val statements = result.declarations.statements
-                val fileStatements = it.accept(IrFileToJsTransformer(), staticContext).statements
+
+        val fileStatements = file.accept(IrFileToJsTransformer(), staticContext).statements
         if (fileStatements.isNotEmpty()) {
             var startComment = ""
 
@@ -306,7 +318,7 @@ class IrModuleToJsTransformer(
             }
 
             if (generateRegionComments || generateFilePaths) {
-                        val originalPath = it.path
+                val originalPath = file.path
                 val path = pathPrefixMap.entries
                     .find { (k, _) -> originalPath.startsWith(k) }
                     ?.let { (k, v) -> v + originalPath.substring(k.length) }
@@ -322,22 +334,39 @@ class IrModuleToJsTransformer(
             statements.addAll(fileStatements)
             statements.endRegion()
         }
+
+        result.classes += staticContext.classModels
+        result.initializers.statements += staticContext.initializerBlock.statements
+
+        staticContext.classModels.clear()
+        staticContext.initializerBlock.statements.clear()
+
+        return result
     }
+
+    private fun merge(fragments: Iterable<JsIrProgramFragment>, staticContext: JsStaticContext): List<JsStatement> {
+        val statements = mutableListOf<JsStatement>().also {
+            if (!generateScriptModule) it += JsStringLiteral("use strict").makeStmt()
+        }
+
+        val preDeclarationBlock = JsGlobalBlock()
+        val postDeclarationBlock = JsGlobalBlock()
+
+        statements.addWithComment("block: pre-declaration", preDeclarationBlock)
+
+        val classModels = mutableMapOf<IrClassSymbol, JsIrClassModel>()
+        val initializerBlock = JsGlobalBlock()
+        fragments.forEach {
+            statements += it.declarations.statements
+            classModels += it.classes
+            initializerBlock.statements += it.initializers.statements
         }
 
         // sort member forwarding code
-        processClassModels(staticContext.classModels, preDeclarationBlock, postDeclarationBlock)
+        processClassModels(classModels, preDeclarationBlock, postDeclarationBlock)
 
         statements.addWithComment("block: post-declaration", postDeclarationBlock.statements)
-        statements.addWithComment("block: init", staticContext.initializerBlock.statements)
+        statements.addWithComment("block: init", initializerBlock.statements)
-
-        modules.forEach {
-            backendContext.testRoots[it]?.let { testContainer ->
-                statements.startRegion("block: tests")
-                statements += JsInvocation(staticContext.getNameForStaticFunction(testContainer).makeRef()).makeStmt()
-                statements.endRegion()
-            }
-        }
 
         return statements
     }

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

@@ -0,0 +1,15 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
+
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.js.backend.ast.JsGlobalBlock
+
+class JsIrProgramFragment {
+    val declarations = JsGlobalBlock()
+    val classes = mutableMapOf<IrClassSymbol, JsIrClassModel>()
+    val initializers = JsGlobalBlock()
+}