From c2acb77b84c39c56c7d7412fdfc57396b925b925 Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Fri, 31 Mar 2023 16:32:33 +0200
Subject: [PATCH] [LL FIR] KT-57220 LLFirSessionCache: Clean up symbol provider
 merging

- `SymbolProviderMerger` encapsulates some boilerplate, which should
  make `mergeDependencySymbolProvidersInto` easier to read and also much
  easier to extend.
---
 .../api/fir/sessions/LLFirSessionCache.kt     | 29 +++++++++++++------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
index ae64a13537b..cd913feeeab 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
@@ -537,18 +537,29 @@ internal class LLFirSessionCache(private val project: Project) {
         session: FirSession,
         destination: MutableList<FirSymbolProvider>,
     ) {
-        val (kotlinSymbolProviders, remainingSymbolProviders1) = partitionIsInstance<_, LLFirProvider.SymbolProvider>()
-        destination.addIfNotNull(LLFirCombinedKotlinSymbolProvider.merge(session, project, kotlinSymbolProviders))
+        SymbolProviderMerger(this, destination).apply {
+            merge<LLFirProvider.SymbolProvider> { LLFirCombinedKotlinSymbolProvider.merge(session, project, it) }
+            merge<FirExtensionSyntheticFunctionInterfaceProvider> { LLFirCombinedSyntheticFunctionSymbolProvider.merge(session, it) }
+            finish()
+        }
+    }
 
-        val (syntheticFunctionSymbolProviders, remainingSymbolProviders2) =
-            remainingSymbolProviders1.partitionIsInstance<_, FirExtensionSyntheticFunctionInterfaceProvider>()
+    private class SymbolProviderMerger(
+        symbolProviders: List<FirSymbolProvider>,
+        private val destination: MutableList<FirSymbolProvider>
+    ) {
+        private var remainingSymbolProviders = symbolProviders
 
-        destination.addAll(remainingSymbolProviders2)
+        inline fun <reified A : FirSymbolProvider> merge(create: (List<A>) -> FirSymbolProvider?) {
+            val (specificSymbolProviders, remainingSymbolProviders) = remainingSymbolProviders.partitionIsInstance<_, A>()
+            destination.addIfNotNull(create(specificSymbolProviders))
+            this.remainingSymbolProviders = remainingSymbolProviders
+        }
 
-        // Unfortunately, the functions that an extension synthetic function symbol provider might provide differ between sessions because
-        // they depend on compiler plugins. However, only extension providers that are affected by compiler plugins are added in
-        // `createSourcesSession`. We can still combine these, because the `ClassId` heuristics only need to be checked once.
-        destination.addIfNotNull(LLFirCombinedSyntheticFunctionSymbolProvider.merge(session, syntheticFunctionSymbolProviders))
+        fun finish() {
+            destination.addAll(remainingSymbolProviders)
+            remainingSymbolProviders = emptyList()
+        }
     }
 }