[Wasm] Don't squish modules together

compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JsIrCompiler.kt

@@ -321,7 +321,7 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
 
 
             if (arguments.wasm) {
-                val (moduleFragment, backendContext) = compileToLoweredIr(
+                val (allModules, backendContext) = compileToLoweredIr(
                     depsDescriptors = module,
                     phaseConfig = PhaseConfig(wasmPhases),
                     irFactory = IrFactoryImpl,
@@ -329,7 +329,7 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                     propertyLazyInitialization = arguments.irPropertyLazyInitialization,
                 )
                 val res = compileWasm(
-                    moduleFragment = moduleFragment,
+                    allModules = allModules,
                     backendContext = backendContext,
                     emitNameSection = arguments.wasmDebug,
                     dceEnabled = arguments.irDce,

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/Lower.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.backend.common
 
+import org.jetbrains.kotlin.backend.common.phaser.Action
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBody
@@ -311,3 +312,11 @@ interface DeclarationTransformer : FileLoweringPass {
         }
     }
 }
+
+fun <C> Action<IrElement, C>.toMultiModuleAction(): Action<Iterable<IrModuleFragment>, C> {
+    return { state, modules, context ->
+        modules.forEach { module ->
+            this(state, module, context)
+        }
+    }
+}

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

@@ -66,14 +66,6 @@ private fun makeCustomJsModulePhase(
     actions = setOf(defaultDumper.toMultiModuleAction(), validationAction.toMultiModuleAction()),
 )
 
-private fun <C> Action<IrElement, C>.toMultiModuleAction(): Action<Iterable<IrModuleFragment>, C> {
-    return { state, modules, context ->
-        modules.forEach { module ->
-            this(state, module, context)
-        }
-    }
-}
-
 sealed class Lowering(val name: String) {
     abstract val modulePhase: NamedCompilerPhase<JsIrBackendContext, Iterable<IrModuleFragment>>
 }

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

@@ -6,12 +6,15 @@
 package org.jetbrains.kotlin.backend.wasm
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.lower
 import org.jetbrains.kotlin.backend.common.lower.*
 import org.jetbrains.kotlin.backend.common.lower.inline.FunctionInlining
 import org.jetbrains.kotlin.backend.common.lower.loops.ForLoopsLowering
 import org.jetbrains.kotlin.backend.common.lower.optimizations.PropertyAccessorInlineLowering
 import org.jetbrains.kotlin.backend.common.phaser.*
+import org.jetbrains.kotlin.backend.common.toMultiModuleAction
 import org.jetbrains.kotlin.backend.wasm.lower.*
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.backend.js.lower.*
 import org.jetbrains.kotlin.ir.backend.js.lower.coroutines.AddContinuationToFunctionCallsLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.coroutines.AddContinuationToNonLocalSuspendFunctionsLowering
@@ -27,9 +30,12 @@ private fun makeWasmModulePhase(
     name: String,
     description: String,
     prerequisite: Set<NamedCompilerPhase<WasmBackendContext, *>> = emptySet()
-): NamedCompilerPhase<WasmBackendContext, IrModuleFragment> =
+): NamedCompilerPhase<WasmBackendContext, Iterable<IrModuleFragment>> =
-    makeIrModulePhase(
+    makeCustomWasmModulePhase(
-        lowering, name, description, prerequisite, actions = setOf(validationAction, defaultDumper)
+        op = { context, modules -> lowering(context).lower(modules) },
+        name = name,
+        description = description,
+        prerequisite = prerequisite
     )
 
 private fun makeCustomWasmModulePhase(
@@ -37,9 +43,25 @@ private fun makeCustomWasmModulePhase(
     description: String,
     name: String,
     prerequisite: Set<NamedCompilerPhase<WasmBackendContext, *>> = emptySet()
-): NamedCompilerPhase<WasmBackendContext, IrModuleFragment> =
+): NamedCompilerPhase<WasmBackendContext, Iterable<IrModuleFragment>> =
-    makeCustomPhase(
+    NamedCompilerPhase(
-        op, name, description, prerequisite, actions = setOf(defaultDumper, validationAction), nlevels = 0,
+        name = name,
+        description = description,
+        prerequisite = prerequisite,
+        lower = object : SameTypeCompilerPhase<WasmBackendContext, Iterable<IrModuleFragment>> {
+            override fun invoke(
+                phaseConfig: PhaseConfig,
+                phaserState: PhaserState<Iterable<IrModuleFragment>>,
+                context: WasmBackendContext,
+                input: Iterable<IrModuleFragment>
+            ): Iterable<IrModuleFragment> {
+                input.forEach { module ->
+                    op(context, module)
+                }
+                return input
+            }
+        },
+        actions = setOf(defaultDumper.toMultiModuleAction(), validationAction.toMultiModuleAction())
     )
 
 private val validateIrBeforeLowering = makeCustomWasmModulePhase(

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

@@ -32,7 +32,7 @@ fun compileToLoweredIr(
     irFactory: IrFactory,
     exportedDeclarations: Set<FqName> = emptySet(),
     propertyLazyInitialization: Boolean,
-): Pair<IrModuleFragment, WasmBackendContext> {
+): Pair<List<IrModuleFragment>, WasmBackendContext> {
     val mainModule = depsDescriptors.mainModule
     val configuration = depsDescriptors.compilerConfiguration
     val (moduleFragment, dependencyModules, irBuiltIns, symbolTable, deserializer) = loadIr(
@@ -55,38 +55,36 @@ fun compileToLoweredIr(
         ExternalDependenciesGenerator(symbolTable, listOf(deserializer)).generateUnboundSymbolsAsDependencies()
     }
 
-    val irFiles = allModules.flatMap { it.files }
-    moduleFragment.files.clear()
-    moduleFragment.files += irFiles
-
     // Create stubs
     ExternalDependenciesGenerator(symbolTable, listOf(deserializer)).generateUnboundSymbolsAsDependencies()
-    moduleFragment.patchDeclarationParents()
+    allModules.forEach { it.patchDeclarationParents() }
 
     deserializer.postProcess()
     symbolTable.noUnboundLeft("Unbound symbols at the end of linker")
 
-    moduleFragment.files.forEach { irFile -> markExportedDeclarations(context, irFile, exportedDeclarations) }
+    for (module in allModules)
+        for (file in module.files)
+            markExportedDeclarations(context, file, exportedDeclarations)
 
-    wasmPhases.invokeToplevel(phaseConfig, context, moduleFragment)
+    wasmPhases.invokeToplevel(phaseConfig, context, allModules)
 
-    return Pair(moduleFragment, context)
+    return Pair(allModules, context)
 }
 
 fun compileWasm(
-    moduleFragment: IrModuleFragment,
+    allModules: List<IrModuleFragment>,
     backendContext: WasmBackendContext,
     emitNameSection: Boolean = false,
     dceEnabled: Boolean = false,
 ): WasmCompilerResult {
 
     if (dceEnabled) {
-        eliminateDeadDeclarations(listOf(moduleFragment), backendContext)
+        eliminateDeadDeclarations(allModules, backendContext)
     }
 
     val compiledWasmModule = WasmCompiledModuleFragment(backendContext.irBuiltIns)
     val codeGenerator = WasmModuleFragmentGenerator(backendContext, compiledWasmModule, allowIncompleteImplementations = dceEnabled)
-    codeGenerator.generateModule(moduleFragment)
+    allModules.forEach { codeGenerator.generateModule(it) }
 
     val linkedModule = compiledWasmModule.linkWasmCompiledFragments()
     val watGenerator = WasmIrToText()
@@ -96,7 +94,7 @@ fun compileWasm(
     val js = compiledWasmModule.generateJs()
 
     val os = ByteArrayOutputStream()
-    WasmIrToBinary(os, linkedModule, moduleFragment.descriptor.name.asString(), emitNameSection).appendWasmModule()
+    WasmIrToBinary(os, linkedModule, allModules.last().descriptor.name.asString(), emitNameSection).appendWasmModule()
     val byteArray = os.toByteArray()
 
     return WasmCompilerResult(

js/js.tests/test/org/jetbrains/kotlin/js/testOld/BasicWasmBoxTest.kt

@@ -200,7 +200,7 @@ abstract class BasicWasmBoxTest(
         jsFilesBefore: List<String>,
         jsFilesAfter: List<String>,
     ) {
-        val (moduleFragment, backendContext) = compileToLoweredIr(
+        val (allModules, backendContext) = compileToLoweredIr(
             depsDescriptors = sourceModule,
             phaseConfig = phaseConfig,
             irFactory = IrFactoryImpl,
@@ -209,7 +209,7 @@ abstract class BasicWasmBoxTest(
         )
 
         val compilerResult = compileWasm(
-            moduleFragment = moduleFragment,
+            allModules = allModules,
             backendContext = backendContext,
             emitNameSection = true,
             dceEnabled = dceEnabled,