From 48229f7faad87a8c4985dee41d4b35bce88a839a Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Wed, 28 Feb 2024 21:19:39 +0100
Subject: [PATCH] [AA Standalone] Implement `getRefinementDependents`

^KT-66013
---
 .../KtStaticModuleDependentsProvider.kt       | 37 +++++++++++++------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleDependentsProvider.kt b/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleDependentsProvider.kt
index f2cf864772b..29f57b8dbcd 100644
--- a/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleDependentsProvider.kt
+++ b/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleDependentsProvider.kt
@@ -15,23 +15,36 @@ class KtStaticModuleDependentsProvider(private val modules: List<KtModule>) : Ko
         // Direct dependencies should be computed lazily, because the built-ins module will be reachable via module dependencies. Getting
         // the built-ins module relies on the built-ins session, which may depend on services that are registered after
         // `KtStaticModuleDependentsProvider`.
-        buildMap<KtModule, MutableSet<KtModule>> {
-            for (module in modules) {
-                for (dependency in module.allDirectDependencies()) {
-                    // `module` should not be part of its own dependents, per the contract of `KotlinModuleDependentsProvider`.
-                    if (dependency == module) continue
-
-                    val dependents = computeIfAbsent(dependency) { mutableSetOf<KtModule>() }
-                    dependents.add(module)
-                }
-            }
-        }
+        buildDependentsMap(modules) { it.allDirectDependencies() }
     }
 
     private val transitiveDependentsByKtModule = createConcurrentSoftMap<KtModule, Set<KtModule>>()
 
-    override fun getDirectDependents(module: KtModule): Set<KtModule> = directDependentsByKtModule[module] ?: emptySet()
+    private val refinementDependentsByKtModule: Map<KtModule, Set<KtModule>> by lazy {
+        // Refinement dependents will usually only be requested for expect classes, so it's better to compute them lazily as they're not
+        // needed for all projects.
+        buildDependentsMap(modules) { it.transitiveDependsOnDependencies.asSequence() }
+    }
+
+    override fun getDirectDependents(module: KtModule): Set<KtModule> = directDependentsByKtModule[module].orEmpty()
 
     override fun getTransitiveDependents(module: KtModule): Set<KtModule> =
         transitiveDependentsByKtModule.computeIfAbsent(module) { computeTransitiveDependents(it) }
+
+    override fun getRefinementDependents(module: KtModule): Set<KtModule> = refinementDependentsByKtModule[module].orEmpty()
+}
+
+private inline fun buildDependentsMap(
+    modules: List<KtModule>,
+    getDependencies: (KtModule) -> Sequence<KtModule>,
+): Map<KtModule, MutableSet<KtModule>> = buildMap<KtModule, MutableSet<KtModule>> {
+    for (module in modules) {
+        for (dependency in getDependencies(module)) {
+            // `module` should not be part of its own dependents, per the contract of `KotlinModuleDependentsProvider`.
+            if (dependency == module) continue
+
+            val dependents = computeIfAbsent(dependency) { mutableSetOf() }
+            dependents.add(module)
+        }
+    }
 }