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

No need to retain all metadata modules in memory -> lesser memory consumption.
2021-02-03 12:57:11 +03:00
parent 53204661e3
commit 73113c1041
9 changed files with 137 additions and 81 deletions
@@ -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
+    }
 }
@@ -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)
 }
@@ -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)
-
-    val commonizedModules: List<ModuleResult.Commonized> = metadataModules.map { metadataModule ->
+    MetadataBuilder.build(mergedTree, targetIndex, parameters.statsCollector) { 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)
-        mergeResult.missingModuleInfos.getValue(target).map { ModuleResult.Missing(it.originalLocation) }
-    else emptyList()
+    if (target is LeafTarget) {
+        mergeResult.missingModuleInfos.getValue(target).forEach {
+            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()
@@ -23,8 +23,8 @@ internal class NativeDistributionModulesProvider(
    private val storageManager: StorageManager,
    private val librariesToCommonize: NativeLibrariesToCommonize
 ) : ModulesProvider {
-    override fun loadModuleInfos(): Map<String, ModuleInfo> {
-        return librariesToCommonize.libraries.associate { library ->
+    private val moduleInfos: Map<String, ModuleInfo> by lazy {
+        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" }

@@ -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>) {
-        val added = consumedTargets.add(target)
-        check(added)
+    private val allLeafCommonizedTargets = originalLibraries.librariesByTargets.filterValues { it.libraries.isNotEmpty() }.keys
+    private val sharedTarget = SharedTarget(allLeafCommonizedTargets)

-        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 = leafTarget.librariesDestination
-        librariesDestination.mkdirs() // always create an empty directory even if there is nothing to copy
+        val librariesDestination = cachedLibrariesDestination.getValue(leafTarget)

        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>) {
-        val librariesDestination = target.librariesDestination
-        librariesDestination.mkdirs() // always create an empty directory even if there is nothing to copy
+    private fun serializeModule(target: CommonizerTarget, moduleResult: ModuleResult) {
+        val librariesDestination = cachedLibrariesDestination.getValue(target)

-        val manifestProvider = when (target) {
-            is LeafTarget -> originalLibraries.librariesByTargets.getValue(target)
-            is SharedTarget -> CommonNativeManifestDataProvider(originalLibraries.librariesByTargets.values)
-        }
+        when (moduleResult) {
+            is ModuleResult.Commonized -> {
+                val libraryName = moduleResult.libraryName

-        for (moduleResult in moduleResults) {
-            when (moduleResult) {
-                is ModuleResult.Commonized -> {
-                    val libraryName = moduleResult.libraryName
+                val manifestData = cachedManifestProviders.getValue(target).getManifest(libraryName)
+                val libraryDestination = librariesDestination.resolve(libraryName)

-                    val manifestData = manifestProvider.getManifest(libraryName)
-                    val libraryDestination = librariesDestination.resolve(libraryName)
+                writeLibrary(moduleResult.metadata, manifestData, libraryDestination)
+            }
+            is ModuleResult.Missing -> {
+                val libraryName = moduleResult.libraryName
+                val missingModuleLocation = moduleResult.originalLocation

-                    writeLibrary(moduleResult.metadata, manifestData, libraryDestination)
-                }
-                is ModuleResult.Missing -> {
-                    val libraryName = moduleResult.libraryName
-                    val missingModuleLocation = moduleResult.originalLocation
-
-                    missingModuleLocation.copyRecursively(librariesDestination.resolve(libraryName))
-                }
+                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)
-        }
 }
@@ -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)

@@ -27,31 +27,35 @@ internal object MetadataBuilder {
    fun build(
        node: CirRootNode,
        targetIndex: Int,
-        statsCollector: StatsCollector?
-    ): Pair<CommonizerTarget, Collection<KlibModuleMetadata>> {
-        return node.accept(
-            MetadataBuildingVisitor(statsCollector),
+        statsCollector: StatsCollector?,
+        moduleConsumer: (KlibModuleMetadata) -> Unit
+    ) {
+        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 {
-            val isCommon = rootNode.indexOfCommon == targetIndex
-            val target = (if (isCommon) rootNode.commonDeclaration() else rootNode.targetDeclarations[targetIndex])!!.target
-
-            return MetadataBuildingVisitorContext(
+        fun rootContext(rootNode: CirRootNode, targetIndex: Int): MetadataBuildingVisitorContext =
+            MetadataBuildingVisitorContext(
                targetIndex = targetIndex,
-                target = target,
-                isCommon = isCommon,
+                target = rootNode.getTarget(targetIndex),
+                isCommon = rootNode.indexOfCommon == targetIndex,
                typeParameterIndexOffset = 0,
                currentPath = Path.Empty
            )
-        }
    }
 }

@@ -46,7 +46,7 @@ class CommonizerFacadeTest {
    )

    @Test
-    fun commonizedWithDifferentModules() = doTestSuccessfulCommonization(
+    fun commonizedWithDifferentModules() = doTestNothingToCommonize(
        mapOf(
            "target1" to listOf("foo"),
            "target2" to listOf("bar")
@@ -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>) {
-        val oldValue = _modulesByTargets.put(target, moduleResults)
-        check(oldValue == null)
+    override fun consume(target: CommonizerTarget, moduleResult: ModuleResult) {
+        check(!this::status.isInitialized)
+        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
    }
 }