[FIR IDE] Add typed lazy resolve implementation

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

@@ -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

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

@@ -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

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

@@ -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 ->

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt

@@ -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

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/LockProvider.kt

@@ -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")

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

@@ -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)
                 }
             }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureUtil.kt

@@ -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
     }
 

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

@@ -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,
+            )
+        }
+    }
+}

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

@@ -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.FirElement
-import org.jetbrains.kotlin.fir.*
 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.api.FirDeclarationDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
+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.idea.util.ifTrue
-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,
-        )
-    }
 
+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) {
+        resolveFileAnnotationsWithoutLock(firFile, moduleFileCache, firFile.annotations, scopeSession, collector)
-            resolveFileAnnotations(firFile, firFile.annotations, moduleFileCache, scopeSession, checkPCE, 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(
+            if (stopAheadOfPhase(firFile)) break
-                firDeclarationToResolve = declaration,
+
-                moduleFileCache = moduleFileCache,
+            val transformersToApply = designations.mapNotNull {
-                scopeSession = scopeSession,
+                val needToResolve = it.resolvePhaseForAllDeclarations(includeDeclarationPhase = false) < currentPhase
-                toPhase = toPhase,
+                needToResolve.ifTrue {
-                checkPCE = checkPCE
+                    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}")
-            }
-        }
-    }
 }
+
+

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

@@ -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}")
+    }
+}

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

@@ -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)
             }
         }

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

@@ -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")
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractFirLazyDeclarationResolveTest.kt

@@ -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" }

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractFirLibraryModuleDeclarationResolveTest.kt

@@ -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)
             }
 

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractPartialRawFirBuilderTestCase.kt

@@ -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

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirTypeAndAnnotations.kt

@@ -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)
     }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/annotations/KtFirAnnotationCall.kt

@@ -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())
         }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/annotations/firAnnotationUtils.kt

@@ -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
-}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/FirRefWithValidityCheck.kt

@@ -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