[FIR IDE] Fix local declarations resolve

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForResolveOnAir.kt

@@ -41,6 +41,7 @@ import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.idea.util.ifTrue
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 
 object LowLevelFirApiFacadeForResolveOnAir {
 
@@ -78,7 +79,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         val declaration = runBodyResolveOnAir(
             state = state,
             replacement = RawFirReplacement(place, elementToResolve),
-            resolveWithUnchangedFir = false
+            onAirCreatedDeclaration = true
         )
 
         val expressionLocator = object : FirVisitorVoid() {
@@ -114,7 +115,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
                 runBodyResolveOnAir(
                     state = state,
                     collector = it,
-                    resolveWithUnchangedFir = true,
+                    onAirCreatedDeclaration = false,
                     replacement = RawFirReplacement(validPlace, validPlace),
                 )
             }
@@ -174,7 +175,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         val copiedFirDeclaration = runBodyResolveOnAir(
             originalState,
             replacement = RawFirReplacement(sameDeclarationInOriginalFile, dependencyNonLocalDeclaration),
-            resolveWithUnchangedFir = false,
+            onAirCreatedDeclaration = true,
             collector = collector,
         )
 
@@ -210,7 +211,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
     private fun runBodyResolveOnAir(
         state: FirModuleResolveStateImpl,
         replacement: RawFirReplacement,
-        resolveWithUnchangedFir: Boolean,
+        onAirCreatedDeclaration: Boolean,
         collector: FirTowerDataContextCollector? = null,
     ): FirElement {
 
@@ -271,7 +272,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
                 designation = onAirDesignation,
                 moduleFileCache = state.rootModuleSession.cache,
                 checkPCE = true,
-                resolveWithUnchangedFir = resolveWithUnchangedFir,
+                onAirCreatedDeclaration = onAirCreatedDeclaration,
                 towerDataContextCollector = collector,
             )
             copiedFirDeclaration
@@ -281,18 +282,22 @@ object LowLevelFirApiFacadeForResolveOnAir {
 
     private fun isInBodyReplacement(ktDeclaration: KtDeclaration, replacement: RawFirReplacement): Boolean = when (ktDeclaration) {
         is KtNamedFunction ->
-            ktDeclaration.bodyBlockExpression?.let { PsiTreeUtil.isAncestor(it, replacement.from, true) } ?: false
+            ktDeclaration.bodyBlockExpression?.let { it.isAncestor(replacement.from, true) } ?: false
         is KtProperty -> {
             val insideGetterBody = ktDeclaration.getter?.bodyBlockExpression?.let {
-                PsiTreeUtil.isAncestor(it, replacement.from, true)
+                it.isAncestor(replacement.from, true)
             } ?: false
 
-            insideGetterBody || ktDeclaration.setter?.bodyBlockExpression?.let {
+            val insideGetterOrSetterBody = insideGetterBody || ktDeclaration.setter?.bodyBlockExpression?.let {
-                PsiTreeUtil.isAncestor(it, replacement.from, true)
+                it.isAncestor(replacement.from, true)
+            } ?: false
+
+            insideGetterOrSetterBody || ktDeclaration.initializer?.let {
+                it.isAncestor(replacement.from, true)
             } ?: false
         }
         is KtClassOrObject ->
-            ktDeclaration.body?.let { PsiTreeUtil.isAncestor(it, replacement.from, true) } ?: false
+            ktDeclaration.body?.let { it.isAncestor(replacement.from, true) } ?: false
         is KtTypeAlias -> false
         else -> error("Not supported type ${ktDeclaration::class.simpleName}")
     }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolveToType.kt

@@ -5,10 +5,17 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 
+import org.jetbrains.kotlin.fir.containingClass
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
+import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.runCustomResolveUnderLock
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.containingKtFileIfAny
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
 
 enum class ResolveType {
     NoResolve,
@@ -39,17 +46,40 @@ internal fun FirLazyDeclarationResolver.lazyResolveDeclaration(
             require(firDeclaration is FirCallableDeclaration<*>) {
                 "CallableReturnType type cannot be applied to ${firDeclaration::class.qualifiedName}"
             }
-            val stopAheadOfPhase = { declaration: FirDeclaration ->
+            //Need to be sync
-                declaration is FirCallableDeclaration<*> && declaration.returnTypeRef is FirResolvedTypeRef
+            if (firDeclaration.returnTypeRef is FirResolvedTypeRef) return
+            val containingFile = firDeclaration.getContainingFile()
+            if (containingFile != null) {
+                moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(containingFile, checkPCE) {
+                    if (firDeclaration.returnTypeRef !is FirResolvedTypeRef) {
+                        lazyResolveDeclaration(
+                            firDeclarationToResolve = firDeclaration,
+                            moduleFileCache = moduleFileCache,
+                            toPhase = FirResolvePhase.TYPES,
+                            scopeSession = scopeSession,
+                            checkPCE = checkPCE,
+                        )
+                    }
+                    if (firDeclaration.returnTypeRef !is FirResolvedTypeRef) {
+                        lazyResolveDeclaration(
+                            firDeclarationToResolve = firDeclaration,
+                            moduleFileCache = moduleFileCache,
+                            toPhase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+                            scopeSession = scopeSession,
+                            checkPCE = checkPCE,
+                        )
+                    }
+                    check(firDeclaration.returnTypeRef is FirResolvedTypeRef)
+                }
+            } else {
+                lazyResolveDeclaration(
+                    firDeclarationToResolve = firDeclaration,
+                    moduleFileCache = moduleFileCache,
+                    toPhase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+                    scopeSession = scopeSession,
+                    checkPCE = checkPCE,
+                )
             }
-            lazyResolveDeclaration(
-                firDeclarationToResolve = firDeclaration,
-                moduleFileCache = moduleFileCache,
-                toPhase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
-                scopeSession = scopeSession,
-                checkPCE = checkPCE,
-                stopAheadOfPhase = stopAheadOfPhase,
-            )
         }
         ResolveType.BodyResolveWithChildren, ResolveType.CallableBodyResolve -> {
             require(firDeclaration is FirCallableDeclaration<*> || toResolveType != ResolveType.CallableBodyResolve) {

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

@@ -8,12 +8,10 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
-import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirImportResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.runCustomResolveUnderLock
@@ -24,7 +22,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirProviderInter
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.LazyTransformerFactory
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhase
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
 import org.jetbrains.kotlin.idea.util.ifFalse
 import org.jetbrains.kotlin.idea.util.ifTrue
 
@@ -113,7 +110,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
         toPhase: FirResolvePhase,
         scopeSession: ScopeSession = ScopeSession(),
         checkPCE: Boolean = false,
-        stopAheadOfPhase: (FirDeclaration) -> Boolean = { false },
     ) {
         moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(firFile, checkPCE) {
             lazyResolveFileDeclarationWithoutLock(
@@ -122,7 +118,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                 toPhase = toPhase,
                 scopeSession = scopeSession,
                 checkPCE = checkPCE,
-                stopAheadOfPhase = stopAheadOfPhase,
             )
         }
     }
@@ -134,7 +129,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
         scopeSession: ScopeSession,
         checkPCE: Boolean = false,
         collector: FirTowerDataContextCollector? = null,
-        stopAheadOfPhase: (FirDeclaration) -> Boolean = { false },
     ) {
         if (toPhase == FirResolvePhase.RAW_FIR) return
         if (firFile.resolvePhase == FirResolvePhase.RAW_FIR) {
@@ -143,7 +137,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
         }
         if (checkPCE) checkCanceled()
         if (toPhase == FirResolvePhase.IMPORTS) return
-        if (stopAheadOfPhase(firFile)) return
 
         resolveFileAnnotationsWithoutLock(firFile, moduleFileCache, firFile.annotations, scopeSession, collector)
 
@@ -157,7 +150,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
             currentPhase = currentPhase.next
             if (currentPhase.pluginPhase) continue
             if (checkPCE) checkCanceled()
-            if (stopAheadOfPhase(firFile)) break
 
             val transformersToApply = designations.mapNotNull {
                 val needToResolve = it.resolvePhaseForAllDeclarations(includeDeclarationPhase = false) < currentPhase
@@ -198,7 +190,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
         toPhase: FirResolvePhase,
         checkPCE: Boolean,
         declarationPhaseDowngraded: Boolean = false,
-        stopAheadOfPhase: (FirDeclaration) -> Boolean = { false },
     ) {
         if (toPhase == FirResolvePhase.RAW_FIR) return
         //TODO Should be synchronised
@@ -211,32 +202,34 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                 toPhase = toPhase,
                 scopeSession = scopeSession,
                 checkPCE = checkPCE,
-                stopAheadOfPhase = stopAheadOfPhase,
             )
             return
         }
 
         val provider = firDeclarationToResolve.moduleData.session.firIdeProvider
-        val resolvableDeclaration = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache, firFileBuilder)
+        val (designation, wasInLocalDeclaration) =
-        //TODO Should be synchronised
+            firDeclarationToResolve.getNonLocalDeclarationToResolveAndInLocal(provider, moduleFileCache, firFileBuilder)
-        if (!resolvableDeclaration.isValidForResolve()) return
+
-        val containerFirFile = resolvableDeclaration.getContainingFile()
+        //TODO Should be synchronised
-            ?: error("FirFile was not found for\n${resolvableDeclaration.render()}")
+        if (!designation.declaration.isValidForResolve()) return
 
-        val designation = resolvableDeclaration.collectDesignation(containerFirFile)
         //TODO Should be synchronised
         val resolvePhase = designation.resolvePhaseForAllDeclarations(includeDeclarationPhase = declarationPhaseDowngraded)
-        if (resolvePhase >= toPhase) return
 
-        moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(containerFirFile, checkPCE) {
+        val neededPhase = if (wasInLocalDeclaration) {
+            if (toPhase >= FirResolvePhase.CONTRACTS) FirResolvePhase.BODY_RESOLVE else maxOf(FirResolvePhase.CONTRACTS, toPhase)
+        } else toPhase
+
+        if (resolvePhase >= neededPhase) return
+
+        moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(designation.firFile, checkPCE) {
             runLazyDesignatedResolveWithoutLock(
                 designation = designation,
                 moduleFileCache = moduleFileCache,
                 scopeSession = scopeSession,
-                toPhase = toPhase,
+                toPhase = neededPhase,
                 checkPCE = checkPCE,
                 declarationPhaseDowngraded = declarationPhaseDowngraded,
-                stopAheadOfPhase = stopAheadOfPhase,
             )
         }
     }
@@ -248,7 +241,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
         toPhase: FirResolvePhase,
         checkPCE: Boolean,
         declarationPhaseDowngraded: Boolean,
-        stopAheadOfPhase: (FirDeclaration) -> Boolean = { false },
     ) {
         val filePhase = designation.firFile.resolvePhase
         if (filePhase == FirResolvePhase.RAW_FIR) {
@@ -269,7 +261,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
             currentPhase = currentPhase.next
             if (currentPhase.pluginPhase) continue
             if (checkPCE) checkCanceled()
-            if (stopAheadOfPhase(designation.declaration)) break
 
             LazyTransformerFactory.createLazyTransformer(
                 phase = currentPhase,
@@ -289,13 +280,13 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
         designation: FirDeclarationDesignationWithFile,
         moduleFileCache: ModuleFileCache,
         checkPCE: Boolean,
-        resolveWithUnchangedFir: Boolean,
+        onAirCreatedDeclaration: Boolean,
         towerDataContextCollector: FirTowerDataContextCollector?,
     ) {
         val scopeSession = ScopeSession()
         var currentPhase = maxOf(designation.declaration.resolvePhase, FirResolvePhase.IMPORTS)
 
-        val firProviderInterceptor = resolveWithUnchangedFir.ifFalse {
+        val firProviderInterceptor = onAirCreatedDeclaration.ifTrue {
             FirProviderInterceptorForIDE.createForFirElement(
                 session = designation.firFile.moduleData.session,
                 firFile = designation.firFile,
@@ -321,8 +312,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
             ).transformDeclaration(firFileBuilder.firPhaseRunner)
         }
     }
-
-
 }
 
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/NonLocalDeclarationUtils.kt

@@ -12,45 +12,25 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 
-internal fun FirDeclaration.getNonLocalDeclarationToResolve(
+internal fun FirDeclaration.getNonLocalDeclarationToResolveAndInLocal(
     provider: FirProvider,
     moduleFileCache: ModuleFileCache,
     firFileBuilder: FirFileBuilder
-): FirDeclaration {
+): Pair<FirDeclarationDesignationWithFile, Boolean> {
-    if (this is FirFile) return this
+    require(this !is FirFile)
 
-    if (this is FirPropertyAccessor || this is FirTypeParameter || this is FirValueParameter) {
+    val nonLocalDesignation = tryCollectDesignationWithFile()
-        val ktContainingResolvableDeclaration = when (val psi = this.psi) {
+    if (nonLocalDesignation != null) {
-            is KtPropertyAccessor -> psi.property
+        return nonLocalDesignation to false
-            is KtProperty -> psi
-            is KtParameter, is KtTypeParameter -> psi.getNonLocalContainingOrThisDeclaration()
-                ?: error("Cannot find containing declaration for KtParameter")
-            is KtCallExpression -> {
-                check(this.source?.kind == FirFakeSourceElementKind.DefaultAccessor)
-                val delegationCall = psi.parent as KtPropertyDelegate
-                delegationCall.parent as KtProperty
-            }
-            null -> error("Cannot find containing declaration for KtParameter")
-            else -> error("Invalid source of property accessor ${psi::class}")
-        }
-
-        val targetElement =
-            if (declarationCanBeLazilyResolved(ktContainingResolvableDeclaration)) ktContainingResolvableDeclaration
-            else ktContainingResolvableDeclaration.getNonLocalContainingOrThisDeclaration()
-        check(targetElement != null) { "Container for local declaration cannot be null" }
-
-        return targetElement.findSourceNonLocalFirDeclaration(
-            firFileBuilder = firFileBuilder,
-            firSymbolProvider = moduleData.session.symbolProvider,
-            moduleFileCache = moduleFileCache
-        )
     }
 
     val ktDeclaration = (psi as? KtDeclaration) ?: run {
@@ -60,10 +40,34 @@ internal fun FirDeclaration.getNonLocalDeclarationToResolve(
         "FirDeclaration should have a PSI of type KtDeclaration"
     }
 
-    if (source !is FirFakeSourceElement<*> && declarationCanBeLazilyResolved(ktDeclaration)) return this
+    val declaration = when (this) {
-    val nonLocalPsi = ktDeclaration.getNonLocalContainingOrThisDeclaration()
+        is FirPropertyAccessor, is FirTypeParameter, is FirValueParameter -> {
+            when (ktDeclaration) {
+                is KtPropertyAccessor -> ktDeclaration.property
+                is KtProperty -> ktDeclaration
+                is KtParameter, is KtTypeParameter -> {
+                    val containingDeclaration = ktDeclaration.getParentOfType<KtDeclaration>(true)
+                    if (containingDeclaration !is KtPropertyAccessor) containingDeclaration else containingDeclaration.property
+                }
+                is KtCallExpression -> {
+                    check(this.source?.kind == FirFakeSourceElementKind.DefaultAccessor)
+                    ((ktDeclaration as? KtCallExpression)?.parent as? KtPropertyDelegate)?.parent as? KtProperty
+                }
+                else -> ktDeclaration
+            }
+        }
+        else -> ktDeclaration
+    }
+    check(declaration is KtDeclaration) {
+        "FirDeclaration should have a PSI of type KtDeclaration"
+    }
+
+    val nonLocalDeclaration = declaration.getNonLocalContainingOrThisDeclaration()
         ?: error("Container for local declaration cannot be null")
-    return nonLocalPsi.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache)
+
+    val firDeclaration = nonLocalDeclaration.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache)
+    val needUpgrade = nonLocalDeclaration !== declaration
+    return firDeclaration.collectDesignationWithFile() to needUpgrade
 }
 
 internal fun declarationCanBeLazilyResolved(declaration: KtDeclaration): Boolean {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt

@@ -15,9 +15,6 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
-import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
-import org.jetbrains.kotlin.idea.util.ifTrue
 import org.jetbrains.kotlin.psi.*
 
 internal class RawFirNonLocalDeclarationBuilder private constructor(