From 1d0a696a628f8c0876c385a21f3d8180af77d4b6 Mon Sep 17 00:00:00 2001
From: Dmitry Savvinov <dmitry.savvinov@jetbrains.com>
Date: Wed, 16 Dec 2020 13:20:02 +0300
Subject: [PATCH] Simplify building of source-sets during MPP Import

Build dependsOn closure straight away by using KotlinSourceSetProto
instead of building first KotlinSourceSetImpls with incorrect
dependsOn-names, and then discarding it

^KT-37127
---
 .../src/KotlinMPPGradleModelBuilder.kt        | 54 ++++---------------
 .../src/KotlinMPPGradleModelImpl.kt           | 20 ++++++-
 2 files changed, 29 insertions(+), 45 deletions(-)

diff --git a/idea/kotlin-gradle-tooling/src/KotlinMPPGradleModelBuilder.kt b/idea/kotlin-gradle-tooling/src/KotlinMPPGradleModelBuilder.kt
index d2431961b54..5f7186d7b30 100644
--- a/idea/kotlin-gradle-tooling/src/KotlinMPPGradleModelBuilder.kt
+++ b/idea/kotlin-gradle-tooling/src/KotlinMPPGradleModelBuilder.kt
@@ -108,25 +108,6 @@ class KotlinMPPGradleModelBuilder : ModelBuilderService {
         return getCoroutines(experimentalExt) as? String
     }
 
-
-    private fun calculateDependsOnClosure(
-        sourceSet: KotlinSourceSetImpl?,
-        sourceSetsMap: Map<String, KotlinSourceSetImpl>,
-        cache: MutableMap<String, Set<String>>
-    ): Set<String> {
-        return if (sourceSet == null) {
-            emptySet()
-        } else {
-            cache[sourceSet.name] ?: sourceSet.dependsOnSourceSets.flatMap { name ->
-                calculateDependsOnClosure(
-                    sourceSetsMap[name],
-                    sourceSetsMap,
-                    cache
-                ).union(setOf(name))
-            }.toSet().also { cache[sourceSet.name] = it }
-        }
-    }
-
     private fun buildSourceSets(
         importingContext: MultiplatformModelImportingContext,
         dependencyResolver: DependencyResolver,
@@ -140,35 +121,22 @@ class KotlinMPPGradleModelBuilder : ModelBuilderService {
             (getSourceSets(kotlinExt) as? NamedDomainObjectContainer<Named>)?.asMap?.values ?: emptyList<Named>()
         val androidDeps = buildAndroidDeps(importingContext, kotlinExt.javaClass.classLoader)
 
-        val allSourceSetsProtos = sourceSets.mapNotNull {
+        val allSourceSetsProtosByNames = sourceSets.mapNotNull {
             buildSourceSet(it, dependencyResolver, importingContext.project, dependencyMapper, androidDeps)
-        }
+        }.associateBy { it.name }
+        val dependsOnCache = HashMap<String, Set<String>>()
 
         // Some performance optimisation: do not build metadata dependencies if source set is not common
-        val allSourceSets = if (importingContext.getProperty(BUILD_METADATA_DEPENDENCIES)) {
-            allSourceSetsProtos.map { proto -> proto.buildKotlinSourceSetImpl(true) }
+        return if (importingContext.getProperty(BUILD_METADATA_DEPENDENCIES)) {
+            allSourceSetsProtosByNames.mapValues { (_, proto) ->
+                proto.buildKotlinSourceSetImpl(true, allSourceSetsProtosByNames, dependsOnCache)
+            }
         } else {
-            val unactualizedSourceSets = allSourceSetsProtos.flatMap { it.dependsOnSourceSets }.distinct()
-            allSourceSetsProtos.map { proto -> proto.buildKotlinSourceSetImpl(unactualizedSourceSets.contains(proto.name)) }
+            val unactualizedSourceSets = allSourceSetsProtosByNames.values.flatMap { it.dependsOnSourceSets }.distinct()
+            allSourceSetsProtosByNames.mapValues { (name, proto) ->
+                proto.buildKotlinSourceSetImpl(unactualizedSourceSets.contains(name), allSourceSetsProtosByNames, dependsOnCache)
+            }
         }
-
-        // TODO: Weird quirk, we actually *recreate* all source-sets, because those contain
-        //       wrong (non-closured) dependsOnNames, so we can not save this map in the importingContext straight away
-        //       This will be fixed in later commit
-        val map = allSourceSets.map { it.name to it }.toMap()
-        val dependsOnCache = HashMap<String, Set<String>>()
-        return allSourceSets.map { sourceSet ->
-            sourceSet.name to KotlinSourceSetImpl(
-                sourceSet.name,
-                sourceSet.languageSettings,
-                sourceSet.sourceDirs,
-                sourceSet.resourceDirs,
-                sourceSet.dependencies,
-                calculateDependsOnClosure(sourceSet, map, dependsOnCache),
-                sourceSet.actualPlatforms as KotlinPlatformContainerImpl,
-                sourceSet.isTestModule
-            )
-        }.toMap()
     }
 
     private fun buildAndroidDeps(importingContext: MultiplatformModelImportingContext, classLoader: ClassLoader): Map<String, List<Any>>? {
diff --git a/idea/kotlin-gradle-tooling/src/KotlinMPPGradleModelImpl.kt b/idea/kotlin-gradle-tooling/src/KotlinMPPGradleModelImpl.kt
index de3b3572f08..3bde747398b 100644
--- a/idea/kotlin-gradle-tooling/src/KotlinMPPGradleModelImpl.kt
+++ b/idea/kotlin-gradle-tooling/src/KotlinMPPGradleModelImpl.kt
@@ -17,15 +17,31 @@ class KotlinSourceSetProto(
     val dependsOnSourceSets: Set<String>
 ) {
 
-    fun buildKotlinSourceSetImpl(doBuildDependencies: Boolean) = KotlinSourceSetImpl(
+    fun buildKotlinSourceSetImpl(
+        doBuildDependencies: Boolean,
+        allSourceSetsProtosByNames: Map<String, KotlinSourceSetProto>,
+        dependsOnCache: HashMap<String, Set<String>>
+    ) = KotlinSourceSetImpl(
         name,
         languageSettings,
         sourceDirs,
         resourceDirs,
         if (doBuildDependencies) dependencies.invoke() else emptyArray(),
-        dependsOnSourceSets
+        calculateDependsOnClosure(this, allSourceSetsProtosByNames, dependsOnCache)
     )
 
+    private fun calculateDependsOnClosure(
+        currentSourceSetProto: KotlinSourceSetProto,
+        sourceSetsMap: Map<String, KotlinSourceSetProto>,
+        cache: MutableMap<String, Set<String>>
+    ): Set<String> {
+        return cache.computeIfAbsent(currentSourceSetProto.name) {
+            currentSourceSetProto.dependsOnSourceSets.flatMap { name ->
+                val nextSourceSet = sourceSetsMap[name] ?: error("Source set $name is not found in map $sourceSetsMap")
+                calculateDependsOnClosure(nextSourceSet, sourceSetsMap, cache).union(setOf(name))
+            }.toSet()
+        }
+    }
 }
 
 class KotlinSourceSetImpl(