[FIR IDE] Add typed lazy resolve implementation

2021-06-04 18:20:52 +02:00
parent 7a0391b9b3
commit eded837115
19 changed files with 449 additions and 222 deletions
@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.lazyResolveDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
@@ -135,6 +135,18 @@ fun <D : FirDeclaration, R> D.withFirDeclarationInWriteLock(
    return resolveState.withLock(this, DeclarationLockType.WRITE_LOCK, action)
 }

+/**
+ * Executes [action] with given [FirDeclaration] under write lock, so resolve **is possible** inside [action]
+ */
+fun <D : FirDeclaration, R> D.withFirDeclarationInWriteLock(
+    resolveState: FirModuleResolveState,
+    resolveType: ResolveType = ResolveType.BodyResolveWithChildren,
+    action: (D) -> R,
+): R {
+    resolvedFirToType(resolveType, resolveState)
+    return resolveState.withLock(this, DeclarationLockType.WRITE_LOCK, action)
+}
+
 /**
 * Returns a list of Diagnostics compiler finds for given [KtElement]
 * This operation could be performance affective because it create FIleStructureElement and resolve non-local declaration into BODY phase
@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.DeclarationCopyBuilder.wi
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FileTowerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextAllElementsCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.runCustomResolveUnderLock
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.KtToFirMapping
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
@@ -248,7 +249,7 @@ object LowLevelFirApiFacadeForResolveOnAir {

        val isInBodyReplacement = isInBodyReplacement(nonLocalDeclaration, replacement)

-        return FirLazyDeclarationResolver.runCustomResolveUnderLock(originalFirFile, state.rootModuleSession.cache, true) {
+        return state.rootModuleSession.cache.firFileLockProvider.runCustomResolveUnderLock(originalFirFile, true) {
            val copiedFirDeclaration = isInBodyReplacement.ifTrue {
                when (originalDeclaration) {
                    is FirSimpleFunction ->
@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureElement
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.declarationCanBeLazilyResolved
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.isNonAnonymousClassOrObject
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
@@ -110,7 +111,7 @@ internal inline fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate:
        if (container is KtNamedDeclaration
            && (container.isNonAnonymousClassOrObject() || container is KtDeclarationWithBody || container is KtProperty || container is KtTypeAlias)
            && container !is KtPrimaryConstructor
-            && FirLazyDeclarationResolver.declarationCanBeLazilyResolved(container)
+            && declarationCanBeLazilyResolved(container)
            && container !is KtEnumEntry
            && container !is KtFunctionLiteral
            && container.containingClassOrObject !is KtEnumEntry
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.builder

 import com.google.common.collect.MapMaker
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.PrivateForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DeclarationLockType
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.lockWithPCECheck
@@ -69,4 +70,19 @@ internal class DeadLockGuard {
    }
 }

+/**
+ * Runs [resolve] function (which is considered to do some resolve on [firFile]) under a lock for [firFile]
+ */
+internal inline fun <R> LockProvider<FirFile>.runCustomResolveUnderLock(
+    firFile: FirFile,
+    checkPCE: Boolean,
+    body: () -> R
+): R {
+    return if (checkPCE) {
+        withWriteLockPCECheck(key = firFile, lockingIntervalMs = 50L, action = body)
+    } else {
+        withWriteLock(key = firFile, action = body)
+    }
+}
+
 class ReadWriteDeadLockException : IllegalStateException("Acquiring write lock when read lock hold")
@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.LockProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirNonLocalDeclarationBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.declarationCanBeLazilyResolved
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.resolvePhaseForDeclarationAndChildren
@@ -231,14 +232,14 @@ internal class NonReanalyzableDeclarationStructureElement(
        private val recorder = object : FirElementsRecorder() {
            override fun visitProperty(property: FirProperty, data: MutableMap<KtElement, FirElement>) {
                val psi = property.psi as? KtProperty ?: return super.visitProperty(property, data)
-                if (!FileElementFactory.isReanalyzableContainer(psi) || !FirLazyDeclarationResolver.declarationCanBeLazilyResolved(psi)) {
+                if (!FileElementFactory.isReanalyzableContainer(psi) || !declarationCanBeLazilyResolved(psi)) {
                    super.visitProperty(property, data)
                }
            }

            override fun visitSimpleFunction(simpleFunction: FirSimpleFunction, data: MutableMap<KtElement, FirElement>) {
                val psi = simpleFunction.psi as? KtNamedFunction ?: return super.visitSimpleFunction(simpleFunction, data)
-                if (!FileElementFactory.isReanalyzableContainer(psi) || !FirLazyDeclarationResolver.declarationCanBeLazilyResolved(psi)) {
+                if (!FileElementFactory.isReanalyzableContainer(psi) || !declarationCanBeLazilyResolved(psi)) {
                    super.visitSimpleFunction(simpleFunction, data)
                }
            }
@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.declarationCanBeLazilyResolved
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.replaceFirst
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
@@ -23,7 +24,7 @@ internal object FileStructureUtil {
        ktDeclaration !is KtClassOrObject && ktDeclaration !is KtDeclarationWithBody && ktDeclaration !is KtProperty && ktDeclaration !is KtTypeAlias -> false
        ktDeclaration is KtEnumEntry -> false
        ktDeclaration.containingClassOrObject is KtEnumEntry -> false
-        ktDeclaration is KtNamedDeclaration -> !FirLazyDeclarationResolver.declarationCanBeLazilyResolved(ktDeclaration)
+        ktDeclaration is KtNamedDeclaration -> !declarationCanBeLazilyResolved(ktDeclaration)
        else -> false
    }

@@ -0,0 +1,135 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
+
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+
+enum class ResolveType {
+    NoResolve,
+    BodyResolveWithChildren,
+    CallableBodyResolve,
+    CallableReturnType,
+    AnnotationType,
+    AnnotationsArguments,
+    ClassSuperTypes,
+    CallableContracts,
+    DeclarationStatus,
+    ValueParametersTypes,
+    TypeParametersTypes,
+//    ResolveForMemberScope,
+//    ResolveForSuperMembers,
+}
+
+internal fun FirLazyDeclarationResolver.lazyResolveDeclaration(
+    firDeclaration: FirDeclaration,
+    moduleFileCache: ModuleFileCache,
+    toResolveType: ResolveType,
+    scopeSession: ScopeSession = ScopeSession(),
+    checkPCE: Boolean = false,
+) {
+    when (toResolveType) {
+        ResolveType.NoResolve -> return
+        ResolveType.CallableReturnType -> {
+            require(firDeclaration is FirCallableDeclaration<*>) {
+                "CallableReturn type cannot be applied to ${firDeclaration::class.qualifiedName}"
+            }
+            val stopAheadOfPhase = { declaration: FirDeclaration ->
+                declaration is FirCallableDeclaration<*> && declaration.returnTypeRef is FirResolvedTypeRef
+            }
+            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) {
+                "CallableReturn type cannot be applied to ${firDeclaration::class.qualifiedName}"
+            }
+            lazyResolveDeclaration(
+                firDeclarationToResolve = firDeclaration,
+                moduleFileCache = moduleFileCache,
+                toPhase = FirResolvePhase.BODY_RESOLVE,
+                scopeSession = scopeSession,
+                checkPCE = checkPCE,
+            )
+        }
+        ResolveType.AnnotationType, ResolveType.AnnotationsArguments -> {
+            if (firDeclaration is FirFile) {
+                resolveFileAnnotations(
+                    firFile = firDeclaration,
+                    annotations = firDeclaration.annotations,
+                    moduleFileCache = moduleFileCache,
+                    scopeSession = scopeSession,
+                    checkPCE = checkPCE
+                )
+            } else {
+                val toPhase = if (toResolveType == ResolveType.AnnotationType) FirResolvePhase.TYPES else FirResolvePhase.BODY_RESOLVE
+                lazyResolveDeclaration(
+                    firDeclarationToResolve = firDeclaration,
+                    moduleFileCache = moduleFileCache,
+                    toPhase = toPhase,
+                    scopeSession = scopeSession,
+                    checkPCE = checkPCE,
+                )
+            }
+        }
+        ResolveType.ClassSuperTypes -> {
+            require(firDeclaration is FirClassLikeDeclaration<*>) {
+                "CallableReturn type cannot be applied to ${firDeclaration::class.qualifiedName}"
+            }
+            lazyResolveDeclaration(
+                firDeclarationToResolve = firDeclaration,
+                moduleFileCache = moduleFileCache,
+                toPhase = FirResolvePhase.SUPER_TYPES,
+                scopeSession = scopeSession,
+                checkPCE = checkPCE,
+            )
+        }
+        ResolveType.CallableContracts -> {
+            require(firDeclaration is FirCallableDeclaration<*>) {
+                "CallableReturn type cannot be applied to ${firDeclaration::class.qualifiedName}"
+            }
+            lazyResolveDeclaration(
+                firDeclarationToResolve = firDeclaration,
+                moduleFileCache = moduleFileCache,
+                toPhase = FirResolvePhase.CONTRACTS,
+                scopeSession = scopeSession,
+                checkPCE = checkPCE,
+            )
+        }
+        ResolveType.DeclarationStatus -> {
+            require(firDeclaration !is FirFile) {
+                "CallableReturn type cannot be applied to ${firDeclaration::class.qualifiedName}"
+            }
+            lazyResolveDeclaration(
+                firDeclarationToResolve = firDeclaration,
+                moduleFileCache = moduleFileCache,
+                toPhase = FirResolvePhase.STATUS,
+                scopeSession = scopeSession,
+                checkPCE = checkPCE,
+            )
+        }
+        ResolveType.ValueParametersTypes, ResolveType.TypeParametersTypes -> {
+            require(firDeclaration !is FirFile) {
+                "CallableReturn type cannot be applied to ${firDeclaration::class.qualifiedName}"
+            }
+            lazyResolveDeclaration(
+                firDeclarationToResolve = firDeclaration,
+                moduleFileCache = moduleFileCache,
+                toPhase = FirResolvePhase.TYPES,
+                scopeSession = scopeSession,
+                checkPCE = checkPCE,
+            )
+        }
+    }
+}
@@ -5,64 +5,30 @@

 package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve

-import com.intellij.psi.util.parentOfType
-import org.jetbrains.kotlin.fir.*
+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.providers.FirProvider
-import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.FirImportResolveTransformer
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
-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.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
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirFileAnnotationsResolveTransformer
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.resolvePhaseForAllDeclarations
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirProviderInterceptorForIDE
+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.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
-
-enum class ResolveType {
-    FileAnnotations,
-    CallableReturnType,
-    ClassSuperTypes,
-    DeclarationStatus,
-    ValueParametersTypes,
-    TypeParametersTypes,
-    AnnotationType,
-    AnnotationParameters,
-    CallableBodyResolve,
-    ResolveForMemberScope,
-    ResolveForSuperMembers,
-    CallableContracts,
-    NoResolve,
-}
-
-internal class FirLazyDeclarationResolver(
-    private val firFileBuilder: FirFileBuilder
-) {
-    fun lazyResolveDeclaration(
-        firDeclaration: FirDeclaration,
-        moduleFileCache: ModuleFileCache,
-        toResolveType: ResolveType,
-        scopeSession: ScopeSession = ScopeSession(),
-        checkPCE: Boolean = false,
-    ) {
-        check(toResolveType == ResolveType.CallableReturnType)
-        lazyResolveDeclaration(
-            firDeclarationToResolve = firDeclaration,
-            moduleFileCache = moduleFileCache,
-            toPhase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
-            scopeSession = scopeSession,
-            checkPCE = checkPCE,
-        )
-    }
+import org.jetbrains.kotlin.idea.util.ifTrue

+internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBuilder) {
    /**
     * Fully resolve file annotations (synchronized)
     * @see resolveFileAnnotationsWithoutLock not synchronized
@@ -75,7 +41,7 @@ internal class FirLazyDeclarationResolver(
        checkPCE: Boolean,
        collector: FirTowerDataContextCollector? = null,
    ) {
-        runCustomResolveUnderLock(firFile, moduleFileCache, checkPCE) {
+        moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(firFile, checkPCE) {
            resolveFileAnnotationsWithoutLock(
                firFile = firFile,
                moduleFileCache = moduleFileCache,
@@ -147,14 +113,16 @@ internal class FirLazyDeclarationResolver(
        toPhase: FirResolvePhase,
        scopeSession: ScopeSession = ScopeSession(),
        checkPCE: Boolean = false,
+        stopAheadOfPhase: (FirDeclaration) -> Boolean = { false },
    ) {
-        runCustomResolveUnderLock(firFile, moduleFileCache, checkPCE) {
+        moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(firFile, checkPCE) {
            lazyResolveFileDeclarationWithoutLock(
                firFile = firFile,
                moduleFileCache = moduleFileCache,
                toPhase = toPhase,
                scopeSession = scopeSession,
                checkPCE = checkPCE,
+                stopAheadOfPhase = stopAheadOfPhase,
            )
        }
    }
@@ -166,27 +134,54 @@ internal class FirLazyDeclarationResolver(
        scopeSession: ScopeSession,
        checkPCE: Boolean = false,
        collector: FirTowerDataContextCollector? = null,
+        stopAheadOfPhase: (FirDeclaration) -> Boolean = { false },
    ) {
+        if (toPhase == FirResolvePhase.RAW_FIR) return
        if (firFile.resolvePhase == FirResolvePhase.RAW_FIR) {
            firFile.transform<FirElement, Any?>(FirImportResolveTransformer(firFile.moduleData.session), null)
            firFile.ensurePhase(FirResolvePhase.IMPORTS)
        }
        if (checkPCE) checkCanceled()
        if (toPhase == FirResolvePhase.IMPORTS) return
+        if (stopAheadOfPhase(firFile)) return

-        if (toPhase > FirResolvePhase.IMPORTS) {
-            resolveFileAnnotations(firFile, firFile.annotations, moduleFileCache, scopeSession, checkPCE, collector)
+        resolveFileAnnotationsWithoutLock(firFile, moduleFileCache, firFile.annotations, scopeSession, collector)
+
+        if (firFile.declarations.isEmpty()) return
+        val designations = firFile.declarations.map {
+            FirDeclarationDesignationWithFile(path = emptyList(), declaration = it, firFile = firFile)
        }

-        for (declaration in firFile.declarations) {
+        var currentPhase = FirResolvePhase.IMPORTS
+        while (currentPhase < toPhase) {
+            currentPhase = currentPhase.next
+            if (currentPhase.pluginPhase) continue
            if (checkPCE) checkCanceled()
-            lazyResolveDeclaration(
-                firDeclarationToResolve = declaration,
-                moduleFileCache = moduleFileCache,
-                scopeSession = scopeSession,
-                toPhase = toPhase,
-                checkPCE = checkPCE
-            )
+            if (stopAheadOfPhase(firFile)) break
+
+            val transformersToApply = designations.mapNotNull {
+                val needToResolve = it.resolvePhaseForAllDeclarations(includeDeclarationPhase = false) < currentPhase
+                needToResolve.ifTrue {
+                    LazyTransformerFactory.createLazyTransformer(
+                        phase = currentPhase,
+                        designation = it,
+                        scopeSession = scopeSession,
+                        declarationPhaseDowngraded = false,
+                        moduleFileCache = moduleFileCache,
+                        lazyDeclarationResolver = this,
+                        towerDataContextCollector = null,
+                        firProviderInterceptor = null,
+                        checkPCE = checkPCE,
+                    )
+                }
+            }
+            if (transformersToApply.isEmpty()) continue
+
+            firFileBuilder.firPhaseRunner.runPhaseWithCustomResolve(currentPhase) {
+                for (currentTransformer in transformersToApply) {
+                    currentTransformer.transformDeclaration(firFileBuilder.firPhaseRunner)
+                }
+            }
        }
    }

@@ -203,6 +198,7 @@ internal class FirLazyDeclarationResolver(
        toPhase: FirResolvePhase,
        checkPCE: Boolean,
        declarationPhaseDowngraded: Boolean = false,
+        stopAheadOfPhase: (FirDeclaration) -> Boolean = { false },
    ) {
        if (toPhase == FirResolvePhase.RAW_FIR) return
        //TODO Should be synchronised
@@ -215,12 +211,13 @@ internal class FirLazyDeclarationResolver(
                toPhase = toPhase,
                scopeSession = scopeSession,
                checkPCE = checkPCE,
+                stopAheadOfPhase = stopAheadOfPhase,
            )
            return
        }

        val provider = firDeclarationToResolve.moduleData.session.firIdeProvider
-        val resolvableDeclaration = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
+        val resolvableDeclaration = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache, firFileBuilder)
        //TODO Should be synchronised
        if (!resolvableDeclaration.isValidForResolve()) return
        val containerFirFile = resolvableDeclaration.getContainingFile()
@@ -231,7 +228,7 @@ internal class FirLazyDeclarationResolver(
        val resolvePhase = designation.resolvePhaseForAllDeclarations(includeDeclarationPhase = declarationPhaseDowngraded)
        if (resolvePhase >= toPhase) return

-        runCustomResolveUnderLock(containerFirFile, moduleFileCache, checkPCE) {
+        moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(containerFirFile, checkPCE) {
            runLazyDesignatedResolveWithoutLock(
                designation = designation,
                moduleFileCache = moduleFileCache,
@@ -239,6 +236,7 @@ internal class FirLazyDeclarationResolver(
                toPhase = toPhase,
                checkPCE = checkPCE,
                declarationPhaseDowngraded = declarationPhaseDowngraded,
+                stopAheadOfPhase = stopAheadOfPhase,
            )
        }
    }
@@ -250,6 +248,7 @@ internal class FirLazyDeclarationResolver(
        toPhase: FirResolvePhase,
        checkPCE: Boolean,
        declarationPhaseDowngraded: Boolean,
+        stopAheadOfPhase: (FirDeclaration) -> Boolean = { false },
    ) {
        val filePhase = designation.firFile.resolvePhase
        if (filePhase == FirResolvePhase.RAW_FIR) {
@@ -270,12 +269,15 @@ internal class FirLazyDeclarationResolver(
            currentPhase = currentPhase.next
            if (currentPhase.pluginPhase) continue
            if (checkPCE) checkCanceled()
+            if (stopAheadOfPhase(designation.declaration)) break

-            currentPhase.createLazyTransformer(
+            LazyTransformerFactory.createLazyTransformer(
+                phase = currentPhase,
                designation = designation,
                scopeSession = scopeSession,
                declarationPhaseDowngraded = declarationPhaseDowngraded,
                moduleFileCache = moduleFileCache,
+                lazyDeclarationResolver = this,
                towerDataContextCollector = null,
                firProviderInterceptor = null,
                checkPCE = checkPCE,
@@ -306,11 +308,13 @@ internal class FirLazyDeclarationResolver(
            if (currentPhase.pluginPhase) continue
            if (checkPCE) checkCanceled()

-            currentPhase.createLazyTransformer(
+            LazyTransformerFactory.createLazyTransformer(
+                phase = currentPhase,
                designation = designation,
                scopeSession = scopeSession,
                declarationPhaseDowngraded = true,
                moduleFileCache = moduleFileCache,
+                lazyDeclarationResolver = this,
                towerDataContextCollector = towerDataContextCollector,
                firProviderInterceptor = firProviderInterceptor,
                checkPCE = checkPCE,
@@ -318,135 +322,7 @@ internal class FirLazyDeclarationResolver(
        }
    }

-    private fun FirResolvePhase.createLazyTransformer(
-        designation: FirDeclarationDesignationWithFile,
-        scopeSession: ScopeSession,
-        declarationPhaseDowngraded: Boolean,
-        moduleFileCache: ModuleFileCache,
-        towerDataContextCollector: FirTowerDataContextCollector?,
-        firProviderInterceptor: FirProviderInterceptor?,
-        checkPCE: Boolean,
-    ): FirLazyTransformerForIDE = when (this) {
-        FirResolvePhase.SEALED_CLASS_INHERITORS -> FirLazyTransformerForIDE.DUMMY
-        FirResolvePhase.SUPER_TYPES -> FirDesignatedSupertypeResolverTransformerForIDE(
-            designation = designation,
-            session = designation.firFile.moduleData.session,
-            scopeSession = scopeSession,
-            declarationPhaseDowngraded = declarationPhaseDowngraded,
-            moduleFileCache = moduleFileCache,
-            firLazyDeclarationResolver = this@FirLazyDeclarationResolver,
-            firProviderInterceptor = firProviderInterceptor,
-            checkPCE = checkPCE,
-        )
-        FirResolvePhase.TYPES -> FirDesignatedTypeResolverTransformerForIDE(
-            designation,
-            designation.firFile.moduleData.session,
-            scopeSession,
-            declarationPhaseDowngraded,
-        )
-        FirResolvePhase.STATUS -> FirDesignatedStatusResolveTransformerForIDE(
-            designation,
-            designation.firFile.moduleData.session,
-            scopeSession,
-            declarationPhaseDowngraded,
-        )
-        FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
-            designation,
-            designation.firFile.moduleData.session,
-            scopeSession,
-            declarationPhaseDowngraded,
-        )
-        FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
-            designation,
-            designation.firFile.moduleData.session,
-            scopeSession,
-            declarationPhaseDowngraded,
-            towerDataContextCollector
-        )
-        FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
-            designation,
-            designation.firFile.moduleData.session,
-            scopeSession,
-            declarationPhaseDowngraded,
-            towerDataContextCollector,
-            firProviderInterceptor,
-        )
-        else -> error("Non-lazy phase $this")
-    }

-    private fun FirDeclaration.getNonLocalDeclarationToResolve(provider: FirProvider, moduleFileCache: ModuleFileCache): FirDeclaration {
-        if (this is FirFile) return this
-
-        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 ktDeclaration = (psi as? KtDeclaration) ?: run {
-            (source as? FirFakeSourceElement<*>).psi?.parentOfType()
-        }
-        check(ktDeclaration is KtDeclaration) {
-            "FirDeclaration should have a PSI of type KtDeclaration"
-        }
-
-        if (source !is FirFakeSourceElement<*> && declarationCanBeLazilyResolved(ktDeclaration)) return this
-        val nonLocalPsi = ktDeclaration.getNonLocalContainingOrThisDeclaration()
-            ?: error("Container for local declaration cannot be null")
-        return nonLocalPsi.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache)
-    }
-
-    companion object {
-        /**
-         * Runs [resolve] function (which is considered to do some resolve on [firFile]) under a lock for [firFile]
-         */
-        internal inline fun <R> runCustomResolveUnderLock(firFile: FirFile, cache: ModuleFileCache, checkPCE: Boolean, body: () -> R): R {
-            return if (checkPCE) {
-                cache.firFileLockProvider.withWriteLockPCECheck(firFile, LOCKING_INTERVAL_MS, body)
-            } else {
-                cache.firFileLockProvider.withWriteLock(firFile, body)
-            }
-        }
-
-        private const val LOCKING_INTERVAL_MS = 500L
-
-        fun declarationCanBeLazilyResolved(declaration: KtDeclaration): Boolean {
-            return when (declaration) {
-                !is KtNamedDeclaration -> false
-                is KtDestructuringDeclarationEntry, is KtFunctionLiteral, is KtTypeParameter -> false
-                is KtPrimaryConstructor -> false
-                is KtParameter -> declaration.hasValOrVar() && declaration.containingClassOrObject?.getClassId() != null
-                is KtCallableDeclaration, is KtEnumEntry -> {
-                    when (val parent = declaration.parent) {
-                        is KtFile -> true
-                        is KtClassBody -> (parent.parent as? KtClassOrObject)?.getClassId() != null
-                        else -> false
-                    }
-                }
-                is KtClassLikeDeclaration -> declaration.getClassId() != null
-                else -> error("Unexpected ${declaration::class.qualifiedName}")
-            }
-        }
-    }
 }
+
+
@@ -0,0 +1,85 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
+
+import com.intellij.psi.util.parentOfType
+import org.jetbrains.kotlin.fir.FirFakeSourceElement
+import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+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.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
+
+internal fun FirDeclaration.getNonLocalDeclarationToResolve(
+    provider: FirProvider,
+    moduleFileCache: ModuleFileCache,
+    firFileBuilder: FirFileBuilder
+): FirDeclaration {
+    if (this is FirFile) return this
+
+    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 ktDeclaration = (psi as? KtDeclaration) ?: run {
+        (source as? FirFakeSourceElement<*>).psi?.parentOfType()
+    }
+    check(ktDeclaration is KtDeclaration) {
+        "FirDeclaration should have a PSI of type KtDeclaration"
+    }
+
+    if (source !is FirFakeSourceElement<*> && declarationCanBeLazilyResolved(ktDeclaration)) return this
+    val nonLocalPsi = ktDeclaration.getNonLocalContainingOrThisDeclaration()
+        ?: error("Container for local declaration cannot be null")
+    return nonLocalPsi.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache)
+}
+
+internal fun declarationCanBeLazilyResolved(declaration: KtDeclaration): Boolean {
+    return when (declaration) {
+        !is KtNamedDeclaration -> false
+        is KtDestructuringDeclarationEntry, is KtFunctionLiteral, is KtTypeParameter -> false
+        is KtPrimaryConstructor -> false
+        is KtParameter -> declaration.hasValOrVar() && declaration.containingClassOrObject?.getClassId() != null
+        is KtCallableDeclaration, is KtEnumEntry -> {
+            when (val parent = declaration.parent) {
+                is KtFile -> true
+                is KtClassBody -> (parent.parent as? KtClassOrObject)?.getClassId() != null
+                else -> false
+            }
+        }
+        is KtClassLikeDeclaration -> declaration.getClassId() != null
+        else -> error("Unexpected ${declaration::class.qualifiedName}")
+    }
+}
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.java.declarations.FirJavaClass
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
@@ -18,8 +19,8 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 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.ModuleFileCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.runCustomResolveUnderLock
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
-import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver.Companion.runCustomResolveUnderLock
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.isResolvedForAllDeclarations
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.updateResolvedPhaseForDeclarationAndChildren
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
@@ -96,19 +97,19 @@ internal class FirDesignatedSupertypeResolverTransformerForIDE(
    private fun collect(designation: FirDeclarationDesignationWithFile): Collection<FirDeclarationDesignationWithFile> {
        val visited = mutableMapOf<FirDeclaration, FirDeclarationDesignationWithFile>()
        val toVisit = mutableListOf<FirDeclarationDesignationWithFile>()
-        toVisit.add(designation)

+        toVisit.add(designation)
        while (toVisit.isNotEmpty()) {
            for (nowVisit in toVisit) {
                if (checkPCE) checkCanceled()
                val resolver = DesignatedFirSupertypeResolverVisitor(nowVisit)
-                runCustomResolveUnderLock(nowVisit.firFile, moduleFileCache, checkPCE) {
+                moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(nowVisit.firFile, checkPCE) {
                    firLazyDeclarationResolver.lazyResolveFileDeclaration(
                        firFile = nowVisit.firFile,
                        moduleFileCache = moduleFileCache,
                        toPhase = FirResolvePhase.IMPORTS,
                        scopeSession = scopeSession,
-                        checkPCE = false,
+                        checkPCE = true,
                    )
                    nowVisit.firFile.accept(resolver, null)
                }
@@ -122,6 +123,7 @@ internal class FirDesignatedSupertypeResolverTransformerForIDE(
                for (reference in value.supertypeRefs) {
                    val classLikeDeclaration = reference.type.toSymbol(session)?.fir
                    if (classLikeDeclaration !is FirClassLikeDeclaration<*>) continue
+                    if (classLikeDeclaration is FirJavaClass) continue
                    if (visited.containsKey(classLikeDeclaration)) continue
                    val containingFile = moduleFileCache.getContainerFirFile(classLikeDeclaration) ?: continue
                    toVisit.add(classLikeDeclaration.collectDesignation(containingFile))
@@ -144,7 +146,7 @@ internal class FirDesignatedSupertypeResolverTransformerForIDE(
        val filesToDesignations = visited.groupBy { it.firFile }
        for (designationsPerFile in filesToDesignations) {
            if (checkPCE) checkCanceled()
-            runCustomResolveUnderLock(designationsPerFile.key, moduleFileCache, checkPCE) {
+            moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(designationsPerFile.key, checkPCE) {
                applyToFileSymbols(designationsPerFile.value)
            }
        }
@@ -0,0 +1,74 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
+
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
+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.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+
+internal object LazyTransformerFactory {
+    fun createLazyTransformer(
+        phase: FirResolvePhase,
+        designation: FirDeclarationDesignationWithFile,
+        scopeSession: ScopeSession,
+        declarationPhaseDowngraded: Boolean,
+        moduleFileCache: ModuleFileCache,
+        lazyDeclarationResolver: FirLazyDeclarationResolver,
+        towerDataContextCollector: FirTowerDataContextCollector?,
+        firProviderInterceptor: FirProviderInterceptor?,
+        checkPCE: Boolean,
+    ): FirLazyTransformerForIDE = when (phase) {
+        FirResolvePhase.SEALED_CLASS_INHERITORS -> FirLazyTransformerForIDE.DUMMY
+        FirResolvePhase.SUPER_TYPES -> FirDesignatedSupertypeResolverTransformerForIDE(
+            designation = designation,
+            session = designation.firFile.moduleData.session,
+            scopeSession = scopeSession,
+            declarationPhaseDowngraded = declarationPhaseDowngraded,
+            moduleFileCache = moduleFileCache,
+            firLazyDeclarationResolver = lazyDeclarationResolver,
+            firProviderInterceptor = firProviderInterceptor,
+            checkPCE = checkPCE,
+        )
+        FirResolvePhase.TYPES -> FirDesignatedTypeResolverTransformerForIDE(
+            designation,
+            designation.firFile.moduleData.session,
+            scopeSession,
+            declarationPhaseDowngraded,
+        )
+        FirResolvePhase.STATUS -> FirDesignatedStatusResolveTransformerForIDE(
+            designation,
+            designation.firFile.moduleData.session,
+            scopeSession,
+            declarationPhaseDowngraded,
+        )
+        FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
+            designation,
+            designation.firFile.moduleData.session,
+            scopeSession,
+            declarationPhaseDowngraded,
+        )
+        FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
+            designation,
+            designation.firFile.moduleData.session,
+            scopeSession,
+            declarationPhaseDowngraded,
+            towerDataContextCollector
+        )
+        FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
+            designation,
+            designation.firFile.moduleData.session,
+            scopeSession,
+            declarationPhaseDowngraded,
+            towerDataContextCollector,
+            firProviderInterceptor,
+        )
+        else -> error("Non-lazy phase $this")
+    }
+}
@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.declarationCanBeLazilyResolved
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.psi.KtDeclaration
@@ -32,7 +33,7 @@ abstract class AbstractFirLazyDeclarationResolveTest : KotlinLightCodeInsightFix
        val testDataFile = File(path)
        val ktFile = myFixture.configureByText(testDataFile.name, FileUtil.loadFile(testDataFile)) as KtFile
        val lazyDeclarations = ktFile.collectDescendantsOfType<KtDeclaration> { ktDeclaration ->
-            FirLazyDeclarationResolver.declarationCanBeLazilyResolved(ktDeclaration)
+            declarationCanBeLazilyResolved(ktDeclaration)
        }

        val declarationToResolve = lazyDeclarations.firstOrNull { it.name?.lowercase() == "resolveme" }
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api

 import com.intellij.testFramework.LightProjectDescriptor
 import org.jetbrains.kotlin.fir.FirRenderer
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclaration
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
@@ -48,7 +49,7 @@ abstract class AbstractFirLibraryModuleDeclarationResolveTest : KotlinLightCodeI
        // We intentionally use ktFile here as a context element, because resolving
        // from compiled PSI-elements (e.g. caretResolutionTarget) is not yet supported
        resolveWithClearCaches(ktFile) { resolveState ->
-            val renderedDeclaration = caretResolutionTarget.withFirDeclaration(resolveState) { firDeclaration ->
+            val renderedDeclaration = caretResolutionTarget.withFirDeclaration(resolveState, FirResolvePhase.RAW_FIR) { firDeclaration ->
                firDeclaration.render(FirRenderer.RenderMode.WithResolvePhases)
            }

@@ -17,10 +17,7 @@ import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
 import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirNonLocalDeclarationBuilder
-import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtAnnotationCall
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtTypeAndAnnotations
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
+import org.jetbrains.kotlin.idea.util.ifTrue

 internal class KtFirTypeAndAnnotations<T : FirDeclaration>(
    private val containingDeclaration: FirRefWithValidityCheck<T>,
@@ -64,10 +65,10 @@ internal fun FirRefWithValidityCheck<FirClass<*>>.superTypesAndAnnotationsList(b

 internal fun FirRefWithValidityCheck<FirRegularClass>.superTypesAndAnnotationsListForRegularClass(builder: KtSymbolByFirBuilder): List<KtTypeAndAnnotations> {
    return withFir { fir ->
-        if (fir.resolvePhase >= FirResolvePhase.SUPER_TYPES) {
+        (fir.resolvePhase >= FirResolvePhase.SUPER_TYPES).ifTrue {
            fir.superTypeRefs.mapToTypeAndAnnotations(this, builder)
-        } else null
-    } ?: withFirWithPossibleResolveInside { fir ->
+        }
+    } ?: withFirWithPossibleResolveInside(ResolveType.NoResolve) { fir ->
        fir.resolveSupertypesInTheAir(builder.rootSession).mapToTypeAndAnnotations(this, builder)
    }
 }
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRefWithValidityCheck
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.convertConstantExpression
@@ -35,14 +36,14 @@ internal class KtFirAnnotationCall(
    }

    override val classId: ClassId? by cached {
-        containingDeclaration.withFirWithPossibleResolveInside(AnnotationPhases.PHASE_FOR_ANNOTATION_CLASS_ID) { fir ->
+        containingDeclaration.withFirWithPossibleResolveInside(ResolveType.AnnotationType) { fir ->
            annotationCallRef.getClassId(fir.moduleData.session)
        }
    }

    override val useSiteTarget: AnnotationUseSiteTarget? get() = annotationCallRef.useSiteTarget

-    override val arguments: List<KtNamedConstantValue> by containingDeclaration.withFirAndCache(FirResolvePhase.TYPES) { fir ->
+    override val arguments: List<KtNamedConstantValue> by containingDeclaration.withFirAndCache(ResolveType.AnnotationsArguments) { fir ->
        mapAnnotationParameters(annotationCallRef, fir.moduleData.session).map { (name, expression) ->
            KtNamedConstantValue(name, expression.convertConstantExpression())
        }
@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.expressions.coneClassLikeType
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.transformers.ensureResolved
 import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRefWithValidityCheck
 import org.jetbrains.kotlin.name.ClassId

@@ -29,16 +30,11 @@ internal fun FirRefWithValidityCheck<FirAnnotatedDeclaration>.toAnnotationsList(
 }

 internal fun FirRefWithValidityCheck<FirAnnotatedDeclaration>.containsAnnotation(classId: ClassId): Boolean =
-    withFir(AnnotationPhases.PHASE_FOR_ANNOTATION_CLASS_ID) { fir ->
+    withFir(ResolveType.AnnotationType) { fir ->
        fir.annotations.any { it.getClassId(fir.moduleData.session) == classId }
    }

 internal fun FirRefWithValidityCheck<FirAnnotatedDeclaration>.getAnnotationClassIds(): Collection<ClassId> =
-    withFir(AnnotationPhases.PHASE_FOR_ANNOTATION_CLASS_ID) { fir ->
+    withFir(ResolveType.AnnotationType) { fir ->
        fir.annotations.mapNotNull { it.getClassId(fir.moduleData.session) }
-    }
-
-
-internal object AnnotationPhases {
-    val PHASE_FOR_ANNOTATION_CLASS_ID = FirResolvePhase.TYPES
-}
+    }
@@ -80,6 +80,31 @@ internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveSt
        }
    }

+    /**
+     * Runs [action] with fir element with write action hold
+     * Consider using this then [action] may call some resolve
+     */
+    inline fun <R> withFirWithPossibleResolveInside(
+        resolveType: ResolveType = ResolveType.NoResolve,
+        crossinline action: (fir: D) -> R
+    ): R {
+        token.assertIsValidAndAccessible()
+        val fir = firWeakRef.get()
+            ?: throw EntityWasGarbageCollectedException("FirElement")
+        val resolveState = resolveStateWeakRef.get()
+            ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
+        return when (resolveType) {
+            ResolveType.BodyResolveWithChildren -> {
+                /*
+                 The ResolveType type is the maximum possible phase we can resolve our declaration to
+                 So there is not need to run whole `action` under write lock
+                 */
+                action(fir.withFirDeclarationInWriteLock(resolveState, resolveType) { it })
+            }
+            else -> fir.withFirDeclarationInWriteLock(resolveState, resolveType) { action(it) }
+        }
+    }
+
    val resolveState
        get() = resolveStateWeakRef.get() ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")

@@ -95,7 +120,7 @@ internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveSt
        val resolveState = resolveStateWeakRef.get()
            ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
        return when (type) {
-            ResolveType.CallableBodyResolve -> {
+            ResolveType.BodyResolveWithChildren -> {
                /*
                 The CallableBodyResolve type is the maximum possible phase we can resolve our declaration to
                 So there is not need to run whole `action` under read lock