From c5788290f3da32d7aef5f314768c4565fa26bf60 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Thu, 17 Dec 2020 17:47:20 +0100
Subject: [PATCH] FIR IDE: refactor: move designation & transformer creation in
 lazy resolve to separate functions

---
 .../resolve/FirLazyDeclarationResolver.kt     | 81 +++++++++++--------
 1 file changed, 49 insertions(+), 32 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt
index 2951a26d363..862a907b527 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt
@@ -149,45 +149,62 @@ internal class FirLazyDeclarationResolver(
     ) {
         val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
 
-        val designation = mutableListOf<FirDeclaration>(containerFirFile)
+        val designation = nonLocalDeclarationToResolve.getDesignation(containerFirFile, provider, moduleFileCache)
 
-        if (nonLocalDeclarationToResolve !is FirFile) {
-            val ktDeclaration = nonLocalDeclarationToResolve.ktDeclaration
-            designation += ktDeclaration.parentsOfType<KtClassOrObject>()
+        if (designation.all { it.resolvePhase >= phase }) {
+            return
+        }
+
+        val transformer = phase.createLazyTransformer(designation, containerFirFile, scopeSession, towerDataContextCollector)
+
+        firFileBuilder.firPhaseRunner.runPhaseWithCustomResolve(phase) {
+            containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
+        }
+    }
+
+    private fun FirResolvePhase.createLazyTransformer(
+        designation: List<FirDeclaration>,
+        containerFirFile: FirFile,
+        scopeSession: ScopeSession,
+        towerDataContextCollector: FirTowerDataContextCollector?
+    ) = when (this) {
+        FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
+            designation.iterator(),
+            containerFirFile.session,
+            scopeSession,
+        )
+        FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
+            designation.iterator(),
+            containerFirFile.session,
+            scopeSession,
+        )
+        FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
+            designation.iterator(),
+            containerFirFile.session,
+            scopeSession,
+            towerDataContextCollector
+        )
+        else -> error("Non-lazy phase $this")
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun FirDeclaration.getDesignation(
+        containerFirFile: FirFile,
+        provider: FirProvider,
+        moduleFileCache: ModuleFileCache
+    ): List<FirDeclaration> = buildList {
+        if (this !is FirFile) {
+            val ktDeclaration = ktDeclaration
+            ktDeclaration.parentsOfType<KtClassOrObject>(withSelf = true)
                 .filter { it !is KtEnumEntry }
                 .map { it.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache, containerFirFile) }
                 .toList()
                 .asReversed()
-            if (nonLocalDeclarationToResolve is FirCallableDeclaration<*>) {
-                designation += nonLocalDeclarationToResolve
+                .let(::addAll)
+            if (this@getDesignation is FirCallableDeclaration<*>) {
+                add(this@getDesignation)
             }
         }
-        if (designation.all { it.resolvePhase >= phase }) {
-            return
-        }
-        val transformer = when (phase) {
-            FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
-                designation.iterator(),
-                containerFirFile.session,
-                scopeSession,
-            )
-            FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
-                designation.iterator(),
-                containerFirFile.session,
-                scopeSession,
-            )
-            FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
-                designation.iterator(),
-                containerFirFile.session,
-                scopeSession,
-                towerDataContextCollector
-            )
-            else -> error("Non-lazy phase $phase")
-        }
-
-        firFileBuilder.firPhaseRunner.runPhaseWithCustomResolve(phase) {
-            containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
-        }
     }