Add caching to dependency resolvers in PM20 + some refactoring

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/GradleProjectModuleBuilder.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.*
 import org.jetbrains.kotlin.gradle.plugin.sources.DefaultLanguageSettingsBuilder
 import org.jetbrains.kotlin.gradle.plugin.sources.KotlinDependencyScope
 import org.jetbrains.kotlin.gradle.plugin.sources.getVisibleSourceSetsFromAssociateCompilations
+import org.jetbrains.kotlin.gradle.utils.getOrPutRootProjectProperty
 import org.jetbrains.kotlin.project.model.*
 
 class ProjectStructureMetadataModuleBuilder {
@@ -260,22 +261,35 @@ internal fun Dependency.toModuleDependency(
 private fun BasicKotlinModule.fragmentByName(name: String) =
     fragments.single { it.fragmentName == name }
 
-class GradleModuleVariantResolver(val project: Project) : ModuleVariantResolver {
+class CachingModuleVariantResolver(private val actualResolver: ModuleVariantResolver) : ModuleVariantResolver {
-    private val resolvedVariantProvider = ResolvedMppVariantsProvider.get(project)
+    private val resultCacheByRequestingVariant: MutableMap<KotlinModuleVariant, MutableMap<KotlinModule, VariantResolution>> = mutableMapOf()
 
+    override fun getChosenVariant(requestingVariant: KotlinModuleVariant, dependencyModule: KotlinModule): VariantResolution {
+        val resultCache = resultCacheByRequestingVariant.getOrPut(requestingVariant) { mutableMapOf() }
+        return resultCache.getOrPut(dependencyModule) { actualResolver.getChosenVariant(requestingVariant, dependencyModule) }
+    }
+}
+
+class GradleModuleVariantResolver : ModuleVariantResolver {
     override fun getChosenVariant(requestingVariant: KotlinModuleVariant, dependencyModule: KotlinModule): VariantResolution {
         // TODO maybe improve this behavior? Currently it contradicts dependency resolution in that it may return a chosen variant for an
         //  unrequested dependency. This workaround is needed for synthetic modules which were not produced from module metadata, so maybe
         //  those modules should be marked somehow
-        if (dependencyModule is ExternalSyntheticKotlinModule) {
+        if (dependencyModule is ExternalPlainKotlinModule) {
             return VariantResolution.fromMatchingVariants(
                 requestingVariant,
                 dependencyModule,
-                listOf(dependencyModule.variants.single())
+                listOf(dependencyModule.singleVariant)
             )
         }
 
+        if (requestingVariant !is KotlinGradleVariant) {
+            return VariantResolution.Unknown(requestingVariant, dependencyModule)
+        }
+
         val module = requestingVariant.containingModule
+        val project = module.project
+        val resolvedVariantProvider = ResolvedMppVariantsProvider.get(project)
 
         // This implementation can only resolve variants for the current project's KotlinModule
         require(module.representsProject(project))
@@ -324,4 +338,13 @@ class GradleModuleVariantResolver(val project: Project) : ModuleVariantResolver
             }
             else -> error("could not find the compile dependencies configuration for variant $requestingVariant")
         }
+
+    companion object {
+        fun getForCurrentBuild(project: Project): ModuleVariantResolver {
+            val extraPropertyName = "org.jetbrains.kotlin.dependencyResolution.variantResolver.${project.getKotlinPluginVersion()}"
+            return project.getOrPutRootProjectProperty(extraPropertyName) {
+                CachingModuleVariantResolver(GradleModuleVariantResolver())
+            }
+        }
+    }
 }

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/GranularMetadataTransformation.kt

@@ -433,6 +433,7 @@ internal open class JarArtifactMppDependencyMetadataExtractor(
     val primaryArtifact: File
 ) : MppDependencyMetadataExtractor(project, dependency) {
 
+    // TODO: add proper API to make an artifact extractor with just the composite artifact "evolve" into one with host-specific artifacts
     val metadataArtifactBySourceSet: MutableMap<String, File> = mutableMapOf()
 
     private fun parseJsonProjectStructureMetadata(input: InputStream) =

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/ResolvedMppVariantsProvider.kt

@@ -17,19 +17,17 @@ import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.toModuleIdentifiers
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.toSingleModuleIdentifier
 import org.jetbrains.kotlin.gradle.plugin.usageByName
+import org.jetbrains.kotlin.gradle.utils.getOrPut
 import org.jetbrains.kotlin.project.model.KotlinModuleIdentifier
 import java.io.File
 
 internal class ResolvedMppVariantsProvider private constructor(private val project: Project) {
     companion object {
-
+        fun get(project: Project): ResolvedMppVariantsProvider {
-        fun get(project: Project): ResolvedMppVariantsProvider = with(project.extensions.extraProperties) {
             val propertyName = "kotlin.mpp.internal.resolvedModuleVariantsProvider"
-            if (!has(propertyName)) {
+            return project.extensions.extraProperties.getOrPut(propertyName) {
-                set(propertyName, ResolvedMppVariantsProvider(project))
+                ResolvedMppVariantsProvider(project)
             }
-            @Suppress("UNCHECKED_CAST")
-            get(propertyName) as ResolvedMppVariantsProvider
         }
     }
 

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/pm20/DependencyGraphResolver.kt

@@ -44,8 +44,6 @@ class GradleKotlinDependencyGraphResolver(
     }
 
     private fun resolveAsGraph(requestingModule: KotlinGradleModule): GradleDependencyGraph {
-        val excludeLegacyMetadataModulesFromResult = mutableSetOf<ResolvedComponentResult>()
-
         val nodeByModuleId = mutableMapOf<KotlinModuleIdentifier, GradleDependencyGraphNode>()
 
         fun getKotlinModuleFromComponentResult(component: ResolvedComponentResult): KotlinModule =
@@ -100,6 +98,8 @@ class GradleKotlinDependencyGraphResolver(
 class GradleDependencyGraphNode(
     override val module: KotlinModule,
     val selectedComponent: ResolvedComponentResult,
+    /** If the Kotlin module description was provided by a different component, such as with legacy publishing layout using *-metadata
+     * modules, then this property points to the other component. */
     val metadataSourceComponent: ResolvedComponentResult?,
     override val dependenciesByFragment: Map<KotlinModuleFragment, Iterable<GradleDependencyGraphNode>>
 ) : DependencyGraphNode(module, dependenciesByFragment)

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/pm20/DependencyResolution.kt

@@ -11,31 +11,69 @@ import org.gradle.api.artifacts.component.*
 import org.gradle.api.artifacts.result.ResolvedComponentResult
 import org.gradle.api.artifacts.result.ResolvedDependencyResult
 import org.gradle.api.capabilities.Capability
-import org.jetbrains.kotlin.gradle.dsl.kotlinExtension
+import org.jetbrains.kotlin.gradle.plugin.getKotlinPluginVersion
-import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
 import org.jetbrains.kotlin.gradle.plugin.mpp.*
 import org.jetbrains.kotlin.gradle.plugin.mpp.getProjectStructureMetadata
-import org.jetbrains.kotlin.gradle.plugin.mpp.resolvableMetadataConfiguration
 import org.jetbrains.kotlin.gradle.plugin.sources.DefaultLanguageSettingsBuilder
-import org.jetbrains.kotlin.gradle.plugin.sources.KotlinDependencyScope
+import org.jetbrains.kotlin.gradle.utils.getOrPut
+import org.jetbrains.kotlin.gradle.utils.getOrPutRootProjectProperty
 import org.jetbrains.kotlin.project.model.*
-import java.util.ArrayDeque
+import java.util.*
+
+class CachingModuleDependencyResolver(private val actualResolver: ModuleDependencyResolver) : ModuleDependencyResolver {
+    private val cacheByRequestingModule = WeakHashMap<KotlinModule, MutableMap<KotlinModuleDependency, KotlinModule?>>()
+
+    private fun cacheForRequestingModule(requestingModule: KotlinModule) =
+        cacheByRequestingModule.getOrPut(requestingModule) { mutableMapOf() }
+
+    override fun resolveDependency(requestingModule: KotlinModule, moduleDependency: KotlinModuleDependency): KotlinModule? =
+        cacheForRequestingModule(requestingModule).getOrPut(moduleDependency) {
+            actualResolver.resolveDependency(requestingModule, moduleDependency)
+        }
+}
+
+open class GradleComponentResultCachingResolver {
+    private val cachedResultsByRequestingModule = mutableMapOf<KotlinGradleModule, Map<KotlinModuleIdentifier, ResolvedComponentResult>>()
+
+    protected open fun configurationToResolve(requestingModule: KotlinGradleModule): Configuration =
+        configurationToResolveMetadataDependencies(requestingModule.project, requestingModule)
+
+    protected open fun resolveDependencies(module: KotlinGradleModule): Map<KotlinModuleIdentifier, ResolvedComponentResult> {
+        val allComponents = configurationToResolve(module).incoming.resolutionResult.allComponents
+        // FIXME handle multi-component results
+        return allComponents.flatMap { component -> component.toModuleIdentifiers().map { it to component } }.toMap()
+    }
+
+    private fun getResultsForModule(module: KotlinGradleModule): Map<KotlinModuleIdentifier, ResolvedComponentResult> =
+        cachedResultsByRequestingModule.getOrPut(module) { resolveDependencies(module) }
+
+    fun resolveModuleDependencyAsComponentResult(
+        requestingModule: KotlinGradleModule,
+        moduleDependency: KotlinModuleDependency
+    ): ResolvedComponentResult? =
+        getResultsForModule(requestingModule)[moduleDependency.moduleIdentifier]
+
+    companion object {
+        fun getForCurrentBuild(project: Project): GradleComponentResultCachingResolver {
+            val extraPropertyName = "org.jetbrains.kotlin.dependencyResolution.gradleComponentResolver.${project.getKotlinPluginVersion()}"
+            return project.getOrPutRootProjectProperty(extraPropertyName) {
+                GradleComponentResultCachingResolver()
+            }
+        }
+    }
+}
 
 class GradleModuleDependencyResolver(
+    private val gradleComponentResultResolver: GradleComponentResultCachingResolver,
     private val projectStructureMetadataModuleBuilder: ProjectStructureMetadataModuleBuilder,
     private val projectModuleBuilder: GradleProjectModuleBuilder
 ) : ModuleDependencyResolver {
 
-    private fun configurationToResolve(requestingModule: KotlinGradleModule): Configuration =
-        configurationToResolveMetadataDependencies(requestingModule.project, requestingModule)
-
     override fun resolveDependency(requestingModule: KotlinModule, moduleDependency: KotlinModuleDependency): KotlinModule? {
         require(requestingModule is KotlinGradleModule)
         val project = requestingModule.project
 
-        val allComponents = configurationToResolve(requestingModule).incoming.resolutionResult.allComponents
+        val component = gradleComponentResultResolver.resolveModuleDependencyAsComponentResult(requestingModule, moduleDependency)
-        // TODO: optimize O(n) search, store the resolved components with dependency keys?
-        val component = allComponents.find { it.id.matchesModuleDependency(moduleDependency) }
         val id = component?.id
 
         //FIXME multiple?
@@ -49,7 +87,8 @@ class GradleModuleDependencyResolver(
                 val metadata = getProjectStructureMetadata(
                     project,
                     component,
-                    configurationToResolve(requestingModule),
+                    // TODO: consistent choice of configurations across multiple resolvers?
+                    configurationToResolveMetadataDependencies(requestingModule.project, requestingModule),
                     moduleDependency.moduleIdentifier
                 ) ?: return null
                 val result = projectStructureMetadataModuleBuilder.getModule(component, metadata)
@@ -58,13 +97,30 @@ class GradleModuleDependencyResolver(
             else -> null
         }
     }
+
+    companion object {
+        fun getForCurrentBuild(project: Project): ModuleDependencyResolver {
+            val extraPropertyName = "org.jetbrains.kotlin.dependencyResolution.moduleResolver.${project.getKotlinPluginVersion()}"
+            return project.getOrPutRootProjectProperty(extraPropertyName) {
+                val componentResultResolver = GradleComponentResultCachingResolver.getForCurrentBuild(project)
+                val metadataModuleBuilder = ProjectStructureMetadataModuleBuilder()
+                val projectModuleBuilder = GradleProjectModuleBuilder(true)
+                CachingModuleDependencyResolver(
+                    GradleModuleDependencyResolver(componentResultResolver, metadataModuleBuilder, projectModuleBuilder)
+                )
+            }
+        }
+    }
 }
 
 // refactor extract to a separate class/interface
 // TODO think about multi-variant stub modules for non-Kotlin modules which got more than one chosen variant
-internal fun buildSyntheticModule(resolvedComponentResult: ResolvedComponentResult, singleVariantName: String): ExternalSyntheticKotlinModule {
+internal fun buildSyntheticPlainModule(
+    resolvedComponentResult: ResolvedComponentResult,
+    singleVariantName: String
+): ExternalPlainKotlinModule {
     val moduleDependency = resolvedComponentResult.toModuleDependency()
-    return ExternalSyntheticKotlinModule(BasicKotlinModule(moduleDependency.moduleIdentifier).apply {
+    return ExternalPlainKotlinModule(BasicKotlinModule(moduleDependency.moduleIdentifier).apply {
         BasicKotlinModuleVariant(this@apply, singleVariantName, DefaultLanguageSettingsBuilder()).apply {
             fragments.add(this)
             this.declaredModuleDependencies.addAll(
@@ -76,8 +132,12 @@ internal fun buildSyntheticModule(resolvedComponentResult: ResolvedComponentResu
     })
 }
 
-internal class ExternalSyntheticKotlinModule(private val moduleData: BasicKotlinModule) : KotlinModule by moduleData {
+internal class ExternalPlainKotlinModule(private val moduleData: BasicKotlinModule) : KotlinModule by moduleData {
-    override fun toString(): String = "synthetic $moduleData"
+    override fun toString(): String = "external plain $moduleData"
+
+    val singleVariant: KotlinModuleVariant
+        get() = moduleData.variants.singleOrNull()
+            ?: error("synthetic $moduleData was expected to have a single variant, got: ${moduleData.variants}")
 }
 
 internal class ExternalImportedKotlinModule(

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/pm20/FragmentGranularMetadataResolver.kt

@@ -31,10 +31,8 @@ internal class FragmentGranularMetadataResolver(
         refinesParentResolvers.value.flatMap { it.resolutions }.groupBy { it.dependency.toSingleModuleIdentifier() }
     }
 
-    private val metadataModuleBuilder = ProjectStructureMetadataModuleBuilder()
+    private val moduleResolver = GradleModuleDependencyResolver.getForCurrentBuild(project)
-    private val projectModuleBuilder = GradleProjectModuleBuilder(true)
+    private val variantResolver = GradleModuleVariantResolver.getForCurrentBuild(project)
-    private val moduleResolver = GradleModuleDependencyResolver(metadataModuleBuilder, projectModuleBuilder)
-    private val variantResolver = GradleModuleVariantResolver(project)
     private val fragmentResolver = DefaultModuleFragmentsResolver(variantResolver)
     private val dependencyGraphResolver = GradleKotlinDependencyGraphResolver(moduleResolver)
 
@@ -78,22 +76,12 @@ internal class FragmentGranularMetadataResolver(
             val resolvedComponentResult = dependencyNode.selectedComponent
             val isResolvedAsProject = resolvedComponentResult.toProjectOrNull(project)
             val result = when (dependencyModule) {
-                is ExternalSyntheticKotlinModule -> {
+                is ExternalPlainKotlinModule -> {
                     MetadataDependencyResolution.KeepOriginalDependency(resolvedComponentResult, isResolvedAsProject)
                 }
                 else -> run {
-                    val projectStructureMetadata = (dependencyModule as? ExternalImportedKotlinModule)?.projectStructureMetadata
-                        ?: checkNotNull(getProjectStructureMetadata(project, resolvedComponentResult, configurationToResolve))
-
                     val metadataSourceComponent = dependencyNode.run { metadataSourceComponent ?: selectedComponent }
 
-                    val parentResolutionsForDependency =
-                        parentResultsByModuleIdentifier[metadataSourceComponent.toSingleModuleIdentifier()].orEmpty()
-                    val fragmentsVisibleByParents =
-                        parentResolutionsForDependency.filterIsInstance<ChooseVisibleSourceSetsImpl>()
-                            .flatMapTo(mutableSetOf()) { it.allVisibleSourceSetNames }
-                    val visibleFragmentNames = visibleFragments.map { it.fragmentName }.toSet()
-
                     val metadataExtractor = getMetadataExtractor(project, resolvedComponentResult, configurationToResolve, true)
 
                     if (dependencyModule is ExternalImportedKotlinModule &&
@@ -103,12 +91,20 @@ internal class FragmentGranularMetadataResolver(
                         resolveHostSpecificMetadataArtifacts(dependencyModule, chosenFragments, metadataExtractor)
                     }
 
+                    val projectStructureMetadata = (dependencyModule as? ExternalImportedKotlinModule)?.projectStructureMetadata
+                        ?: checkNotNull(metadataExtractor?.getProjectStructureMetadata())
+
+                    val visibleFragmentNames = visibleFragments.map { it.fragmentName }.toSet()
+                    val visibleFragmentNamesExcludingVisibleByParents =
+                        visibleFragmentNames
+                            .minus(fragmentsNamesVisibleByParents(metadataSourceComponent.toSingleModuleIdentifier()))
+
                     ChooseVisibleSourceSetsImpl(
                         metadataSourceComponent,
                         isResolvedAsProject,
                         projectStructureMetadata,
                         visibleFragmentNames,
-                        visibleFragmentNames.minus(fragmentsVisibleByParents),
+                        visibleFragmentNamesExcludingVisibleByParents,
                         visibleTransitiveDependencies.map { resolvedDependenciesByModuleId.getValue(it.module.moduleIdentifier) }.toSet(),
                         checkNotNull(metadataExtractor)
                     )
@@ -126,6 +122,12 @@ internal class FragmentGranularMetadataResolver(
         return results
     }
 
+    private fun fragmentsNamesVisibleByParents(kotlinModuleIdentifier: KotlinModuleIdentifier): MutableSet<String> {
+        val parentResolutionsForDependency = parentResultsByModuleIdentifier[kotlinModuleIdentifier].orEmpty()
+        return parentResolutionsForDependency.filterIsInstance<ChooseVisibleSourceSetsImpl>()
+            .flatMapTo(mutableSetOf()) { it.allVisibleSourceSetNames }
+    }
+
     private fun resolveHostSpecificMetadataArtifacts(
         dependencyModule: ExternalImportedKotlinModule,
         chosenFragments: FragmentResolution.ChosenFragments,
@@ -135,19 +137,18 @@ internal class FragmentGranularMetadataResolver(
         val variantResolutions = chosenFragments.variantResolutions
         val hostSpecificFragments = dependencyModule.hostSpecificFragments
         val hostSpecificFragmentToArtifact = visibleFragments.intersect(hostSpecificFragments).mapNotNull { hostSpecificFragment ->
-            val relevantResolution = variantResolutions
+            val relevantVariantResolution = variantResolutions
                 .filterIsInstance<VariantResolution.VariantMatch>()
                 // find some of our variants that resolved a dependency's variant containing the fragment
                 .find { hostSpecificFragment in it.chosenVariant.refinesClosure }
             // resolve the dependencies of that variant getting the host-specific metadata artifact
-            relevantResolution?.let { resolution ->
+            relevantVariantResolution?.let { resolution ->
-                val variantResolvingPlatformArtifact =
+                val configurationResolvingPlatformVariant =
                     (resolution.requestingVariant as KotlinGradleVariant).compileDependencyConfiguration
-                val configuration = variantResolvingPlatformArtifact
                 val hostSpecificArtifact = ResolvedMppVariantsProvider.get(project)
                     .getHostSpecificMetadataArtifactByRootModule(
                         dependencyModule.moduleIdentifier,
-                        configuration
+                        configurationResolvingPlatformVariant
                     )
                 hostSpecificArtifact?.let { hostSpecificFragment.fragmentName to it }
             }

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/utils/extraProperties.kt

@@ -0,0 +1,21 @@
+/*
+ * 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.gradle.utils
+
+import org.gradle.api.Project
+import org.gradle.api.plugins.ExtraPropertiesExtension
+
+internal inline fun <reified T : Any> ExtraPropertiesExtension.getOrPut(key: String, provideValue: () -> T): T {
+    return synchronized(this) {
+        if (!has(key)) {
+            set(key, provideValue())
+        }
+        get(key) as T
+    }
+}
+
+internal inline fun <reified T : Any> Project.getOrPutRootProjectProperty(key: String, provideValue: () -> T) =
+    rootProject.extensions.extraProperties.getOrPut(key, provideValue)

libraries/tools/kotlin-project-model/src/main/kotlin/DependencyResolution.kt

@@ -13,7 +13,6 @@ interface ModuleDependencyResolver {
 }
 
 interface KotlinDependencyGraphResolver {
-    // TODO add explicit dependency consistency scopes if we decide to keep non-production code in the same variant
     fun resolveDependencyGraph(requestingModule: KotlinModule): DependencyGraphResolution
 }