[FIR IDE] Fix local declarations resolve

2021-06-08 17:30:50 +02:00
parent 4a212dc6a0
commit e900a1fa57
5 changed files with 105 additions and 80 deletions
@@ -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 {
-                PsiTreeUtil.isAncestor(it, replacement.from, true)
+            val insideGetterOrSetterBody = insideGetterBody || ktDeclaration.setter?.bodyBlockExpression?.let {
+                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}")
    }
@@ -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 ->
-                declaration is FirCallableDeclaration<*> && declaration.returnTypeRef is FirResolvedTypeRef
+            //Need to be sync
+            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) {
@@ -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)
-        //TODO Should be synchronised
-        if (!resolvableDeclaration.isValidForResolve()) return
-        val containerFirFile = resolvableDeclaration.getContainingFile()
-            ?: error("FirFile was not found for\n${resolvableDeclaration.render()}")
+        val (designation, wasInLocalDeclaration) =
+            firDeclarationToResolve.getNonLocalDeclarationToResolveAndInLocal(provider, moduleFileCache, firFileBuilder)
+
+        //TODO Should be synchronised
+        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)
        }
    }
-
-
 }


@@ -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 {
-    if (this is FirFile) return this
+): Pair<FirDeclarationDesignationWithFile, Boolean> {
+    require(this !is FirFile)

-    if (this is FirPropertyAccessor || this is FirTypeParameter || this is FirValueParameter) {
-        val ktContainingResolvableDeclaration = when (val psi = this.psi) {
-            is KtPropertyAccessor -> psi.property
-            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 nonLocalDesignation = tryCollectDesignationWithFile()
+    if (nonLocalDesignation != null) {
+        return nonLocalDesignation to false
    }

    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 nonLocalPsi = ktDeclaration.getNonLocalContainingOrThisDeclaration()
+    val declaration = when (this) {
+        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 {
@@ -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(