LL API: do not lazy resolve local declarations

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirModuleResolveStateDepended.kt

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.KtToFirMapping
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingKtFileIfAny
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalKtFile
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
@@ -62,11 +61,9 @@ internal class FirModuleResolveStateDepended(
     override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
         originalState.getOrBuildFirFile(ktFile)
 
-    override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D =
+    override fun resolveFirToPhase(declaration: FirDeclaration, toPhase: FirResolvePhase) {
         originalState.resolveFirToPhase(declaration, toPhase)
-
+    }
-    override fun <D : FirDeclaration> resolveFirToResolveType(declaration: D, type: ResolveType): D =
-        originalState.resolveFirToResolveType(declaration, type)
 
     override fun tryGetCachedFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {
         val ktFile = declaration.containingKtFileIfAny ?: return null

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirModuleResolveStateImpl.kt

@@ -16,8 +16,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.FirFileBuild
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FileStructureCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyDeclarationResolver
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveType
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.lazyResolveDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.firIdeProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSourcesSession
@@ -28,9 +26,9 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -73,7 +71,7 @@ internal class FirModuleResolveStateImpl(
         )
 
     override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
-        firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache, preferLazyBodies = false)
+        firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache)
 
     override fun tryGetCachedFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? =
         cache.getContainerFirFile(declaration)
@@ -139,13 +137,13 @@ internal class FirModuleResolveStateImpl(
         }
     }
 
-    override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
+    override fun resolveFirToPhase(declaration: FirDeclaration, toPhase: FirResolvePhase) {
-        if (toPhase == FirResolvePhase.RAW_FIR) return declaration
+        if (toPhase == FirResolvePhase.RAW_FIR) return
         val fileCache = when (val session = declaration.moduleData.session) {
             is FirIdeSourcesSession -> session.cache
-            else -> return declaration
+            else -> return
         }
-        return firLazyDeclarationResolver.lazyResolveDeclaration(
+        firLazyDeclarationResolver. lazyResolveDeclaration(
             firDeclarationToResolve = declaration,
             moduleFileCache = fileCache,
             scopeSession = ScopeSession(),
@@ -153,19 +151,4 @@ internal class FirModuleResolveStateImpl(
             checkPCE = true,
         )
     }
-
-    override fun <D : FirDeclaration> resolveFirToResolveType(declaration: D, type: ResolveType): D {
-        if (type == ResolveType.NoResolve) return declaration
-        val fileCache = when (val session = declaration.moduleData.session) {
-            is FirIdeSourcesSession -> session.cache
-            else -> return declaration
-        }
-        return firLazyDeclarationResolver.lazyResolveDeclaration(
-            firDeclaration = declaration,
-            moduleFileCache = fileCache,
-            scopeSession = ScopeSession(),
-            toResolveType = type,
-            checkPCE = true,
-        )
-    }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/IdeFirPhaseManager.kt

@@ -33,7 +33,6 @@ internal class IdeFirPhaseManager(
                 scopeSession = ScopeSession(),
                 toPhase = requiredPhase,
                 checkPCE = true,
-                skipLocalDeclaration = true,
             )
         } catch (e: Throwable) {
             sessionInvalidator.invalidate(fir.moduleData.session)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructure.kt

@@ -116,7 +116,7 @@ internal class FileStructure private constructor(
             moduleFileCache,
             firFile
         )
-        val resolvedDeclaration = firLazyDeclarationResolver.lazyResolveDeclaration(
+        firLazyDeclarationResolver.lazyResolveDeclaration(
             firDeclarationToResolve = firDeclaration,
             moduleFileCache = moduleFileCache,
             scopeSession = ScopeSession(),
@@ -124,7 +124,7 @@ internal class FileStructure private constructor(
             checkPCE = true,
         )
         return FileElementFactory.createFileStructureElement(
-            firDeclaration = resolvedDeclaration,
+            firDeclaration = firDeclaration,
             ktDeclaration = declaration,
             firFile = firFile,
             firFileLockProvider = moduleFileCache.firFileLockProvider
@@ -133,7 +133,7 @@ internal class FileStructure private constructor(
 
     private fun createStructureElement(container: KtAnnotated): FileStructureElement = when (container) {
         is KtFile -> {
-            val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, preferLazyBodies = true)
+            val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
             firLazyDeclarationResolver.resolveFileAnnotations(
                 firFile = firFile,
                 annotations = firFile.annotations,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt

@@ -140,16 +140,13 @@ internal class ReanalyzableFunctionStructureElement(
                 it.replaceResolvePhase(minOf(it.resolvePhase, upgradedPhase))
             }
 
-            val resolvedDeclaration = firLazyDeclarationResolver.lazyResolveDeclaration(
+            firLazyDeclarationResolver.lazyResolveDeclaration(
                 firDeclarationToResolve = originalFunction,
                 moduleFileCache = cache,
                 scopeSession = ScopeSession(),
                 toPhase = FirResolvePhase.BODY_RESOLVE,
                 checkPCE = true,
             )
-            check(resolvedDeclaration === originalFunction) {
-                "Reanalysed declaration not expected to be updated"
-            }
 
             ReanalyzableFunctionStructureElement(
                 firFile,
@@ -203,16 +200,14 @@ internal class ReanalyzablePropertyStructureElement(
                 replaceBodyResolveState(FirPropertyBodyResolveState.NOTHING_RESOLVED)
             }
 
-            val resolvedDeclaration = firLazyDeclarationResolver.lazyResolveDeclaration(
+            firLazyDeclarationResolver.lazyResolveDeclaration(
                 firDeclarationToResolve = originalProperty,
                 moduleFileCache = cache,
                 scopeSession = ScopeSession(),
                 toPhase = FirResolvePhase.BODY_RESOLVE,
                 checkPCE = true,
             )
-            check(resolvedDeclaration === originalProperty) {
+
-                "Reanalysed declaration not expected to be updated"
-            }
 
             ReanalyzablePropertyStructureElement(
                 firFile,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -5,6 +5,9 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
 
+import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.persistentListOf
+import kotlinx.collections.immutable.toPersistentList
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.firIdeProvider
 import org.jetbrains.kotlin.fir.FirElement
@@ -23,14 +26,14 @@ import org.jetbrains.kotlin.psi.KtSecondaryConstructor
 
 internal object FirLazyBodiesCalculator {
     fun calculateLazyBodiesInside(designation: FirDeclarationDesignation) {
-        designation.declaration.transform<FirElement, MutableList<FirDeclaration>>(
+        designation.declaration.transform<FirElement, PersistentList<FirDeclaration>>(
             FirLazyBodiesCalculatorTransformer,
-            designation.toSequence(includeTarget = true).toMutableList()
+            designation.toSequence(includeTarget = false).toList().toPersistentList()
         )
     }
 
     fun calculateLazyBodies(firFile: FirFile) {
-        firFile.transform<FirElement, MutableList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, mutableListOf())
+        firFile.transform<FirElement, PersistentList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, persistentListOf())
     }
 
     fun calculateLazyBodiesForFunction(designation: FirDeclarationDesignation) {
@@ -122,30 +125,29 @@ internal object FirLazyBodiesCalculator {
                 || firProperty.getExplicitBackingField()?.initializer is FirLazyExpression
 }
 
-private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<FirDeclaration>>() {
+private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentList<FirDeclaration>>() {
 
-    override fun transformFile(file: FirFile, data: MutableList<FirDeclaration>): FirFile {
+    override fun transformFile(file: FirFile, data: PersistentList<FirDeclaration>): FirFile {
         file.declarations.forEach {
             it.transformSingle(this, data)
         }
         return file
     }
 
-    override fun <E : FirElement> transformElement(element: E, data: MutableList<FirDeclaration>): E {
+    override fun <E : FirElement> transformElement(element: E, data: PersistentList<FirDeclaration>): E {
         if (element is FirRegularClass) {
-            data.add(element)
+            val newList = data.add(element)
             element.declarations.forEach {
-                it.transformSingle(this, data)
+                it.transformSingle(this, newList)
             }
-            element.transformChildren(this, data)
+            element.transformChildren(this, newList)
-            data.removeLast()
         }
         return element
     }
 
     override fun transformSimpleFunction(
         simpleFunction: FirSimpleFunction,
-        data: MutableList<FirDeclaration>
+        data: PersistentList<FirDeclaration>
     ): FirSimpleFunction {
         if (simpleFunction.body is FirLazyBlock) {
             val designation = FirDeclarationDesignation(data, simpleFunction)
@@ -156,7 +158,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<F
 
     override fun transformConstructor(
         constructor: FirConstructor,
-        data: MutableList<FirDeclaration>
+        data: PersistentList<FirDeclaration>
     ): FirConstructor {
         if (constructor.body is FirLazyBlock) {
             val designation = FirDeclarationDesignation(data, constructor)
@@ -165,7 +167,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<F
         return constructor
     }
 
-    override fun transformProperty(property: FirProperty, data: MutableList<FirDeclaration>): FirProperty {
+    override fun transformProperty(property: FirProperty, data: PersistentList<FirDeclaration>): FirProperty {
         if (FirLazyBodiesCalculator.needCalculatingLazyBodyForProperty(property)) {
             val designation = FirDeclarationDesignation(data, property)
             FirLazyBodiesCalculator.calculateLazyBodyForProperty(designation)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.firIdeProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.FirFileAnnotationsResolveTransformer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.FirProviderInterceptorForIDE
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LazyTransformerFactory
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.fir.FirElement
@@ -24,15 +23,13 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.FirImportResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
-import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtClassBody
 import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtEnumEntry
 
 internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBuilder) {
     /**
@@ -227,18 +224,17 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
      * @see lazyResolveDeclaration for ordinary resolve
      * @param firDeclarationToResolve target non-local declaration
      */
-    fun <D : FirDeclaration> lazyResolveDeclaration(
+    fun lazyResolveDeclaration(
-        firDeclarationToResolve: D,
+        firDeclarationToResolve: FirDeclaration,
         moduleFileCache: ModuleFileCache,
         scopeSession: ScopeSession,
         toPhase: FirResolvePhase,
         checkPCE: Boolean,
-        skipLocalDeclaration: Boolean = false,
+    ) {
-    ): D {
+        if (toPhase == FirResolvePhase.RAW_FIR) return
-        if (toPhase == FirResolvePhase.RAW_FIR) return firDeclarationToResolve
         if (toPhase == FirResolvePhase.IMPORTS) {
             if (fastTrackForImportsPhase(firDeclarationToResolve, moduleFileCache, checkPCE)) {
-                return firDeclarationToResolve
+                return
             }
         }
         when (firDeclarationToResolve) {
@@ -249,9 +245,8 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                     scopeSession,
                     toPhase,
                     checkPCE,
-                    skipLocalDeclaration
                 )
-                return firDeclarationToResolve
+                return
             }
 
             is FirBackingField -> {
@@ -261,7 +256,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                     scopeSession,
                     toPhase,
                     checkPCE,
-                    skipLocalDeclaration
                 )
             }
 
@@ -273,51 +267,61 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                     scopeSession = scopeSession,
                     checkPCE = checkPCE,
                 )
-                return firDeclarationToResolve
+                return
             }
+            else -> {}
         }
 
-        if (!firDeclarationToResolve.isValidForResolve()) return firDeclarationToResolve
+        if (!firDeclarationToResolve.isValidForResolve()) return
-        if (firDeclarationToResolve.resolvePhase >= toPhase) return firDeclarationToResolve
+        if (firDeclarationToResolve.resolvePhase >= toPhase) return
 
 
         val requestedDeclarationDesignation = firDeclarationToResolve.tryCollectDesignationWithFile()
 
         val designation: FirDeclarationDesignationWithFile
         val neededPhase: FirResolvePhase
-        val isLocalDeclarationResolveRequested: Boolean
+
         if (requestedDeclarationDesignation != null) {
             designation = requestedDeclarationDesignation
             neededPhase = toPhase
-            isLocalDeclarationResolveRequested = false
         } else {
             val possiblyLocalDeclaration = firDeclarationToResolve.getKtDeclarationForFirElement()
             val nonLocalDeclaration = possiblyLocalDeclaration.getNonLocalContainingOrThisDeclaration()
                 ?: error("Container for local declaration cannot be null")
-
+            val isLocalDeclarationResolveRequested =
-            isLocalDeclarationResolveRequested =
                 possiblyLocalDeclaration != nonLocalDeclaration
-            if (isLocalDeclarationResolveRequested && skipLocalDeclaration) return firDeclarationToResolve
 
-            val nonLocalFirDeclaration = nonLocalDeclaration.findSourceNonLocalFirDeclaration(
+            val declarationToResolve: FirDeclaration
-                firFileBuilder,
-                firDeclarationToResolve.moduleData.session.firIdeProvider.symbolProvider,
-                moduleFileCache
-            )
 
-            neededPhase = if (isLocalDeclarationResolveRequested) FirResolvePhase.BODY_RESOLVE else toPhase
+            if (isLocalDeclarationResolveRequested) {
+                val enumEntry = possiblyLocalDeclaration.getContainingEnumEntryAsMemberOfEnumEntry() ?: return
 
-            if (nonLocalFirDeclaration.resolvePhase >= neededPhase) return firDeclarationToResolve
+                declarationToResolve = enumEntry.findSourceNonLocalFirDeclaration(
-            if (!nonLocalFirDeclaration.isValidForResolve()) return firDeclarationToResolve
+                    firFileBuilder,
+                    firDeclarationToResolve.moduleData.session.firIdeProvider.symbolProvider,
+                    moduleFileCache
+                )
+                neededPhase = FirResolvePhase.BODY_RESOLVE
+            } else {
+                declarationToResolve = nonLocalDeclaration.findSourceNonLocalFirDeclaration(
+                    firFileBuilder,
+                    firDeclarationToResolve.moduleData.session.firIdeProvider.symbolProvider,
+                    moduleFileCache
+                )
+                neededPhase = toPhase
+            }
 
-            designation = nonLocalFirDeclaration.collectDesignationWithFile()
+            if (declarationToResolve.resolvePhase >= neededPhase) return
+            if (!declarationToResolve.isValidForResolve()) return
+
+            designation = declarationToResolve.collectDesignationWithFile()
         }
 
-        if (designation.declaration.resolvePhase >= neededPhase) return firDeclarationToResolve
+        if (designation.declaration.resolvePhase >= neededPhase) return
 
         if (neededPhase == FirResolvePhase.IMPORTS) {
             resolveFileToImports(designation.firFile, moduleFileCache, checkPCE)
-            return firDeclarationToResolve
+            return
         }
 
         moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(designation.firFile, checkPCE) {
@@ -332,33 +336,8 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                 designation.declaration
             }
         }
-
-        if (!isLocalDeclarationResolveRequested) return firDeclarationToResolve
-        return remapDeclarationInContainerIfNeeded(firDeclarationToResolve, designation.declaration, toPhase)
     }
 
-    private fun <D : FirDeclaration> remapDeclarationInContainerIfNeeded(
-        declarationToRemap: D,
-        firContainer: FirDeclaration,
-        firResolvePhase: FirResolvePhase
-    ): D {
-        if (declarationToRemap.resolvePhase >= firResolvePhase) return declarationToRemap
-        val realPsi = declarationToRemap.realPsi
-        check(realPsi != null) {
-            "Cannot remap element without PSI"
-        }
-        @Suppress("UNCHECKED_CAST")
-        val firDeclaration = FirElementFinder.findElementIn<FirDeclaration>(firContainer) {
-            it.realPsi == realPsi
-        } as? D
-        check(firDeclaration != null) {
-            "Containing declaration was resolved but local didn't found in it"
-        }
-        check(firDeclaration.resolvePhase >= firResolvePhase) {
-            "Found local declaration wasn't completely resolved"
-        }
-        return firDeclaration
-    }
 
     private fun runLazyDesignatedResolveWithoutLock(
         designation: FirDeclarationDesignationWithFile,
@@ -433,10 +412,8 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
     }
 }
 
-private fun KtDeclaration.isMemberOfEnumEntry(container: KtDeclaration) : Boolean {
+private fun KtDeclaration.getContainingEnumEntryAsMemberOfEnumEntry(): KtEnumEntry? {
-    if (container !is KtClass) return false
+    val body = parent as? KtClassBody ?: return null
-    if (!container.isEnum()) return false
+    return body.parent as? KtEnumEntry
-    val enumEntryBody = this.parent as? KtClassBody ?: return false
-    return enumEntryBody.parent == container
 }