[Commonizer] Dump every module to disk immediately when it's ready

No need to retain all metadata modules in memory -> lesser memory consumption.

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/CommonizerParameters.kt

@@ -38,5 +38,14 @@ class CommonizerParameters(
             _resultsConsumer = value
         }
 
-    fun hasAnythingToCommonize(): Boolean = _targetProviders.size >= 2
+    fun hasAnythingToCommonize(): Boolean {
+        if (_targetProviders.size < 2) return false // too few targets
+
+        val allModuleNames: List<Set<String>> = _targetProviders.values.map { targetProvider ->
+            targetProvider.modulesProvider.loadModuleInfos().keys
+        }
+        val commonModuleNames: Set<String> = allModuleNames.reduce { a, b -> a intersect b }
+
+        return commonModuleNames.isNotEmpty() // there are modules that are present in every target
+    }
 }

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/ResultsConsumer.kt

@@ -21,7 +21,20 @@ interface ResultsConsumer {
         class Commonized(override val libraryName: String, val metadata: SerializedMetadata) : ModuleResult()
     }
 
-    fun consumeResults(target: CommonizerTarget, moduleResults: Collection<ModuleResult>)
+    /**
+     * Consume a single [ModuleResult] for the specified [CommonizerTarget].
+     */
+    fun consume(target: CommonizerTarget, moduleResult: ModuleResult)
 
-    fun successfullyFinished(status: Status)
+    /**
+     * Mark the specified [CommonizerTarget] as fully consumed.
+     * It's forbidden to make subsequent [consume] calls for fully consumed targets.
+     */
+    fun targetConsumed(target: CommonizerTarget)
+
+    /**
+     * Notify that all results have been consumed.
+     * It's forbidden to make any subsequent [consume] and [targetConsumed] calls after this call.
+     */
+    fun allConsumed(status: Status)
 }

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/facade.kt

@@ -19,7 +19,7 @@ import org.jetbrains.kotlin.storage.StorageManager
 
 fun runCommonization(parameters: CommonizerParameters) {
     if (!parameters.hasAnythingToCommonize()) {
-        parameters.resultsConsumer.successfullyFinished(Status.NOTHING_TO_DO)
+        parameters.resultsConsumer.allConsumed(Status.NOTHING_TO_DO)
         return
     }
 
@@ -31,10 +31,9 @@ fun runCommonization(parameters: CommonizerParameters) {
     // build resulting declarations:
     for (targetIndex in 0 until mergedTree.dimension) {
         serializeTarget(mergeResult, targetIndex, parameters)
-        System.gc()
     }
 
-    parameters.resultsConsumer.successfullyFinished(Status.DONE)
+    parameters.resultsConsumer.allConsumed(Status.DONE)
 }
 
 private fun mergeAndCommonize(storageManager: StorageManager, parameters: CommonizerParameters): CirTreeMergeResult {
@@ -60,23 +59,25 @@ private fun mergeAndCommonize(storageManager: StorageManager, parameters: Common
 }
 
 private fun serializeTarget(mergeResult: CirTreeMergeResult, targetIndex: Int, parameters: CommonizerParameters) {
+    val mergedTree = mergeResult.root
+    val target = mergedTree.getTarget(targetIndex)
 
-    val (target, metadataModules) = MetadataBuilder.build(mergeResult.root, targetIndex, parameters.statsCollector)
+    MetadataBuilder.build(mergedTree, targetIndex, parameters.statsCollector) { metadataModule ->
-
-    val commonizedModules: List<ModuleResult.Commonized> = metadataModules.map { metadataModule ->
         val libraryName = metadataModule.name
         val serializedMetadata = with(metadataModule.write(KLIB_FRAGMENT_WRITE_STRATEGY)) {
             SerializedMetadata(header, fragments, fragmentNames)
         }
 
-        ModuleResult.Commonized(libraryName, serializedMetadata)
+        parameters.resultsConsumer.consume(target, ModuleResult.Commonized(libraryName, serializedMetadata))
     }
 
-    val missingModules: List<ModuleResult.Missing> = if (target is LeafTarget)
+    if (target is LeafTarget) {
-        mergeResult.missingModuleInfos.getValue(target).map { ModuleResult.Missing(it.originalLocation) }
+        mergeResult.missingModuleInfos.getValue(target).forEach {
-    else emptyList()
+            parameters.resultsConsumer.consume(target, ModuleResult.Missing(it.originalLocation))
+        }
+    }
 
-    parameters.resultsConsumer.consumeResults(target, commonizedModules + missingModules)
+    parameters.resultsConsumer.targetConsumed(target)
 }
 
 private val KLIB_FRAGMENT_WRITE_STRATEGY = ChunkedKlibModuleFragmentWriteStrategy()

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/konan/NativeDistributionModulesProvider.kt

@@ -23,8 +23,8 @@ internal class NativeDistributionModulesProvider(
     private val storageManager: StorageManager,
     private val librariesToCommonize: NativeLibrariesToCommonize
 ) : ModulesProvider {
-    override fun loadModuleInfos(): Map<String, ModuleInfo> {
+    private val moduleInfos: Map<String, ModuleInfo> by lazy {
-        return librariesToCommonize.libraries.associate { library ->
+        librariesToCommonize.libraries.associate { library ->
             val manifestData = library.manifestData
 
             val name = manifestData.uniqueName
@@ -42,6 +42,8 @@ internal class NativeDistributionModulesProvider(
         }
     }
 
+    override fun loadModuleInfos(): Map<String, ModuleInfo> = moduleInfos
+
     override fun loadModules(dependencies: Collection<ModuleDescriptor>): Map<String, ModuleDescriptor> {
         check(dependencies.isNotEmpty()) { "At least Kotlin/Native stdlib should be provided" }
 

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/konan/NativeDistributionResultsConsumer.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.descriptors.commonizer.konan
 
+import com.intellij.util.containers.FactoryMap
 import org.jetbrains.kotlin.descriptors.commonizer.ResultsConsumer
 import org.jetbrains.kotlin.descriptors.commonizer.ResultsConsumer.ModuleResult
 import org.jetbrains.kotlin.descriptors.commonizer.ResultsConsumer.Status
@@ -31,18 +32,44 @@ internal class NativeDistributionResultsConsumer(
     private val copyEndorsedLibs: Boolean,
     private val logProgress: (String) -> Unit
 ) : ResultsConsumer {
-    private val consumedTargets = LinkedHashSet<CommonizerTarget>()
     private val allLeafTargets = originalLibraries.librariesByTargets.keys
-    private val sharedTarget = SharedTarget(allLeafTargets)
 
-    override fun consumeResults(target: CommonizerTarget, moduleResults: Collection<ModuleResult>) {
+    private val allLeafCommonizedTargets = originalLibraries.librariesByTargets.filterValues { it.libraries.isNotEmpty() }.keys
-        val added = consumedTargets.add(target)
+    private val sharedTarget = SharedTarget(allLeafCommonizedTargets)
-        check(added)
 
-        serializeTarget(target, moduleResults)
+    private val consumedTargets = LinkedHashSet<CommonizerTarget>()
+
+    private val cachedManifestProviders = FactoryMap.create<CommonizerTarget, NativeManifestDataProvider> { target ->
+        when (target) {
+            is LeafTarget -> originalLibraries.librariesByTargets.getValue(target)
+            is SharedTarget -> CommonNativeManifestDataProvider(originalLibraries.librariesByTargets.values)
+        }
     }
 
-    override fun successfullyFinished(status: Status) {
+    private val cachedLibrariesDestination = FactoryMap.create<CommonizerTarget, File> { target ->
+        val librariesDestination = when (target) {
+            is LeafTarget -> destination.resolve(KONAN_DISTRIBUTION_PLATFORM_LIBS_DIR).resolve(target.name)
+            is SharedTarget -> destination.resolve(KONAN_DISTRIBUTION_COMMON_LIBS_DIR)
+        }
+
+        librariesDestination.mkdirs() // always create an empty directory even if there is nothing to copy
+        librariesDestination
+    }
+
+    override fun consume(target: CommonizerTarget, moduleResult: ModuleResult) {
+        check(target in allLeafCommonizedTargets || target == sharedTarget)
+        consumedTargets += target
+
+        serializeModule(target, moduleResult)
+    }
+
+    override fun targetConsumed(target: CommonizerTarget) {
+        check(target in consumedTargets)
+
+        logProgress("Written libraries for ${target.prettyCommonizedName(sharedTarget)}")
+    }
+
+    override fun allConsumed(status: Status) {
         // optimization: stdlib and endorsed libraries effectively remain the same across all Kotlin/Native targets,
         // so they can be just copied to the new destination without running serializer
         copyCommonStandardLibraries()
@@ -89,9 +116,7 @@ internal class NativeDistributionResultsConsumer(
 
     private fun copyTargetAsIs(leafTarget: LeafTarget) {
         val librariesCount = originalLibraries.librariesByTargets.getValue(leafTarget).libraries.size
-
+        val librariesDestination = cachedLibrariesDestination.getValue(leafTarget)
-        val librariesDestination = leafTarget.librariesDestination
-        librariesDestination.mkdirs() // always create an empty directory even if there is nothing to copy
 
         val librariesSource = leafTarget.platformLibrariesSource
         if (librariesSource.isDirectory) librariesSource.copyRecursively(librariesDestination)
@@ -99,35 +124,25 @@ internal class NativeDistributionResultsConsumer(
         logProgress("Copied $librariesCount libraries for ${leafTarget.prettyName}")
     }
 
-    private fun serializeTarget(target: CommonizerTarget, moduleResults: Collection<ModuleResult>) {
+    private fun serializeModule(target: CommonizerTarget, moduleResult: ModuleResult) {
-        val librariesDestination = target.librariesDestination
+        val librariesDestination = cachedLibrariesDestination.getValue(target)
-        librariesDestination.mkdirs() // always create an empty directory even if there is nothing to copy
 
-        val manifestProvider = when (target) {
+        when (moduleResult) {
-            is LeafTarget -> originalLibraries.librariesByTargets.getValue(target)
+            is ModuleResult.Commonized -> {
-            is SharedTarget -> CommonNativeManifestDataProvider(originalLibraries.librariesByTargets.values)
+                val libraryName = moduleResult.libraryName
-        }
 
-        for (moduleResult in moduleResults) {
+                val manifestData = cachedManifestProviders.getValue(target).getManifest(libraryName)
-            when (moduleResult) {
+                val libraryDestination = librariesDestination.resolve(libraryName)
-                is ModuleResult.Commonized -> {
-                    val libraryName = moduleResult.libraryName
 
-                    val manifestData = manifestProvider.getManifest(libraryName)
+                writeLibrary(moduleResult.metadata, manifestData, libraryDestination)
-                    val libraryDestination = librariesDestination.resolve(libraryName)
+            }
+            is ModuleResult.Missing -> {
+                val libraryName = moduleResult.libraryName
+                val missingModuleLocation = moduleResult.originalLocation
 
-                    writeLibrary(moduleResult.metadata, manifestData, libraryDestination)
+                missingModuleLocation.copyRecursively(librariesDestination.resolve(libraryName))
-                }
-                is ModuleResult.Missing -> {
-                    val libraryName = moduleResult.libraryName
-                    val missingModuleLocation = moduleResult.originalLocation
-
-                    missingModuleLocation.copyRecursively(librariesDestination.resolve(libraryName))
-                }
             }
         }
-
-        logProgress("Written libraries for ${target.prettyCommonizedName(sharedTarget)}")
     }
 
     private fun writeLibrary(
@@ -154,10 +169,4 @@ internal class NativeDistributionResultsConsumer(
         get() = repository.resolve(KONAN_DISTRIBUTION_KLIB_DIR)
             .resolve(KONAN_DISTRIBUTION_PLATFORM_LIBS_DIR)
             .resolve(name)
-
-    private val CommonizerTarget.librariesDestination: File
-        get() = when (this) {
-            is LeafTarget -> destination.resolve(KONAN_DISTRIBUTION_PLATFORM_LIBS_DIR).resolve(name)
-            is SharedTarget -> destination.resolve(KONAN_DISTRIBUTION_COMMON_LIBS_DIR)
-        }
 }

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/mergedtree/CirRootNode.kt

@@ -6,7 +6,9 @@
 package org.jetbrains.kotlin.descriptors.commonizer.mergedtree
 
 import gnu.trove.THashMap
+import org.jetbrains.kotlin.descriptors.commonizer.CommonizerTarget
 import org.jetbrains.kotlin.descriptors.commonizer.cir.CirRoot
+import org.jetbrains.kotlin.descriptors.commonizer.mergedtree.CirNode.Companion.indexOfCommon
 import org.jetbrains.kotlin.descriptors.commonizer.utils.CommonizedGroup
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.storage.NullableLazyValue
@@ -17,6 +19,9 @@ class CirRootNode(
 ) : CirNode<CirRoot, CirRoot> {
     val modules: MutableMap<Name, CirModuleNode> = THashMap()
 
+    fun getTarget(targetIndex: Int): CommonizerTarget =
+        (if (targetIndex == indexOfCommon) commonDeclaration() else targetDeclarations[targetIndex])!!.target
+
     override fun <T, R> accept(visitor: CirNodeVisitor<T, R>, data: T): R =
         visitor.visitRootNode(this, data)
 

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/metadata/buildingVisitor.kt

@@ -27,31 +27,35 @@ internal object MetadataBuilder {
     fun build(
         node: CirRootNode,
         targetIndex: Int,
-        statsCollector: StatsCollector?
+        statsCollector: StatsCollector?,
-    ): Pair<CommonizerTarget, Collection<KlibModuleMetadata>> {
+        moduleConsumer: (KlibModuleMetadata) -> Unit
-        return node.accept(
+    ) {
-            MetadataBuildingVisitor(statsCollector),
+        node.accept(
+            MetadataBuildingVisitor(statsCollector, moduleConsumer),
             MetadataBuildingVisitorContext.rootContext(node, targetIndex)
-        ).cast()
+        )
     }
 }
 
 @Suppress("PARAMETER_NAME_CHANGED_ON_OVERRIDE")
-private class MetadataBuildingVisitor(private val statsCollector: StatsCollector?) : CirNodeVisitor<MetadataBuildingVisitorContext, Any?> {
+private class MetadataBuildingVisitor(
+    private val statsCollector: StatsCollector?,
+    private val moduleConsumer: (KlibModuleMetadata) -> Unit
+) : CirNodeVisitor<MetadataBuildingVisitorContext, Any?> {
     private val classConsumer = ClassConsumer()
 
     override fun visitRootNode(
         node: CirRootNode,
         rootContext: MetadataBuildingVisitorContext
-    ): Pair<CommonizerTarget, Collection<KlibModuleMetadata>> {
+    ) {
-        val modules: List<KlibModuleMetadata> = node.modules.mapNotNull { (moduleName, moduleNode) ->
+        node.modules.forEach { (moduleName, moduleNode) ->
             val moduleContext = rootContext.moduleContext(moduleName)
-            val module: KlibModuleMetadata? = moduleNode.accept(this, moduleContext)?.cast()
+            val module: KlibModuleMetadata = moduleNode.accept(this, moduleContext)?.cast() ?: return@forEach
             statsCollector?.logModule(moduleContext)
-            module
+            moduleConsumer(module)
         }
 
-        return rootContext.target to modules
+        System.gc()
     }
 
     override fun visitModuleNode(
@@ -415,18 +419,14 @@ internal data class MetadataBuildingVisitorContext(
     }
 
     companion object {
-        fun rootContext(rootNode: CirRootNode, targetIndex: Int): MetadataBuildingVisitorContext {
+        fun rootContext(rootNode: CirRootNode, targetIndex: Int): MetadataBuildingVisitorContext =
-            val isCommon = rootNode.indexOfCommon == targetIndex
+            MetadataBuildingVisitorContext(
-            val target = (if (isCommon) rootNode.commonDeclaration() else rootNode.targetDeclarations[targetIndex])!!.target
-
-            return MetadataBuildingVisitorContext(
                 targetIndex = targetIndex,
-                target = target,
+                target = rootNode.getTarget(targetIndex),
-                isCommon = isCommon,
+                isCommon = rootNode.indexOfCommon == targetIndex,
                 typeParameterIndexOffset = 0,
                 currentPath = Path.Empty
             )
-        }
     }
 }
 

native/commonizer/tests/org/jetbrains/kotlin/descriptors/commonizer/CommonizerFacadeTest.kt

@@ -46,7 +46,7 @@ class CommonizerFacadeTest {
     )
 
     @Test
-    fun commonizedWithDifferentModules() = doTestSuccessfulCommonization(
+    fun commonizedWithDifferentModules() = doTestNothingToCommonize(
         mapOf(
             "target1" to listOf("foo"),
             "target2" to listOf("bar")

native/commonizer/tests/org/jetbrains/kotlin/descriptors/commonizer/utils/mocks.kt

@@ -177,22 +177,39 @@ internal class MockModulesProvider private constructor(
     }
 }
 
+private typealias ModuleName = String
+private typealias ModuleResults = HashMap<ModuleName, ModuleResult>
+
 internal class MockResultsConsumer : ResultsConsumer {
-    private val _modulesByTargets = LinkedHashMap<CommonizerTarget, Collection<ModuleResult>>() // use linked hash map to preserve order
+    private val _modulesByTargets = LinkedHashMap<CommonizerTarget, ModuleResults>() // use linked hash map to preserve order
     val modulesByTargets: Map<CommonizerTarget, Collection<ModuleResult>>
-        get() = _modulesByTargets
+        get() = _modulesByTargets.mapValues { it.value.values }
 
     val sharedTarget: SharedTarget by lazy { modulesByTargets.keys.filterIsInstance<SharedTarget>().single() }
     val leafTargets: Set<LeafTarget> by lazy { modulesByTargets.keys.filterIsInstance<LeafTarget>().toSet() }
 
+    private val finishedTargets = mutableSetOf<CommonizerTarget>()
+
     lateinit var status: ResultsConsumer.Status
 
-    override fun consumeResults(target: CommonizerTarget, moduleResults: Collection<ModuleResult>) {
+    override fun consume(target: CommonizerTarget, moduleResult: ModuleResult) {
-        val oldValue = _modulesByTargets.put(target, moduleResults)
+        check(!this::status.isInitialized)
-        check(oldValue == null)
+        check(target !in finishedTargets)
+        val moduleResults: ModuleResults = _modulesByTargets.getOrPut(target) { ModuleResults() }
+        val oldResult = moduleResults.put(moduleResult.libraryName, moduleResult)
+        check(oldResult == null) // to avoid accidental overwriting
     }
 
-    override fun successfullyFinished(status: ResultsConsumer.Status) {
+    override fun targetConsumed(target: CommonizerTarget) {
+        check(!this::status.isInitialized)
+        check(target in _modulesByTargets.keys)
+        check(target !in finishedTargets)
+        finishedTargets += target
+    }
+
+    override fun allConsumed(status: ResultsConsumer.Status) {
+        check(!this::status.isInitialized)
+        check(finishedTargets.containsAll(_modulesByTargets.keys))
         this.status = status
     }
 }