[FIR IDE] Implement lazy designated resolve for all resolve FIR phases

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -56,6 +56,8 @@ open class RawFirBuilder(
 
     private val stubMode get() = mode == RawFirBuilderMode.STUBS
 
+    protected open fun bindFunctionTarget(target: FirFunctionTarget, function: FirFunction<*>) = target.bind(function)
+
     var mode: RawFirBuilderMode = builderMode
         private set
 
@@ -375,7 +377,7 @@ open class RawFirBuilder(
                         }
                     }.also {
                         it.initContainingClassAttr()
-                        accessorTarget.bind(it)
+                        bindFunctionTarget(accessorTarget, it)
                         this@RawFirBuilder.context.firFunctionTargets.removeLast()
                     }
                 }
@@ -1072,7 +1074,7 @@ open class RawFirBuilder(
                 }
                 context.firFunctionTargets.removeLast()
             }.build().also {
-                target.bind(it)
+                bindFunctionTarget(target, it)
                 if (it is FirSimpleFunction) {
                     function.fillDanglingConstraintsTo(it)
                 }
@@ -1180,7 +1182,7 @@ open class RawFirBuilder(
                 }
                 context.firFunctionTargets.removeLast()
             }.also {
-                target.bind(it)
+                bindFunctionTarget(target, it)
             }
         }
 
@@ -1220,7 +1222,7 @@ open class RawFirBuilder(
                 this@RawFirBuilder.context.firFunctionTargets.removeLast()
             }.also {
                 it.containingClassAttr = currentDispatchReceiverType()!!.lookupTag
-                target.bind(it)
+                bindFunctionTarget(target, it)
             }
         }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/TypeExpansionUtils.kt

@@ -6,9 +6,11 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.declarations.expandedConeType
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
+import org.jetbrains.kotlin.fir.resolve.transformers.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirAbstractPhaseTransformer.kt

@@ -27,14 +27,18 @@ abstract class FirAbstractPhaseTransformer<D>(
 
     override fun transformFile(file: FirFile, data: D): FirFile {
         checkSessionConsistency(file)
-        file.replaceResolvePhase(transformerPhase)
+        if (file.resolvePhase < transformerPhase) {
+            file.replaceResolvePhase(transformerPhase)
+        }
 
         @Suppress("UNCHECKED_CAST")
         return super.transformFile(file, data) as FirFile
     }
 
     override fun transformDeclaration(declaration: FirDeclaration, data: D): FirDeclaration {
-        declaration.replaceResolvePhase(transformerPhase)
+        if (declaration.resolvePhase < transformerPhase) {
+            declaration.replaceResolvePhase(transformerPhase)
+        }
 
         return super.transformDeclaration(declaration, data)
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirAbstractTreeTransformerWithSuperTypes.kt

@@ -33,6 +33,8 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
     protected val scopes = mutableListOf<FirScope>()
     protected val towerScope = FirCompositeScope(scopes.asReversed())
 
+    protected open fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = transformerPhase > firDeclaration.resolvePhase
+
     protected inline fun <T> withScopeCleanup(crossinline l: () -> T): T {
         val sizeBefore = scopes.size
         val result = l()
@@ -48,7 +50,9 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
         firClass: FirClass<*>,
         data: Any?
     ): FirStatement {
-        firClass.replaceResolvePhase(transformerPhase)
+        if (needReplacePhase(firClass)) {
+            firClass.replaceResolvePhase(transformerPhase)
+        }
         return withScopeCleanup {
             // Otherwise annotations may try to resolve
             // themselves as inner classes of the `firClass`

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolveTransformer.kt

@@ -67,7 +67,7 @@ abstract class ResolvedStatusCalculator {
     }
 }
 
-class FirStatusResolveTransformer(
+open class FirStatusResolveTransformer(
     session: FirSession,
     scopeSession: ScopeSession,
     statusComputationSession: StatusComputationSession,
@@ -253,22 +253,13 @@ abstract class AbstractFirStatusResolveTransformer(
     protected abstract fun FirDeclaration.needResolveMembers(): Boolean
     protected abstract fun FirDeclaration.needResolveNestedClassifiers(): Boolean
 
+    protected open fun needReplacePhase(firDeclaration: FirDeclaration) = transformerPhase > firDeclaration.resolvePhase
+
     override fun transformFile(file: FirFile, data: FirResolvedDeclarationStatus?): FirFile {
-        file.replaceResolvePhase(transformerPhase)
-        if (file.needResolveMembers()) {
-            for (declaration in file.declarations) {
-                if (declaration !is FirClassLikeDeclaration<*>) {
-                    declaration.transformSingle(this, data)
-                }
-            }
-        }
-        if (file.needResolveNestedClassifiers()) {
-            for (declaration in file.declarations) {
-                if (declaration is FirClassLikeDeclaration<*>) {
-                    declaration.transformSingle(this, data)
-                }
-            }
+        if (needReplacePhase(file)) {
+            file.replaceResolvePhase(transformerPhase)
         }
+        transformDeclarationContent(file, data)
         return file
     }
 
@@ -293,7 +284,9 @@ abstract class AbstractFirStatusResolveTransformer(
         declaration: FirDeclaration,
         data: FirResolvedDeclarationStatus?
     ): FirDeclaration {
-        declaration.replaceResolvePhase(transformerPhase)
+        if (needReplacePhase(declaration)) {
+            declaration.replaceResolvePhase(transformerPhase)
+        }
         return when (declaration) {
             is FirCallableDeclaration<*> -> {
                 when (declaration) {
@@ -333,6 +326,36 @@ abstract class AbstractFirStatusResolveTransformer(
         return transformClass(anonymousObject, data)
     }
 
+    open fun transformDeclarationContent(
+        declaration: FirDeclaration,
+        data: FirResolvedDeclarationStatus?
+    ): FirDeclaration {
+
+        val declarations = when (declaration) {
+            is FirRegularClass -> declaration.declarations
+            is FirAnonymousObject -> declaration.declarations
+            is FirFile -> declaration.declarations
+            else -> error("Not supported declaration ${declaration::class.simpleName}")
+        }
+
+        if (declaration.needResolveMembers()) {
+            val members = declarations.filter { it !is FirClassLikeDeclaration<*> }
+            members.forEach { member ->
+                if (needReplacePhase(member)) {
+                    member.replaceResolvePhase(transformerPhase)
+                }
+            }
+            members.forEach { member -> member.transformSingle(this, data) }
+        }
+        if (declaration.needResolveNestedClassifiers()) {
+            val members = declarations.filterIsInstance<FirClassLikeDeclaration<*>>()
+            for (klass in members) {
+                klass.transformSingle(this, data)
+            }
+        }
+        return declaration
+    }
+
     @Suppress("UNCHECKED_CAST")
     override fun <F : FirClass<F>> transformClass(
         klass: FirClass<F>,
@@ -340,27 +363,19 @@ abstract class AbstractFirStatusResolveTransformer(
     ): FirStatement {
         return storeClass(klass) {
             klass.typeParameters.forEach { it.transformSingle(this, data) }
-            klass.replaceResolvePhase(transformerPhase)
-            if (klass.needResolveMembers()) {
-                val members = klass.declarations.filter { it !is FirClassLikeDeclaration<*> }
-                members.forEach { it.replaceResolvePhase(transformerPhase) }
-                members.forEach { it.transformSingle(this, data) }
+            if (needReplacePhase(klass)) {
+                klass.replaceResolvePhase(transformerPhase)
             }
-            if (klass.needResolveNestedClassifiers()) {
-                for (declaration in klass.declarations) {
-                    if (declaration is FirClassLikeDeclaration<*>) {
-                        declaration.transformSingle(this, data)
-                    }
-                }
-            }
-            klass
+            transformDeclarationContent(klass, data)
         } as FirStatement
     }
 
     protected fun updateResolvePhaseOfMembers(regularClass: FirRegularClass) {
         for (declaration in regularClass.declarations) {
             if (declaration is FirProperty || declaration is FirSimpleFunction) {
-                declaration.replaceResolvePhase(transformerPhase)
+                if (needReplacePhase(declaration)) {
+                    declaration.replaceResolvePhase(transformerPhase)
+                }
             }
         }
     }
@@ -372,7 +387,9 @@ abstract class AbstractFirStatusResolveTransformer(
     }
 
     private fun forceResolveStatusOfCorrespondingClass(typeRef: FirTypeRef) {
-        when (val superClass = typeRef.coneType.toSymbol(session)?.fir) {
+        val superClass = typeRef.coneType.toSymbol(session)?.fir
+        superClass?.ensureResolved(FirResolvePhase.SUPER_TYPES, session)
+        when (superClass) {
             is FirRegularClass -> forceResolveStatusesOfClass(superClass)
             is FirTypeAlias -> forceResolveStatusOfCorrespondingClass(superClass.expandedTypeRef)
         }
@@ -439,7 +456,9 @@ abstract class AbstractFirStatusResolveTransformer(
             statusResolver.resolveStatus(propertyAccessor, containingClass, containingProperty, isLocal = false)
         )
 
-        propertyAccessor.replaceResolvePhase(transformerPhase)
+        if (needReplacePhase(propertyAccessor)) {
+            propertyAccessor.replaceResolvePhase(transformerPhase)
+        }
     }
 
     override fun transformConstructor(
@@ -454,7 +473,9 @@ abstract class AbstractFirStatusResolveTransformer(
         simpleFunction: FirSimpleFunction,
         data: FirResolvedDeclarationStatus?
     ): FirDeclaration {
-        simpleFunction.replaceResolvePhase(transformerPhase)
+        if (needReplacePhase(simpleFunction)) {
+            simpleFunction.replaceResolvePhase(transformerPhase)
+        }
         simpleFunction.transformStatus(this, statusResolver.resolveStatus(simpleFunction, containingClass, isLocal = false))
         return transformDeclaration(simpleFunction, data)
     }
@@ -463,7 +484,9 @@ abstract class AbstractFirStatusResolveTransformer(
         property: FirProperty,
         data: FirResolvedDeclarationStatus?
     ): FirDeclaration {
-        property.replaceResolvePhase(transformerPhase)
+        if (needReplacePhase(property)) {
+            property.replaceResolvePhase(transformerPhase)
+        }
         property.transformStatus(this, statusResolver.resolveStatus(property, containingClass, isLocal = false))
 
         property.getter?.let { transformPropertyAccessor(it, property) }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -95,23 +95,31 @@ fun <F : FirClassLikeDeclaration<F>> F.runSupertypeResolvePhaseForLocalClass(
     return this.transform<F, Nothing?>(applySupertypesTransformer, null)
 }
 
-class FirApplySupertypesTransformer(
+open class FirApplySupertypesTransformer(
     private val supertypeComputationSession: SupertypeComputationSession
 ) : FirDefaultTransformer<Any?>() {
     override fun <E : FirElement> transformElement(element: E, data: Any?): E {
         return element
     }
 
-    override fun transformFile(file: FirFile, data: Any?): FirDeclaration {
-        file.replaceResolvePhase(FirResolvePhase.SUPER_TYPES)
+    protected open fun needReplacePhase(firDeclaration: FirDeclaration) = FirResolvePhase.SUPER_TYPES > firDeclaration.resolvePhase
 
-        return (file.transformChildren(this, null) as FirFile)
+    protected open fun transformDeclarationContent(declaration: FirDeclaration, data: Any?): FirDeclaration {
+        return declaration.transformChildren(this, null) as FirDeclaration
+    }
+
+    override fun transformFile(file: FirFile, data: Any?): FirDeclaration {
+        if (needReplacePhase(file)) {
+            file.replaceResolvePhase(FirResolvePhase.SUPER_TYPES)
+        }
+
+        return transformDeclarationContent(file, null) as FirFile
     }
 
     override fun transformRegularClass(regularClass: FirRegularClass, data: Any?): FirStatement {
         applyResolvedSupertypesToClass(regularClass)
 
-        return (regularClass.transformChildren(this, null) as FirRegularClass)
+        return transformDeclarationContent(regularClass, null) as FirRegularClass
     }
 
     private fun applyResolvedSupertypesToClass(firClass: FirClass<*>) {
@@ -121,7 +129,9 @@ class FirApplySupertypesTransformer(
             // TODO: Replace with an immutable version or transformer
             firClass.replaceSuperTypeRefs(supertypeRefs)
         }
-        firClass.replaceResolvePhase(FirResolvePhase.SUPER_TYPES)
+        if (needReplacePhase(firClass)) {
+            firClass.replaceResolvePhase(FirResolvePhase.SUPER_TYPES)
+        }
     }
 
     override fun transformAnonymousObject(anonymousObject: FirAnonymousObject, data: Any?): FirStatement {
@@ -139,7 +149,12 @@ class FirApplySupertypesTransformer(
     }
 
     override fun transformTypeAlias(typeAlias: FirTypeAlias, data: Any?): FirDeclaration {
-        if (typeAlias.expandedTypeRef is FirResolvedTypeRef) return typeAlias
+        if (typeAlias.expandedTypeRef is FirResolvedTypeRef) {
+            if (needReplacePhase(typeAlias)) {
+                typeAlias.replaceResolvePhase(FirResolvePhase.SUPER_TYPES)
+            }
+            return typeAlias
+        }
         val supertypeRefs = getResolvedSupertypeRefs(typeAlias)
 
         assert(supertypeRefs.size == 1) {
@@ -148,8 +163,9 @@ class FirApplySupertypesTransformer(
 
         // TODO: Replace with an immutable version or transformer
         typeAlias.replaceExpandedTypeRef(supertypeRefs[0])
-        typeAlias.replaceResolvePhase(FirResolvePhase.SUPER_TYPES)
-
+        if (needReplacePhase(typeAlias)) {
+            typeAlias.replaceResolvePhase(FirResolvePhase.SUPER_TYPES)
+        }
         return typeAlias
     }
 }
@@ -187,7 +203,7 @@ fun FirRegularClass.resolveSupertypesInTheAir(session: FirSession): List<FirType
         .resolveSpecificClassLikeSupertypes(this, superTypeRefs)
 }
 
-class FirSupertypeResolverVisitor(
+open class FirSupertypeResolverVisitor(
     private val session: FirSession,
     private val supertypeComputationSession: SupertypeComputationSession,
     private val scopeSession: ScopeSession,
@@ -295,14 +311,18 @@ class FirSupertypeResolverVisitor(
         return resolvedTypesRefs
     }
 
+    open fun visitDeclarationContent(declaration: FirDeclaration, data: Any?) {
+        declaration.acceptChildren(this, null)
+    }
+
     override fun visitRegularClass(regularClass: FirRegularClass, data: Any?) {
         resolveSpecificClassLikeSupertypes(regularClass, regularClass.superTypeRefs)
-        regularClass.acceptChildren(this, null)
+        visitDeclarationContent(regularClass, null)
     }
 
     override fun visitAnonymousObject(anonymousObject: FirAnonymousObject, data: Any?) {
         resolveSpecificClassLikeSupertypes(anonymousObject, anonymousObject.superTypeRefs)
-        anonymousObject.acceptChildren(this, null)
+        visitDeclarationContent(anonymousObject, null)
     }
 
     fun resolveSpecificClassLikeSupertypes(
@@ -397,7 +417,7 @@ class FirSupertypeResolverVisitor(
     }
 
     override fun visitFile(file: FirFile, data: Any?) {
-        file.acceptChildren(this, null)
+        visitDeclarationContent(file, null)
     }
 }
 
@@ -409,7 +429,7 @@ private fun createErrorTypeRef(fir: FirElement, message: String, kind: Diagnosti
 class SupertypeComputationSession {
     private val fileScopesMap = hashMapOf<FirFile, ScopePersistentList>()
     private val scopesForNestedClassesMap = hashMapOf<FirClass<*>, ScopePersistentList>()
-    private val supertypeStatusMap = linkedMapOf<FirClassLikeDeclaration<*>, SupertypeComputationStatus>()
+    val supertypeStatusMap = linkedMapOf<FirClassLikeDeclaration<*>, SupertypeComputationStatus>()
 
     val supertypesSupplier: SupertypeSupplier = object : SupertypeSupplier() {
         override fun forClass(firClass: FirClass<*>, useSiteSession: FirSession): List<ConeClassLikeType> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirTypeResolveTransformer.kt

@@ -94,7 +94,9 @@ open class FirTypeResolveTransformer(
     }
 
     override fun transformEnumEntry(enumEntry: FirEnumEntry, data: Any?): FirDeclaration {
-        enumEntry.replaceResolvePhase(FirResolvePhase.TYPES)
+        if (needReplacePhase(enumEntry)) {
+            enumEntry.replaceResolvePhase(FirResolvePhase.TYPES)
+        }
         enumEntry.transformReturnTypeRef(this, data)
         enumEntry.transformTypeParameters(this, data)
         enumEntry.transformAnnotations(this, data)
@@ -104,7 +106,11 @@ open class FirTypeResolveTransformer(
     override fun transformProperty(property: FirProperty, data: Any?): FirDeclaration {
         return withScopeCleanup {
             property.addTypeParametersScope()
-            property.replaceResolvePhase(FirResolvePhase.TYPES)
+            if (needReplacePhase(property)) {
+                property.replaceResolvePhase(FirResolvePhase.TYPES)
+                property.getter?.replaceResolvePhase(FirResolvePhase.TYPES)
+                property.setter?.replaceResolvePhase(FirResolvePhase.TYPES)
+            }
             property.transformTypeParameters(this, data)
                 .transformReturnTypeRef(this, data)
                 .transformReceiverTypeRef(this, data)
@@ -125,7 +131,9 @@ open class FirTypeResolveTransformer(
 
     override fun transformField(field: FirField, data: Any?): FirDeclaration {
         return withScopeCleanup {
-            field.replaceResolvePhase(FirResolvePhase.TYPES)
+            if (needReplacePhase(field)) {
+                field.replaceResolvePhase(FirResolvePhase.TYPES)
+            }
             field.transformReturnTypeRef(this, data).transformAnnotations(this, data)
             field
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt

@@ -43,7 +43,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
 
     final override val session: FirSession get() = components.session
 
-    protected open fun needReplacePhase(firDeclaration: FirDeclaration) = true
+    protected open fun needReplacePhase(firDeclaration: FirDeclaration) = transformerPhase > firDeclaration.resolvePhase
 
     fun replaceDeclarationResolvePhaseIfNeeded(firDeclaration: FirDeclaration, newResolvePhase: FirResolvePhase) {
         if (needReplacePhase(firDeclaration) && newResolvePhase > firDeclaration.resolvePhase) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -378,6 +378,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         typeAlias.transformAnnotations(transformer, data)
         transformer.firTowerDataContextCollector?.addDeclarationContext(typeAlias, context.towerDataContext)
         typeAlias.transformExpandedTypeRef(transformer, data)
+        transformer.replaceDeclarationResolvePhaseIfNeeded(typeAlias, transformerPhase)
         return typeAlias
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirKotlinScopeProvider.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
+import org.jetbrains.kotlin.fir.resolve.transformers.ensureResolved
 import org.jetbrains.kotlin.fir.scopes.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
@@ -108,7 +109,10 @@ fun ConeKotlinType.scopeForSupertype(
     if (this is ConeClassErrorType) return null
     val symbol = lookupTag.toSymbol(useSiteSession)
     return if (symbol is FirRegularClassSymbol) {
-        val delegateField = delegateFields?.find { useSiteSession.typeContext.equalTypes(it.returnTypeRef.coneType, this) }
+        val delegateField = delegateFields?.find {
+            it.ensureResolved(FirResolvePhase.TYPES, useSiteSession)
+            useSiteSession.typeContext.equalTypes(it.returnTypeRef.coneType, this)
+        }
         symbol.fir.scopeForSupertype(
             substitutor(symbol, this, useSiteSession),
             useSiteSession, scopeSession, delegateField,

compiler/fir/tree/src/org/jetbrains/kotlin/fir/EnumClassUtils.kt

@@ -56,7 +56,9 @@ fun FirRegularClassBuilder.generateValuesFunction(
         }
         symbol = FirNamedFunctionSymbol(CallableId(packageFqName, classFqName, ENUM_VALUES))
         resolvePhase = FirResolvePhase.BODY_RESOLVE
-        body = buildEmptyExpressionBlock()
+        body = buildEmptyExpressionBlock().also {
+            it.replaceTypeRef(returnTypeRef)
+        }
     }.apply {
         containingClassAttr = this@generateValuesFunction.symbol.toLookupTag()
     }
@@ -97,7 +99,9 @@ fun FirRegularClassBuilder.generateValueOfFunction(
             isVararg = false
         }
         resolvePhase = FirResolvePhase.BODY_RESOLVE
-        body = buildEmptyExpressionBlock()
+        body = buildEmptyExpressionBlock().also {
+            it.replaceTypeRef(returnTypeRef)
+        }
     }.apply {
         containingClassAttr = this@generateValueOfFunction.symbol.toLookupTag()
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/impl/FirDefaultPropertyAccessor.kt

@@ -12,10 +12,7 @@ import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.contracts.impl.FirEmptyContractDescription
-import org.jetbrains.kotlin.fir.declarations.FirDeclarationAttributes
-import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.declarations.FirValueParameter
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildDefaultSetterValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertyAccessorSymbol
@@ -40,7 +37,7 @@ abstract class FirDefaultPropertyAccessor(
 ) : FirPropertyAccessorImpl(
     source,
     moduleData,
-    resolvePhase = FirResolvePhase.BODY_RESOLVE,
+    resolvePhase = if (effectiveVisibility != null) FirResolvePhase.BODY_RESOLVE else FirResolvePhase.TYPES,
     origin,
     FirDeclarationAttributes(),
     propertyTypeRef,
@@ -59,7 +56,7 @@ abstract class FirDefaultPropertyAccessor(
     typeParameters = mutableListOf(),
 ) {
     override var resolvePhase
-        get() = FirResolvePhase.BODY_RESOLVE
+        get() = if (status is FirResolvedDeclarationStatus) FirResolvePhase.BODY_RESOLVE else FirResolvePhase.TYPES
         set(_) {}
 
     final override var body: FirBlock?

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

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.KtToFirMapping
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.containingKtFileIfAny
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.psi.KtDeclaration
@@ -50,6 +51,7 @@ internal class FirModuleResolveStateDepended(
             firFileBuilder = originalState.firFileBuilder,
             moduleFileCache = originalState.rootModuleSession.cache,
             fileStructureCache = fileStructureCache,
+            firLazyDeclarationResolver = originalState.firLazyDeclarationResolver,
             state = this,
         )
     }
@@ -60,6 +62,9 @@ internal class FirModuleResolveStateDepended(
     override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D =
         originalState.resolveFirToPhase(declaration, toPhase)
 
+    override fun <D : FirDeclaration> resolveFirToResolveType(declaration: D, type: ResolveType): D =
+        originalState.resolveFirToResolveType(declaration, type)
+
     override fun tryGetCachedFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {
         val ktFile = declaration.containingKtFileIfAny ?: return null
         cache.getCachedFirFile(ktFile)?.let { return it }

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

@@ -25,6 +25,7 @@ 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.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.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
@@ -42,7 +43,6 @@ internal class FirModuleResolveStateImpl(
 ) : FirModuleResolveState() {
     override val rootModuleSession: FirIdeSourcesSession get() = sessionProvider.rootModuleSession
 
-
     /**
      * WARNING! This object contains scopes for all statements and declarations that were ever resolved.
      * It can grow unbounded if you never edit the files in the opened project.
@@ -57,7 +57,14 @@ internal class FirModuleResolveStateImpl(
         sessionProvider.getSession(moduleInfo)!!
 
     override fun getOrBuildFirFor(element: KtElement): FirElement =
-        elementBuilder.getOrBuildFirFor(element, firFileBuilder, rootModuleSession.cache, fileStructureCache, this)
+        elementBuilder.getOrBuildFirFor(
+            element = element,
+            firFileBuilder = firFileBuilder,
+            moduleFileCache = rootModuleSession.cache,
+            fileStructureCache = fileStructureCache,
+            firLazyDeclarationResolver = firLazyDeclarationResolver,
+            state = this
+        )
 
     override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
         firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache, lazyBodiesMode = false)
@@ -98,15 +105,12 @@ internal class FirModuleResolveStateImpl(
 
         if (ktDeclaration == nonLocalNamedDeclaration) return nonLocalFirForNamedDeclaration
 
-        if (nonLocalFirForNamedDeclaration.resolvePhase < FirResolvePhase.BODY_RESOLVE) {
-            val cache = (nonLocalFirForNamedDeclaration.moduleData.session as FirIdeSourcesSession).cache
-            firLazyDeclarationResolver.lazyResolveDeclaration(
-                nonLocalFirForNamedDeclaration,
-                cache,
-                FirResolvePhase.BODY_RESOLVE,
-                checkPCE = false, /*TODO*/
-            )
-        }
+        firLazyDeclarationResolver.lazyResolveDeclaration(
+            firDeclarationToResolve = nonLocalFirForNamedDeclaration,
+            moduleFileCache = (nonLocalFirForNamedDeclaration.moduleData.session as FirIdeSourcesSession).cache,
+            toPhase = FirResolvePhase.BODY_RESOLVE,
+            checkPCE = false, /*TODO*/
+        )
         val firDeclaration = FirElementFinder.findElementIn<FirDeclaration>(nonLocalFirForNamedDeclaration) { firDeclaration ->
             when (val realPsi = firDeclaration.realPsi) {
                 is KtObjectLiteralExpression -> realPsi.objectDeclaration == ktDeclaration
@@ -137,14 +141,30 @@ internal class FirModuleResolveStateImpl(
     }
 
     override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
+        if (toPhase == FirResolvePhase.RAW_FIR) return declaration
         val fileCache = when (val session = declaration.moduleData.session) {
             is FirIdeSourcesSession -> session.cache
             else -> return declaration
         }
         firLazyDeclarationResolver.lazyResolveDeclaration(
-            declaration,
-            fileCache,
-            toPhase,
+            firDeclarationToResolve = declaration,
+            moduleFileCache = fileCache,
+            toPhase = toPhase,
+            checkPCE = true,
+        )
+        return declaration
+    }
+
+    override fun <D : FirDeclaration> resolveFirToResolveType(declaration: D, type: ResolveType): D {
+        if (type == ResolveType.NoResolve) return declaration
+        val fileCache = when (val session = declaration.moduleData.session) {
+            is FirIdeSourcesSession -> session.cache
+            else -> return declaration
+        }
+        firLazyDeclarationResolver.lazyResolveDeclaration(
+            firDeclaration = declaration,
+            moduleFileCache = fileCache,
+            toResolveType = type,
             checkPCE = true,
         )
         return declaration

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

@@ -5,18 +5,7 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.*
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitTypeBodyResolveProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirAnnotationArgumentsResolveProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirPluginAnnotationsResolveProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirTransformerBasedExtensionStatusProcessor
-import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
 import java.util.concurrent.locks.ReentrantLock
 import kotlin.concurrent.withLock
@@ -26,21 +15,6 @@ internal class FirPhaseRunner {
     private val statusResolveLock = ReentrantLock()
     private val implicitTypesResolveLock = ReentrantLock()
 
-    fun runPhase(firFile: FirFile, phase: FirResolvePhase, scopeSession: ScopeSession) = when (phase) {
-        FirResolvePhase.SUPER_TYPES -> superTypesBodyResolveLock.withLock {
-            runPhaseWithoutLock(firFile, phase, scopeSession)
-        }
-        FirResolvePhase.STATUS -> statusResolveLock.withLock {
-            runPhaseWithoutLock(firFile, phase, scopeSession)
-        }
-        FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> implicitTypesResolveLock.withLock {
-            runPhaseWithoutLock(firFile, phase, scopeSession)
-        }
-        else -> {
-            runPhaseWithoutLock(firFile, phase, scopeSession)
-        }
-    }
-
     inline fun runPhaseWithCustomResolve(phase: FirResolvePhase, crossinline resolve: () -> Unit) = when (phase) {
         FirResolvePhase.SUPER_TYPES -> superTypesBodyResolveLock.withLock {
             runPhaseWithCustomResolveWithoutLock(resolve)
@@ -61,34 +35,4 @@ internal class FirPhaseRunner {
             resolve()
         }
     }
-
-    private fun runPhaseWithoutLock(firFile: FirFile, phase: FirResolvePhase, scopeSession: ScopeSession) {
-        val phaseProcessor = phase.createTransformerBasedProcessorByPhase(firFile.moduleData.session, scopeSession)
-        executeWithoutPCE {
-            FirLazyBodiesCalculator.calculateLazyBodiesIfPhaseRequires(firFile, phase)
-            phaseProcessor.processFile(firFile)
-        }
-    }
-}
-
-internal fun FirResolvePhase.createTransformerBasedProcessorByPhase(
-    session: FirSession,
-    scopeSession: ScopeSession
-): FirTransformerBasedResolveProcessor {
-    return when (this) {
-        FirResolvePhase.RAW_FIR -> throw IllegalStateException("Raw FIR building phase does not have a transformer")
-        FirResolvePhase.ANNOTATIONS_FOR_PLUGINS -> FirPluginAnnotationsResolveProcessor(session, scopeSession)
-        FirResolvePhase.CLASS_GENERATION -> FirDummyTransformerBasedProcessor(session, scopeSession) // TODO: remove
-        FirResolvePhase.IMPORTS -> FirImportResolveProcessor(session, scopeSession)
-        FirResolvePhase.SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
-        FirResolvePhase.SEALED_CLASS_INHERITORS -> FirDummyTransformerBasedProcessor(session, scopeSession)
-        FirResolvePhase.TYPES -> FirTypeResolveProcessor(session, scopeSession)
-        FirResolvePhase.ARGUMENTS_OF_PLUGIN_ANNOTATIONS -> FirAnnotationArgumentsResolveProcessor(session, scopeSession)
-        FirResolvePhase.EXTENSION_STATUS_UPDATE -> FirTransformerBasedExtensionStatusProcessor(session, scopeSession)
-        FirResolvePhase.STATUS -> FirStatusResolveProcessor(session, scopeSession)
-        FirResolvePhase.CONTRACTS -> FirContractResolveProcessor(session, scopeSession)
-        FirResolvePhase.NEW_MEMBERS_GENERATION -> FirDummyTransformerBasedProcessor(session, scopeSession) // TODO: remove
-        FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirImplicitTypeBodyResolveProcessor(session, scopeSession)
-        FirResolvePhase.BODY_RESOLVE -> FirBodyResolveProcessor(session, scopeSession)
-    }
-}
+}

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

@@ -26,7 +26,14 @@ internal class IdeFirPhaseManager(
     ) {
         val fir = symbol.fir as FirDeclaration
         try {
-            lazyDeclarationResolver.lazyResolveDeclaration(fir, cache, requiredPhase, checkPCE = true)
+            if (fir.resolvePhase < requiredPhase) { //TODO Make thread safe
+                lazyDeclarationResolver.lazyResolveDeclaration(
+                    firDeclarationToResolve = fir,
+                    moduleFileCache = cache,
+                    toPhase = requiredPhase,
+                    checkPCE = true
+                )
+            }
         } catch (e: Throwable) {
             sessionInvalidator.invalidate(fir.moduleData.session)
             throw e

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

@@ -5,171 +5,72 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
-import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.*
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
-import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
-import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirNonLocalDeclarationBuilder
-import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirReplacement
-import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
-import org.jetbrains.kotlin.psi.*
 
 internal object DeclarationCopyBuilder {
-    fun createDeclarationCopy(
-        state: FirModuleResolveState,
-        nonLocalDeclaration: KtDeclaration,
-        replacement: RawFirReplacement
-    ): FirDeclaration {
+    fun FirSimpleFunction.withBodyFrom(
+        functionWithBody: FirSimpleFunction,
+    ): FirSimpleFunction = buildSimpleFunctionCopy(this) {
+        body = functionWithBody.body
+        symbol = functionWithBody.symbol
+        initDeclaration(this@withBodyFrom, functionWithBody)
+    }.apply { reassignAllReturnTargets(functionWithBody) }
 
-        return when (nonLocalDeclaration) {
-            is KtNamedFunction -> createFunctionCopy(
-                nonLocalDeclaration,
-                state,
-                replacement
-            )
-            is KtProperty -> createPropertyCopy(
-                nonLocalDeclaration,
-                state,
-                replacement
-            )
-            is KtClassOrObject -> createClassCopy(
-                nonLocalDeclaration,
-                state,
-                replacement
-            )
-            is KtTypeAlias -> createTypeAliasCopy(
-                nonLocalDeclaration,
-                state,
-                replacement
-            )
-            else -> error("Unsupported declaration ${nonLocalDeclaration::class.simpleName}")
-        }
+    fun FirRegularClass.withBodyFrom(
+        classWithBody: FirRegularClass,
+    ): FirRegularClass = buildRegularClassCopy(this) {
+        declarations.clear()
+        declarations.addAll(classWithBody.declarations)
+        symbol = classWithBody.symbol
+        initDeclaration(this@withBodyFrom, classWithBody)
+        resolvePhase = minOf(this.resolvePhase, FirResolvePhase.IMPORTS) //TODO move into initDeclaration?
     }
 
-    private fun createFunctionCopy(
-        rootNonLocalDeclaration: KtNamedFunction,
-        state: FirModuleResolveState,
-        replacement: RawFirReplacement,
-    ): FirSimpleFunction {
-
-        val originalFunction = rootNonLocalDeclaration.getOrBuildFirOfType<FirSimpleFunction>(state)
-        val builtFunction = createCopy(rootNonLocalDeclaration, originalFunction, replacement)
-
-        val functionBlock = rootNonLocalDeclaration.bodyBlockExpression
-        if (functionBlock == null || !PsiTreeUtil.isAncestor(functionBlock, replacement.from, true)) {
-            return builtFunction
-        }
-
-        // right now we can't resolve builtFunction header properly, as it built right in air,
-        // without file, which is now required for running stages other then body resolve, so we
-        // take original function header (which is resolved) and copy replacing body with body from builtFunction
-        return buildSimpleFunctionCopy(originalFunction) {
-            body = builtFunction.body
-            symbol = builtFunction.symbol
-            initDeclaration(originalFunction, builtFunction, state)
-        }.apply { reassignAllReturnTargets(builtFunction) }
-    }
-
-    private fun createClassCopy(
-        rootNonLocalDeclaration: KtClassOrObject,
-        state: FirModuleResolveState,
-        replacement: RawFirReplacement,
-    ): FirRegularClass {
-        val originalFirClass = rootNonLocalDeclaration.getOrBuildFirOfType<FirRegularClass>(state)
-        val builtClass = createCopy(rootNonLocalDeclaration, originalFirClass, replacement)
-
-        val classBody = rootNonLocalDeclaration.body
-        if (classBody == null || !PsiTreeUtil.isAncestor(classBody, replacement.from, true)) {
-            return builtClass
-        }
-
-        return buildRegularClassCopy(originalFirClass) {
-            declarations.clear()
-            declarations.addAll(builtClass.declarations)
-            symbol = builtClass.symbol
-            initDeclaration(originalFirClass, builtClass, state)
-            resolvePhase = minOf(originalFirClass.resolvePhase, FirResolvePhase.IMPORTS) //TODO move into initDeclaration?
-        }
-    }
-
-    private fun createTypeAliasCopy(
-        rootNonLocalDeclaration: KtTypeAlias,
-        state: FirModuleResolveState,
-        replacement: RawFirReplacement,
-    ): FirTypeAlias {
-        val originalFirTypeAlias = rootNonLocalDeclaration.getOrBuildFirOfType<FirTypeAlias>(state)
-        return createCopy(rootNonLocalDeclaration, originalFirTypeAlias, replacement)
-    }
-
-    private fun createPropertyCopy(
-        rootNonLocalDeclaration: KtProperty,
-        state: FirModuleResolveState,
-        replacement: RawFirReplacement,
-    ): FirProperty {
-        val originalProperty = rootNonLocalDeclaration.getOrBuildFirOfType<FirProperty>(state)
-        val builtProperty = createCopy(rootNonLocalDeclaration, originalProperty, replacement)
-
-        val insideGetterBody = rootNonLocalDeclaration.getter?.bodyBlockExpression?.let {
-            PsiTreeUtil.isAncestor(it, replacement.from, true)
-        } ?: false
-        if (!insideGetterBody) {
-            val insideSetterBody = rootNonLocalDeclaration.setter?.bodyBlockExpression?.let {
-                PsiTreeUtil.isAncestor(it, replacement.from, true)
-            } ?: false
-
-            if (!insideSetterBody) return builtProperty
-        }
-
-        val originalSetter = originalProperty.setter
-        val builtSetter = builtProperty.setter
+    fun FirProperty.withBodyFrom(propertyWithBody: FirProperty): FirProperty {
+        val originalSetter = this@withBodyFrom.setter
+        val replacementSetter = propertyWithBody.setter
 
         // setter has a header with `value` parameter, and we want it type to be resolved
-        val copySetter = if (originalSetter != null && builtSetter != null) {
+        val copySetter = if (originalSetter != null && replacementSetter != null) {
             buildPropertyAccessorCopy(originalSetter) {
-                body = builtSetter.body
-                symbol = builtSetter.symbol
-                initDeclaration(originalSetter, builtSetter, state)
-            }.apply { reassignAllReturnTargets(builtSetter) }
+                body = replacementSetter.body
+                symbol = replacementSetter.symbol
+                initDeclaration(originalSetter, replacementSetter)
+            }.apply { reassignAllReturnTargets(replacementSetter) }
         } else {
-            builtSetter
+            replacementSetter
         }
 
-        return buildPropertyCopy(originalProperty) {
-            symbol = builtProperty.symbol
-            initializer = builtProperty.initializer
+        val propertyResolvePhase = minOf(
+            this@withBodyFrom.resolvePhase,
+            FirResolvePhase.DECLARATIONS,
+            copySetter?.resolvePhase ?: FirResolvePhase.BODY_RESOLVE,
+            propertyWithBody.getter?.resolvePhase ?: FirResolvePhase.BODY_RESOLVE,
+        )
 
-            getter = builtProperty.getter
+        return buildPropertyCopy(this@withBodyFrom) {
+            symbol = propertyWithBody.symbol
+            initializer = propertyWithBody.initializer
+
+            getter = propertyWithBody.getter
             setter = copySetter
 
-            initDeclaration(originalProperty, builtProperty, state)
+            initDeclaration(this@withBodyFrom, propertyWithBody)
+            resolvePhase = propertyResolvePhase
         }
     }
 
     private fun FirDeclarationBuilder.initDeclaration(
         originalDeclaration: FirDeclaration,
         builtDeclaration: FirDeclaration,
-        state: FirModuleResolveState
     ) {
         resolvePhase = minOf(originalDeclaration.resolvePhase, FirResolvePhase.DECLARATIONS)
         source = builtDeclaration.source
-        moduleData = state.rootModuleSession.moduleData
-    }
-
-    internal inline fun <reified D : FirDeclaration> createCopy(
-        rootNonLocalDeclaration: KtDeclaration,
-        originalFirDeclaration: D,
-        replacement: RawFirReplacement? = null,
-    ): D {
-        return RawFirNonLocalDeclarationBuilder.build(
-            session = originalFirDeclaration.moduleData.session,
-            baseScopeProvider = originalFirDeclaration.moduleData.session.firIdeProvider.kotlinScopeProvider,
-            designation = originalFirDeclaration.collectDesignation(),
-            rootNonLocalDeclaration = rootNonLocalDeclaration,
-            replacement = replacement,
-        ) as D
+        moduleData = originalDeclaration.moduleData
     }
 
     private fun FirFunction<*>.reassignAllReturnTargets(from: FirFunction<*>) {

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

@@ -50,6 +50,7 @@ private fun FirClassLikeDeclaration<*>.containingClass(): FirClassLikeDeclaratio
 private fun collectDesignationAndIsLocal(declaration: FirDeclaration): Pair<List<FirDeclaration>, Boolean> {
     val containingClass = when (declaration) {
         is FirCallableDeclaration<*> -> declaration.containingClass()?.toFirRegularClass(declaration.moduleData.session)
+        is FirAnonymousObject -> return emptyList<FirDeclaration>() to true
         is FirClassLikeDeclaration<*> -> declaration.containingClass()
         else -> error("Invalid declaration ${declaration.renderWithType()}")
     } ?: return emptyList<FirDeclaration>() to false

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
@@ -82,4 +83,6 @@ abstract class FirModuleResolveState {
 
 
     internal abstract fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D
+
+    internal abstract fun <D : FirDeclaration> resolveFirToResolveType(declaration: D, type: ResolveType): D
 }

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -109,6 +110,19 @@ fun <D : FirDeclaration, R> D.withFirDeclaration(
     return resolveState.withLock(this, DeclarationLockType.READ_LOCK, action)
 }
 
+/**
+ * Executes [action] with given [FirDeclaration] under read action, so resolve **is not possible** inside [action]
+ * [FirDeclaration] passed to [action] will be resolved at least to [phase] when executing [action] on it
+ */
+fun <D : FirDeclaration, R> D.withFirDeclaration(
+    type: ResolveType,
+    resolveState: FirModuleResolveState,
+    action: (D) -> R,
+): R {
+    resolvedFirToType(type, resolveState)
+    return resolveState.withLock(this, DeclarationLockType.READ_LOCK, action)
+}
+
 /**
  * Executes [action] with given [FirDeclaration] under write lock, so resolve **is possible** inside [action]
  */
@@ -149,6 +163,17 @@ fun <D : FirDeclaration> D.resolvedFirToPhase(
 ): D =
     resolveState.resolveFirToPhase(this, phase)
 
+/**
+ * Resolves a given [FirDeclaration] to [phase] and returns resolved declaration
+ *
+ * Should not be called form [withFirDeclaration], [withFirDeclarationOfType] functions, as it it may cause deadlock
+ */
+fun <D : FirDeclaration> D.resolvedFirToType(
+    type: ResolveType,
+    resolveState: FirModuleResolveState
+): D =
+    resolveState.resolveFirToResolveType(this, type)
+
 
 /**
  * Get a [FirElement] which was created by [KtElement]

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

@@ -8,43 +8,40 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.parentOfType
 import com.intellij.psi.util.parentsOfType
-import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.asTowerDataElement
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTypeResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
-import org.jetbrains.kotlin.fir.resolve.typeResolver
-import org.jetbrains.kotlin.fir.scopes.FirCompositeScope
 import org.jetbrains.kotlin.fir.scopes.createImportingScopes
-import org.jetbrains.kotlin.fir.toFirPsiSourceElement
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.fir.types.toFirResolvedTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateDepended
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.DeclarationCopyBuilder.withBodyFrom
 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.FirFileBuilder
 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.*
+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.RawFirReplacement
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.buildFileFirAnnotation
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.buildFirUserTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirProviderInterceptorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.idea.util.ifTrue
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
@@ -84,7 +81,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         val declaration = runBodyResolveOnAir(
             state = state,
             replacement = RawFirReplacement(place, elementToResolve),
-            isOnAirResolve = true,
+            resolveWithUnchangedFir = false
         )
 
         val expressionLocator = object : FirVisitorVoid() {
@@ -120,25 +117,30 @@ object LowLevelFirApiFacadeForResolveOnAir {
                 runBodyResolveOnAir(
                     state = state,
                     collector = it,
+                    resolveWithUnchangedFir = true,
                     replacement = RawFirReplacement(validPlace, validPlace),
-                    isOnAirResolve = false //isOnAirResolve can be little faster because node resolved in it's context
                 )
             }
         }
     }
 
     private fun onAirGetTowerContextForFile(
-        state: FirModuleResolveState,
+        state: FirModuleResolveStateImpl,
         file: KtFile,
     ): FirTowerDataContext {
         require(file.isPhysical)
         val session = state.getSessionFor(file.getModuleInfo()) as FirIdeSourcesSession
-        val firFile = session.firFileBuilder.getFirFileResolvedToPhaseWithCaching(
-            file,
-            session.cache,
-            FirResolvePhase.IMPORTS,
-            ScopeSession(),
-            checkPCE = false
+
+        val firFile = session.firFileBuilder.buildRawFirFileWithCaching(
+            ktFile = file,
+            cache = session.cache,
+            lazyBodiesMode = true
+        )
+
+        state.firLazyDeclarationResolver.lazyResolveFileDeclaration(
+            firFile = firFile,
+            moduleFileCache = session.cache,
+            toPhase = FirResolvePhase.IMPORTS
         )
 
         val importingScopes = createImportingScopes(firFile, firFile.moduleData.session, ScopeSession(), useCaching = false)
@@ -175,7 +177,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         val copiedFirDeclaration = runBodyResolveOnAir(
             originalState,
             replacement = RawFirReplacement(sameDeclarationInOriginalFile, dependencyNonLocalDeclaration),
-            isOnAirResolve = true,
+            resolveWithUnchangedFir = false,
             collector = collector,
         )
 
@@ -191,17 +193,18 @@ object LowLevelFirApiFacadeForResolveOnAir {
         collector: FirTowerDataContextCollector? = null,
     ): FirAnnotationCall {
         val annotationCall = buildFileFirAnnotation(
-            firFile.moduleData.session,
-            firFile.moduleData.session.firIdeProvider.kotlinScopeProvider,
-            annotationEntry,
-            replacement
+            session = firFile.moduleData.session,
+            baseScopeProvider = firFile.moduleData.session.firIdeProvider.kotlinScopeProvider,
+            fileAnnotation = annotationEntry,
+            replacement = replacement
         )
         state.firLazyDeclarationResolver.resolveFileAnnotations(
-            firFile,
-            listOf(annotationCall),
-            state.rootModuleSession.cache,
-            ScopeSession(),
-            collector
+            firFile = firFile,
+            annotations = listOf(annotationCall),
+            moduleFileCache = state.rootModuleSession.cache,
+            scopeSession = ScopeSession(),
+            checkPCE = true,
+            collector = collector
         )
 
         return annotationCall
@@ -210,7 +213,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
     private fun runBodyResolveOnAir(
         state: FirModuleResolveStateImpl,
         replacement: RawFirReplacement,
-        isOnAirResolve: Boolean,
+        resolveWithUnchangedFir: Boolean,
         collector: FirTowerDataContextCollector? = null,
     ): FirElement {
 
@@ -233,24 +236,69 @@ object LowLevelFirApiFacadeForResolveOnAir {
             }
         }
 
-        val copiedFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(
-            state = state,
-            nonLocalDeclaration = nonLocalDeclaration,
+        val originalDeclaration = nonLocalDeclaration.getOrBuildFir(state)
+        check(originalDeclaration is FirDeclaration) { "Invalid original declaration type ${originalDeclaration::class.simpleName}" }
+
+        val originalDesignation = originalDeclaration.collectDesignation()
+
+        val newDeclarationWithReplacement = RawFirNonLocalDeclarationBuilder.buildWithReplacement(
+            session = originalDeclaration.moduleData.session,
+            scopeProvider = originalDeclaration.moduleData.session.firIdeProvider.kotlinScopeProvider,
+            designation = originalDesignation,
+            rootNonLocalDeclaration = nonLocalDeclaration,
             replacement = replacement,
         )
 
-        state.firLazyDeclarationResolver.lazyDesignatedResolveDeclaration(
-            firDeclarationToResolve = copiedFirDeclaration,
-            moduleFileCache = state.rootModuleSession.cache,
-            containerFirFile = originalFirFile,
-            provider = originalFirFile.moduleData.session.firIdeProvider,
-            toPhase = FirResolvePhase.BODY_RESOLVE,
-            checkPCE = true,
-            isOnAirResolve = isOnAirResolve,
-            towerDataContextCollector = collector,
-        )
+        val isInBodyReplacement = isInBodyReplacement(nonLocalDeclaration, replacement)
 
-        return copiedFirDeclaration
+        return FirLazyDeclarationResolver.runCustomResolveUnderLock(originalFirFile, state.rootModuleSession.cache, true) {
+            val copiedFirDeclaration = isInBodyReplacement.ifTrue {
+                when (originalDeclaration) {
+                    is FirSimpleFunction ->
+                        originalDeclaration.withBodyFrom(newDeclarationWithReplacement as FirSimpleFunction)
+                    is FirProperty ->
+                        originalDeclaration.withBodyFrom(newDeclarationWithReplacement as FirProperty)
+                    is FirRegularClass ->
+                        originalDeclaration.withBodyFrom(newDeclarationWithReplacement as FirRegularClass)
+                    is FirTypeAlias -> newDeclarationWithReplacement
+                    else -> error("Not supported type ${originalDeclaration::class.simpleName}")
+                }
+            } ?: newDeclarationWithReplacement
+
+            val onAirDesignation = FirDeclarationUntypedDesignationWithFile(
+                path = originalDesignation.path,
+                declaration = copiedFirDeclaration,
+                isLocalDesignation = false,
+                firFile = originalFirFile
+            )
+            state.firLazyDeclarationResolver.runLazyDesignatedOnAirResolveToBodyWithoutLock(
+                designation = onAirDesignation,
+                moduleFileCache = state.rootModuleSession.cache,
+                checkPCE = true,
+                resolveWithUnchangedFir = resolveWithUnchangedFir,
+                towerDataContextCollector = collector,
+            )
+            copiedFirDeclaration
+        }
+
+    }
+
+    private fun isInBodyReplacement(ktDeclaration: KtDeclaration, replacement: RawFirReplacement): Boolean = when (ktDeclaration) {
+        is KtNamedFunction ->
+            ktDeclaration.bodyBlockExpression?.let { PsiTreeUtil.isAncestor(it, replacement.from, true) } ?: false
+        is KtProperty -> {
+            val insideGetterBody = ktDeclaration.getter?.bodyBlockExpression?.let {
+                PsiTreeUtil.isAncestor(it, replacement.from, true)
+            } ?: false
+
+            insideGetterBody || ktDeclaration.setter?.bodyBlockExpression?.let {
+                PsiTreeUtil.isAncestor(it, replacement.from, true)
+            } ?: false
+        }
+        is KtClassOrObject ->
+            ktDeclaration.body?.let { PsiTreeUtil.isAncestor(it, replacement.from, true) } ?: false
+        is KtTypeAlias -> false
+        else -> error("Not supported type ${ktDeclaration::class.simpleName}")
     }
 
     fun onAirResolveTypeInPlace(

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

@@ -32,7 +32,7 @@ internal class SingleNonLocalDeclarationDiagnosticRetriever(
         lockProvider: LockProvider<FirFile>
     ): FileStructureElementDiagnosticList {
         val sessionHolder = SessionHolderImpl.createWithEmptyScopeSession(firFile.moduleData.session)
-        val context = lockProvider.withReadLock(firFile) {
+        val context = lockProvider.withWriteLock(firFile) {
             PersistenceContextCollector.collectContext(sessionHolder, firFile, structureElementDeclaration)
         }
         return collector.collectForStructureElement(structureElementDeclaration) { components ->

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

@@ -53,20 +53,29 @@ internal class FirElementBuilder {
         firFileBuilder: FirFileBuilder,
         moduleFileCache: ModuleFileCache,
         fileStructureCache: FileStructureCache,
+        firLazyDeclarationResolver: FirLazyDeclarationResolver,
         state: FirModuleResolveState,
     ): FirElement = when (element) {
-        is KtFile -> getOrBuildFirForKtFile(element, firFileBuilder, moduleFileCache)
+        is KtFile -> getOrBuildFirForKtFile(element, firFileBuilder, moduleFileCache, firLazyDeclarationResolver)
         else -> getOrBuildFirForNonKtFileElement(element, fileStructureCache, moduleFileCache, state)
     }
 
-    private fun getOrBuildFirForKtFile(ktFile: KtFile, firFileBuilder: FirFileBuilder, moduleFileCache: ModuleFileCache): FirFile =
-        firFileBuilder.getFirFileResolvedToPhaseWithCaching(
-            ktFile,
-            moduleFileCache,
-            FirResolvePhase.BODY_RESOLVE,
+    private fun getOrBuildFirForKtFile(
+        ktFile: KtFile,
+        firFileBuilder: FirFileBuilder,
+        moduleFileCache: ModuleFileCache,
+        firLazyDeclarationResolver: FirLazyDeclarationResolver
+    ): FirFile {
+        val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, lazyBodiesMode = false)
+        firLazyDeclarationResolver.lazyResolveFileDeclaration(
+            firFile = firFile,
+            moduleFileCache = moduleFileCache,
+            toPhase = FirResolvePhase.BODY_RESOLVE,
             scopeSession = ScopeSession(),
             checkPCE = true
         )
+        return firFile
+    }
 
     private fun getOrBuildFirForNonKtFileElement(
         element: KtElement,

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

@@ -8,14 +8,10 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.file.builder
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.lockWithPCECheck
+import org.jetbrains.kotlin.psi.KtFile
 
 /**
  * Responsible for building [FirFile] by [KtFile]
@@ -37,70 +33,6 @@ internal class FirFileBuilder(
     ): FirFile = cache.fileCached(ktFile) {
         RawFirBuilder(cache.session, scopeProvider, RawFirBuilderMode.lazyBodies(lazyBodiesMode)).buildFirFile(ktFile)
     }
-
-    fun getBuiltFirFileOrNull(ktFile: KtFile, cache: ModuleFileCache): FirFile? =
-        cache.getCachedFirFile(ktFile)
-
-    fun isFirFileBuilt(
-        ktFile: KtFile,
-        cache: ModuleFileCache
-    ): Boolean = cache.getCachedFirFile(ktFile) != null
-
-    fun getFirFileResolvedToPhaseWithCaching(
-        ktFile: KtFile,
-        cache: ModuleFileCache,
-        @Suppress("SameParameterValue") toPhase: FirResolvePhase,
-        scopeSession: ScopeSession,
-        checkPCE: Boolean
-    ): FirFile {
-        val needResolve = toPhase > FirResolvePhase.RAW_FIR
-        val firFile = buildRawFirFileWithCaching(ktFile, cache, lazyBodiesMode = !needResolve)
-        if (needResolve) {
-            cache.firFileLockProvider.withWriteLock(firFile) {
-                if (firFile.resolvePhase >= toPhase) return@withWriteLock
-                runResolveWithoutLock(
-                    firFile,
-                    fromPhase = firFile.resolvePhase,
-                    toPhase = toPhase,
-                    scopeSession = scopeSession,
-                    checkPCE = checkPCE,
-                )
-            }
-        }
-        return firFile
-    }
-
-    /**
-     * Runs [resolve] function (which is considered to do some resolve on [firFile]) under a lock for [firFile]
-     */
-    inline fun <R> runCustomResolveUnderLock(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R =
-        cache.firFileLockProvider.withWriteLock(firFile) { resolve() }
-
-    inline fun <R : Any> runCustomResolveWithPCECheck(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R =
-        cache.firFileLockProvider.withWriteLockPCECheck(firFile, LOCKING_INTERVAL_MS, resolve)
-
-    fun runResolveWithoutLock(
-        firFile: FirFile,
-        fromPhase: FirResolvePhase,
-        toPhase: FirResolvePhase,
-        scopeSession: ScopeSession,
-        checkPCE: Boolean,
-    ) {
-        assert(fromPhase <= toPhase) {
-            "Trying to resolve file ${firFile.name} from $fromPhase to $toPhase"
-        }
-        var currentPhase = fromPhase
-        while (currentPhase < toPhase) {
-            if (checkPCE) checkCanceled()
-            currentPhase = currentPhase.next
-            firPhaseRunner.runPhase(firFile, currentPhase, scopeSession)
-        }
-    }
-
-
-    companion object {
-        private const val LOCKING_INTERVAL_MS = 500L
-    }
 }
 
 

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

@@ -105,9 +105,9 @@ internal class FileStructure(
             firFile
         )
         firLazyDeclarationResolver.lazyResolveDeclaration(
-            firDeclaration,
-            moduleFileCache,
-            FirResolvePhase.BODY_RESOLVE,
+            firDeclarationToResolve = firDeclaration,
+            moduleFileCache = moduleFileCache,
+            toPhase = FirResolvePhase.BODY_RESOLVE,
             checkPCE = true,
         )
         return moduleFileCache.firFileLockProvider.withReadLock(firFile) {
@@ -117,15 +117,14 @@ internal class FileStructure(
 
     private fun createStructureElement(container: KtAnnotated): FileStructureElement = when (container) {
         is KtFile -> {
-            val scopeSession = ScopeSession()
-            val firFile = firFileBuilder.getFirFileResolvedToPhaseWithCaching(
-                container,
-                moduleFileCache,
-                FirResolvePhase.IMPORTS,
-                scopeSession,
+            val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, lazyBodiesMode = true)
+            firLazyDeclarationResolver.resolveFileAnnotations(
+                firFile = firFile,
+                annotations = firFile.annotations,
+                moduleFileCache = moduleFileCache,
+                scopeSession = ScopeSession(),
                 checkPCE = true
             )
-            firLazyDeclarationResolver.resolveFileAnnotations(firFile, firFile.annotations, moduleFileCache, scopeSession)
             RootStructureElement(
                 firFile,
                 container,

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

@@ -12,17 +12,19 @@ import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.DeclarationCopyBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForResolveOnAir
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileDiagnosticRetriever
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnostics
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
 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.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
 import org.jetbrains.kotlin.psi.*
 import java.util.concurrent.ConcurrentHashMap
 
@@ -115,25 +117,42 @@ internal class ReanalyzableFunctionStructureElement(
         firIdeProvider: FirIdeProvider,
     ): ReanalyzableFunctionStructureElement {
         val originalFunction = firSymbol.fir as FirSimpleFunction
-        val newFunction = DeclarationCopyBuilder.createCopy(newKtDeclaration, originalFunction)
+        val designation = originalFunction.collectDesignation()
 
-        return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalFunction, newFunction) {
-            firLazyDeclarationResolver.lazyResolveDeclaration(
-                newFunction,
-                cache,
-                FirResolvePhase.BODY_RESOLVE,
-                checkPCE = true,
-                reresolveFile = true,
-            )
-            cache.firFileLockProvider.withReadLock(firFile) {
-                ReanalyzableFunctionStructureElement(
-                    firFile,
-                    newKtDeclaration,
-                    newFunction.symbol,
-                    newKtDeclaration.modificationStamp,
-                    lockProvider,
-                )
+        val temporaryFunction = RawFirNonLocalDeclarationBuilder.buildWithRebind(
+            session = originalFunction.moduleData.session,
+            scopeProvider = originalFunction.moduleData.session.firIdeProvider.kotlinScopeProvider,
+            designation = designation,
+            rootNonLocalDeclaration = newKtDeclaration,
+        ) as FirSimpleFunction
+
+        return cache.firFileLockProvider.withWriteLock(firFile) {
+
+            val upgradedPhase = minOf(originalFunction.resolvePhase, FirResolvePhase.DECLARATIONS)
+            with(originalFunction) {
+                replaceBody(temporaryFunction.body)
+                replaceContractDescription(temporaryFunction.contractDescription)
+                replaceResolvePhase(upgradedPhase)
             }
+            designation.toSequence(includeTarget = true).forEach {
+                it.resolvePhaseForDeclarationAndChildren = minOf(it.resolvePhaseForDeclarationAndChildren, upgradedPhase)
+            }
+
+            firLazyDeclarationResolver.lazyResolveDeclaration(
+                firDeclarationToResolve = originalFunction,
+                moduleFileCache = cache,
+                toPhase = FirResolvePhase.BODY_RESOLVE,
+                checkPCE = true,
+                declarationPhaseDowngraded = true,
+            )
+
+            ReanalyzableFunctionStructureElement(
+                firFile,
+                newKtDeclaration,
+                originalFunction.symbol,
+                newKtDeclaration.modificationStamp,
+                lockProvider,
+            )
         }
     }
 }
@@ -154,25 +173,45 @@ internal class ReanalyzablePropertyStructureElement(
         firIdeProvider: FirIdeProvider,
     ): ReanalyzablePropertyStructureElement {
         val originalProperty = firSymbol.fir
-        val newProperty = DeclarationCopyBuilder.createCopy(newKtDeclaration, originalProperty)
+        val designation = originalProperty.collectDesignation()
 
-        return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalProperty, newProperty) {
-            firLazyDeclarationResolver.lazyResolveDeclaration(
-                newProperty,
-                cache,
-                FirResolvePhase.BODY_RESOLVE,
-                checkPCE = true,
-                reresolveFile = true,
-            )
-            cache.firFileLockProvider.withReadLock(firFile) {
-                ReanalyzablePropertyStructureElement(
-                    firFile,
-                    newKtDeclaration,
-                    newProperty.symbol,
-                    newKtDeclaration.modificationStamp,
-                    lockProvider,
-                )
+        val temporaryProperty = RawFirNonLocalDeclarationBuilder.buildWithRebind(
+            session = originalProperty.moduleData.session,
+            scopeProvider = originalProperty.moduleData.session.firIdeProvider.kotlinScopeProvider,
+            designation = designation,
+            rootNonLocalDeclaration = newKtDeclaration,
+        ) as FirProperty
+
+        return cache.firFileLockProvider.withWriteLock(firFile) {
+
+            val getterPhase = originalProperty.getter?.resolvePhase ?: originalProperty.resolvePhase
+            val setterPhase = originalProperty.setter?.resolvePhase ?: originalProperty.resolvePhase
+            val upgradedPhase = minOf(originalProperty.resolvePhase, getterPhase, setterPhase, FirResolvePhase.DECLARATIONS)
+
+            with(originalProperty) {
+                getter?.replaceBody(temporaryProperty.getter?.body)
+                setter?.replaceBody(temporaryProperty.setter?.body)
+                replaceInitializer(temporaryProperty.initializer)
+                getter?.replaceResolvePhase(upgradedPhase)
+                setter?.replaceResolvePhase(upgradedPhase)
+                replaceResolvePhase(upgradedPhase)
             }
+
+            firLazyDeclarationResolver.lazyResolveDeclaration(
+                firDeclarationToResolve = originalProperty,
+                moduleFileCache = cache,
+                toPhase = FirResolvePhase.BODY_RESOLVE,
+                checkPCE = true,
+                declarationPhaseDowngraded = true,
+            )
+
+            ReanalyzablePropertyStructureElement(
+                firFile,
+                newKtDeclaration,
+                originalProperty.symbol,
+                newKtDeclaration.modificationStamp,
+                lockProvider,
+            )
         }
     }
 }

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

@@ -28,21 +28,30 @@ internal object FirLazyBodiesCalculator {
         )
     }
 
+    fun calculateLazyBodiesInsideIfNeeded(designation: FirDeclarationUntypedDesignation, phase: FirResolvePhase) {
+        if (phase == FIRST_PHASE_WHICH_NEEDS_BODIES) {
+            calculateLazyBodiesInside(designation)
+        }
+    }
+
     fun calculateLazyBodiesIfPhaseRequires(firFile: FirFile, phase: FirResolvePhase) {
         if (phase == FIRST_PHASE_WHICH_NEEDS_BODIES) {
             firFile.transform<FirElement, MutableList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, mutableListOf())
         }
     }
 
+    fun calculateLazyBodies(firFile: FirFile) {
+        firFile.transform<FirElement, MutableList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, mutableListOf())
+    }
+
     fun calculateLazyBodiesForFunction(designation: FirDeclarationDesignation<FirSimpleFunction>) {
-        require(!designation.isLocalDesignation) { "Not supported for local designations" }
         val simpleFunction = designation.declaration
         if (simpleFunction.body !is FirLazyBlock) return
-        val newFunction = RawFirNonLocalDeclarationBuilder.build(
+        val newFunction = RawFirNonLocalDeclarationBuilder.buildWithRebind(
             session = simpleFunction.moduleData.session,
-            baseScopeProvider = simpleFunction.moduleData.session.firIdeProvider.kotlinScopeProvider,
+            scopeProvider = simpleFunction.moduleData.session.firIdeProvider.kotlinScopeProvider,
             designation = designation,
-            rootNonLocalDeclaration = simpleFunction.psi as KtNamedFunction
+            rootNonLocalDeclaration = simpleFunction.psi as KtNamedFunction,
         ) as FirSimpleFunction
         simpleFunction.apply {
             replaceBody(newFunction.body)
@@ -56,11 +65,11 @@ internal object FirLazyBodiesCalculator {
         require(!secondaryConstructor.isPrimary)
         if (secondaryConstructor.body !is FirLazyBlock) return
 
-        val newFunction = RawFirNonLocalDeclarationBuilder.build(
+        val newFunction = RawFirNonLocalDeclarationBuilder.buildWithRebind(
             session = secondaryConstructor.moduleData.session,
-            baseScopeProvider = secondaryConstructor.moduleData.session.firIdeProvider.kotlinScopeProvider,
+            scopeProvider = secondaryConstructor.moduleData.session.firIdeProvider.kotlinScopeProvider,
             designation = designation,
-            rootNonLocalDeclaration = secondaryConstructor.psi as KtSecondaryConstructor
+            rootNonLocalDeclaration = secondaryConstructor.psi as KtSecondaryConstructor,
         ) as FirSimpleFunction
 
         secondaryConstructor.apply {
@@ -73,9 +82,9 @@ internal object FirLazyBodiesCalculator {
         val firProperty = designation.declaration
         if (!needCalculatingLazyBodyForProperty(firProperty)) return
 
-        val newProperty = RawFirNonLocalDeclarationBuilder.build(
+        val newProperty = RawFirNonLocalDeclarationBuilder.buildWithRebind(
             session = firProperty.moduleData.session,
-            baseScopeProvider = firProperty.moduleData.session.firIdeProvider.kotlinScopeProvider,
+            scopeProvider = firProperty.moduleData.session.firIdeProvider.kotlinScopeProvider,
             designation = designation,
             rootNonLocalDeclaration = firProperty.psi as KtProperty
         ) as FirProperty

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

@@ -5,32 +5,64 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 
-import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+import com.intellij.psi.util.parentOfType
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
-import org.jetbrains.kotlin.fir.psi
-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.FirDeclarationUntypedDesignation
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.*
+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.util.ifTrue
+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,
+        )
+    }
+
     /**
      * Fully resolve file annotations (synchronized)
      * @see resolveFileAnnotationsWithoutLock not synchronized
@@ -40,17 +72,17 @@ internal class FirLazyDeclarationResolver(
         annotations: List<FirAnnotationCall>,
         moduleFileCache: ModuleFileCache,
         scopeSession: ScopeSession,
+        checkPCE: Boolean,
         collector: FirTowerDataContextCollector? = null,
     ) {
-        lazyResolveDeclaration(
-            declaration = firFile,
-            moduleFileCache = moduleFileCache,
-            toPhase = FirResolvePhase.IMPORTS,
-            checkPCE = false,
-            reresolveFile = false
-        )
-        firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
-            resolveFileAnnotationsWithoutLock(firFile, annotations, scopeSession, collector)
+        runCustomResolveUnderLock(firFile, moduleFileCache, checkPCE) {
+            resolveFileAnnotationsWithoutLock(
+                firFile = firFile,
+                moduleFileCache = moduleFileCache,
+                annotations = annotations,
+                scopeSession = scopeSession,
+                collector = collector
+            )
         }
     }
 
@@ -60,152 +92,102 @@ internal class FirLazyDeclarationResolver(
      */
     private fun resolveFileAnnotationsWithoutLock(
         firFile: FirFile,
+        moduleFileCache: ModuleFileCache,
         annotations: List<FirAnnotationCall>,
         scopeSession: ScopeSession,
         collector: FirTowerDataContextCollector? = null,
     ) {
+        lazyResolveFileDeclarationWithoutLock(
+            firFile = firFile,
+            moduleFileCache = moduleFileCache,
+            toPhase = FirResolvePhase.IMPORTS,
+            scopeSession = scopeSession,
+            checkPCE = false,
+            collector = collector
+        )
+
         FirFileAnnotationsResolveTransformer(
             firFile = firFile,
             annotations = annotations,
             session = firFile.moduleData.session,
             scopeSession = scopeSession,
             firTowerDataContextCollector = collector,
-        ).transformDeclaration()
+        ).transformDeclaration(firFileBuilder.firPhaseRunner)
     }
 
-    private fun getResolvableDeclaration(declaration: FirDeclaration, moduleFileCache: ModuleFileCache): FirDeclaration {
-        if (declaration is FirPropertyAccessor || declaration is FirTypeParameter || declaration is FirValueParameter) {
-            val ktContainingResolvableDeclaration = when (val psi = declaration.psi) {
-                is KtPropertyAccessor -> psi.property
-                is KtProperty -> psi
-                is KtParameter, is KtTypeParameter -> psi.getNonLocalContainingOrThisDeclaration()
-                    ?: error("Cannot find containing declaration for KtParameter")
-                is KtCallExpression -> {
-                    check(declaration.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}")
+    private fun FirDeclaration.isValidForResolve(): Boolean = when (origin) {
+        is FirDeclarationOrigin.Source,
+        is FirDeclarationOrigin.ImportedFromObject,
+        is FirDeclarationOrigin.Delegated,
+        is FirDeclarationOrigin.Synthetic,
+        is FirDeclarationOrigin.SubstitutionOverride,
+        is FirDeclarationOrigin.IntersectionOverride -> {
+            when (this) {
+                is FirFile -> true
+                is FirSimpleFunction,
+                is FirProperty,
+                is FirPropertyAccessor,
+                is FirField,
+                is FirTypeAlias,
+                is FirConstructor -> resolvePhase < FirResolvePhase.BODY_RESOLVE
+                else -> true
             }
-            return ktContainingResolvableDeclaration.findSourceNonLocalFirDeclaration(
-                firFileBuilder = firFileBuilder,
-                firSymbolProvider = declaration.moduleData.session.symbolProvider,
-                moduleFileCache = moduleFileCache
+        }
+        else -> {
+            check(resolvePhase == FirResolvePhase.BODY_RESOLVE) {
+                "Expected body resolve phase for origin $origin but found $resolvePhase"
+            }
+            false
+        }
+    }
+
+    fun lazyResolveFileDeclaration(
+        firFile: FirFile,
+        moduleFileCache: ModuleFileCache,
+        toPhase: FirResolvePhase,
+        scopeSession: ScopeSession = ScopeSession(),
+        checkPCE: Boolean = false,
+    ) {
+        runCustomResolveUnderLock(firFile, moduleFileCache, checkPCE) {
+            lazyResolveFileDeclarationWithoutLock(
+                firFile = firFile,
+                moduleFileCache = moduleFileCache,
+                toPhase = toPhase,
+                scopeSession = scopeSession,
+                checkPCE = checkPCE,
             )
         }
-
-        return declaration
     }
 
-    /**
-     * Run partially designated resolve that resolve declaration into last file-wise resolve and then resolve a designation (synchronized)
-     * @see LAST_NON_LAZY_PHASE is the last file-wise resolve
-     * @see lazyDesignatedResolveDeclaration designated resolve
-     * @see runLazyResolveWithoutLock (not synchronized)
-     */
-    fun lazyResolveDeclaration(
-        declaration: FirDeclaration,
+    private fun lazyResolveFileDeclarationWithoutLock(
+        firFile: FirFile,
         moduleFileCache: ModuleFileCache,
         toPhase: FirResolvePhase,
+        scopeSession: ScopeSession,
         checkPCE: Boolean = false,
-        reresolveFile: Boolean = false,
+        collector: FirTowerDataContextCollector? = null,
     ) {
-        if (declaration.resolvePhase >= toPhase) 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
 
-        val resolvableDeclaration = getResolvableDeclaration(declaration, moduleFileCache)
-
-        val firFile = resolvableDeclaration.getContainingFile()
-            ?: error("FirFile was not found for\n${resolvableDeclaration.render()}")
-        val provider = firFile.moduleData.session.firIdeProvider
-
-        // Lazy since we want to read the resolve phase inside the lock. Otherwise, we may run the same resolve phase multiple times. See
-        // KT-45121
-        val fromPhase: FirResolvePhase by lazy(LazyThreadSafetyMode.NONE) {
-            if (reresolveFile) resolvableDeclaration.resolvePhase else minOf(firFile.resolvePhase, resolvableDeclaration.resolvePhase)
+        if (toPhase > FirResolvePhase.IMPORTS) {
+            resolveFileAnnotations(firFile, firFile.annotations, moduleFileCache, scopeSession, checkPCE, collector)
         }
 
-        if (checkPCE) {
-            firFileBuilder.runCustomResolveWithPCECheck(firFile, moduleFileCache) {
-                runLazyResolveWithoutLock(
-                    firDeclarationToResolve = declaration,
-                    moduleFileCache = moduleFileCache,
-                    containerFirFile = firFile,
-                    provider = provider,
-                    fromPhase = fromPhase,
-                    toPhase = toPhase,
-                    checkPCE = true,
-                )
-            }
-        } else {
-            firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
-                executeWithoutPCE {
-                    runLazyResolveWithoutLock(
-                        firDeclarationToResolve = declaration,
-                        moduleFileCache = moduleFileCache,
-                        containerFirFile = firFile,
-                        provider = provider,
-                        fromPhase = fromPhase,
-                        toPhase = toPhase,
-                        checkPCE = false,
-                    )
-                }
-            }
-        }
-    }
-
-    private fun createLazyBodiesCalculator(designation: FirDeclarationUntypedDesignation): (FirResolvePhase) -> Unit {
-        var calculated = false
-        return { phase: FirResolvePhase ->
-            if (!calculated && phase >= FirResolvePhase.CONTRACTS) {
-                executeWithoutPCE {
-                    FirLazyBodiesCalculator.calculateLazyBodiesInside(designation)
-                }
-                calculated = true
-            }
-        }
-    }
-
-    /**
-     * Designated resolve (not synchronized)
-     * @see runLazyDesignatedResolveWithoutLock for designated resolve
-     * @see lazyResolveDeclaration synchronized version
-     */
-    private fun runLazyResolveWithoutLock(
-        firDeclarationToResolve: FirDeclaration,
-        moduleFileCache: ModuleFileCache,
-        containerFirFile: FirFile,
-        provider: FirProvider,
-        fromPhase: FirResolvePhase,
-        toPhase: FirResolvePhase,
-        checkPCE: Boolean,
-    ) {
-        if (fromPhase >= toPhase) return
-        val nonLazyPhase = minOf(toPhase, LAST_NON_LAZY_PHASE)
-
-        val scopeSession = ScopeSession()
-        if (fromPhase < nonLazyPhase) {
-            firFileBuilder.runResolveWithoutLock(
-                firFile = containerFirFile,
-                fromPhase = fromPhase,
-                toPhase = nonLazyPhase,
+        for (declaration in firFile.declarations) {
+            if (checkPCE) checkCanceled()
+            lazyResolveDeclaration(
+                firDeclarationToResolve = declaration,
+                moduleFileCache = moduleFileCache,
                 scopeSession = scopeSession,
+                toPhase = toPhase,
                 checkPCE = checkPCE
             )
         }
-        if (toPhase <= nonLazyPhase) return
-        resolveFileAnnotationsWithoutLock(containerFirFile, containerFirFile.annotations, scopeSession)
-
-        runLazyDesignatedResolveWithoutLock(
-            firDeclarationToResolve = firDeclarationToResolve,
-            moduleFileCache = moduleFileCache,
-            containerFirFile = containerFirFile,
-            provider = provider,
-            fromPhase = LAST_NON_LAZY_PHASE,
-            toPhase = toPhase,
-            checkPCE = checkPCE,
-            isOnAirResolve = false
-        )
     }
 
     /**
@@ -213,93 +195,108 @@ internal class FirLazyDeclarationResolver(
      * Suitable for body resolve or/and on-air resolve.
      * @see lazyResolveDeclaration for ordinary resolve
      * @param firDeclarationToResolve target non-local declaration
-     * @param isOnAirResolve should be true when node does not belong to it's true designation (OnAir resolve in custom context)
      */
-    fun lazyDesignatedResolveDeclaration(
+    fun lazyResolveDeclaration(
         firDeclarationToResolve: FirDeclaration,
         moduleFileCache: ModuleFileCache,
-        containerFirFile: FirFile,
-        provider: FirProvider,
+        scopeSession: ScopeSession = ScopeSession(),
         toPhase: FirResolvePhase,
         checkPCE: Boolean,
-        isOnAirResolve: Boolean,
-        towerDataContextCollector: FirTowerDataContextCollector? = null,
+        declarationPhaseDowngraded: Boolean = false,
     ) {
-        // Lazy since we want to read the resolve phase inside the lock. Otherwise, we may run the same resolve phase multiple times. See
-        // KT-45121
-        val fromPhase: FirResolvePhase by lazy(LazyThreadSafetyMode.NONE) {
-            minOf(containerFirFile.resolvePhase, firDeclarationToResolve.resolvePhase)
+        if (toPhase == FirResolvePhase.RAW_FIR) return
+        //TODO Should be synchronised
+        if (!firDeclarationToResolve.isValidForResolve()) return
+
+        if (firDeclarationToResolve is FirFile) {
+            lazyResolveFileDeclaration(
+                firFile = firDeclarationToResolve,
+                moduleFileCache = moduleFileCache,
+                toPhase = toPhase,
+                scopeSession = scopeSession,
+                checkPCE = checkPCE,
+            )
+            return
         }
 
-        if (checkPCE) {
-            firFileBuilder.runCustomResolveWithPCECheck(containerFirFile, moduleFileCache) {
-                runLazyDesignatedResolveWithoutLock(
-                    firDeclarationToResolve = firDeclarationToResolve,
-                    moduleFileCache = moduleFileCache,
-                    containerFirFile = containerFirFile,
-                    provider = provider,
-                    fromPhase = fromPhase,
-                    toPhase = toPhase,
-                    checkPCE = checkPCE,
-                    isOnAirResolve = isOnAirResolve,
-                    towerDataContextCollector = towerDataContextCollector,
-                )
-            }
-        } else {
-            firFileBuilder.runCustomResolveUnderLock(containerFirFile, moduleFileCache) {
-                runLazyDesignatedResolveWithoutLock(
-                    firDeclarationToResolve = firDeclarationToResolve,
-                    moduleFileCache = moduleFileCache,
-                    containerFirFile = containerFirFile,
-                    provider = provider,
-                    fromPhase = fromPhase,
-                    toPhase = toPhase,
-                    checkPCE = checkPCE,
-                    isOnAirResolve = isOnAirResolve,
-                    towerDataContextCollector = towerDataContextCollector,
-                )
-            }
+        val provider = firDeclarationToResolve.moduleData.session.firIdeProvider
+        val resolvableDeclaration = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
+        //TODO Should be synchronised
+        if (!resolvableDeclaration.isValidForResolve()) return
+        val containerFirFile = resolvableDeclaration.getContainingFile()
+            ?: error("FirFile was not found for\n${resolvableDeclaration.render()}")
+
+        val designation = resolvableDeclaration.collectDesignation(containerFirFile)
+        //TODO Should be synchronised
+        val resolvePhase = designation.resolvePhaseForAllDeclarations(includeDeclarationPhase = declarationPhaseDowngraded)
+        if (resolvePhase >= toPhase) return
+
+        runCustomResolveUnderLock(containerFirFile, moduleFileCache, checkPCE) {
+            runLazyDesignatedResolveWithoutLock(
+                designation = designation,
+                moduleFileCache = moduleFileCache,
+                scopeSession = scopeSession,
+                toPhase = toPhase,
+                checkPCE = checkPCE,
+                declarationPhaseDowngraded = declarationPhaseDowngraded,
+            )
         }
     }
 
-    /**
-     * Designated resolve (not synchronized)
-     * @see runLazyResolveWithoutLock for ordinary resolve
-     * @see lazyDesignatedResolveDeclaration synchronized version
-     */
     private fun runLazyDesignatedResolveWithoutLock(
-        firDeclarationToResolve: FirDeclaration,
+        designation: FirDeclarationUntypedDesignationWithFile,
         moduleFileCache: ModuleFileCache,
-        containerFirFile: FirFile,
-        provider: FirProvider,
-        fromPhase: FirResolvePhase,
+        scopeSession: ScopeSession,
         toPhase: FirResolvePhase,
         checkPCE: Boolean,
-        isOnAirResolve: Boolean,
-        towerDataContextCollector: FirTowerDataContextCollector? = null
+        declarationPhaseDowngraded: Boolean,
     ) {
-        var currentPhase = fromPhase
-        runLazyResolveWithoutLock(
-            firDeclarationToResolve = firDeclarationToResolve,
-            moduleFileCache = moduleFileCache,
-            containerFirFile = containerFirFile,
-            provider = provider,
-            fromPhase = currentPhase,
-            toPhase = FirResolvePhase.IMPORTS,
-            checkPCE = checkPCE
-        )
-        currentPhase = maxOf(fromPhase, FirResolvePhase.IMPORTS)
-        if (currentPhase >= toPhase) return
-
-        val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
-        val designation = nonLocalDeclarationToResolve.collectDesignation(containerFirFile)
         check(!designation.isLocalDesignation) { "Could not resolve local designation" }
 
-        val lazyBodiesCalculator = createLazyBodiesCalculator(designation)
-        val scopeSession = ScopeSession()
+        val filePhase = designation.firFile.resolvePhase
+        if (filePhase == FirResolvePhase.RAW_FIR) {
+            lazyResolveFileDeclarationWithoutLock(
+                firFile = designation.firFile,
+                moduleFileCache = moduleFileCache,
+                toPhase = FirResolvePhase.IMPORTS,
+                scopeSession = scopeSession,
+                checkPCE = checkPCE
+            )
+        }
+        if (toPhase == FirResolvePhase.IMPORTS) return
 
-        //This needed to override standard symbol resolve in supertype transformer with adding on-air created symbols
-        val firProviderInterceptor = isOnAirResolve.ifTrue {
+        val designationPhase = designation.resolvePhaseForAllDeclarations(includeDeclarationPhase = declarationPhaseDowngraded)
+        var currentPhase = maxOf(designationPhase, FirResolvePhase.IMPORTS)
+
+        while (currentPhase < toPhase) {
+            currentPhase = currentPhase.next
+            if (currentPhase.pluginPhase) continue
+            if (checkPCE) checkCanceled()
+            FirLazyBodiesCalculator.calculateLazyBodiesInsideIfNeeded(designation, currentPhase)
+
+            currentPhase.createLazyTransformer(
+                designation = designation,
+                scopeSession = scopeSession,
+                declarationPhaseDowngraded = declarationPhaseDowngraded,
+                moduleFileCache = moduleFileCache,
+                towerDataContextCollector = null,
+                firProviderInterceptor = null,
+                checkPCE = checkPCE,
+            ).transformDeclaration(firFileBuilder.firPhaseRunner)
+        }
+    }
+
+    internal fun runLazyDesignatedOnAirResolveToBodyWithoutLock(
+        designation: FirDeclarationUntypedDesignationWithFile,
+        moduleFileCache: ModuleFileCache,
+        checkPCE: Boolean,
+        resolveWithUnchangedFir: Boolean,
+        towerDataContextCollector: FirTowerDataContextCollector?,
+    ) {
+        val scopeSession = ScopeSession()
+        var currentPhase = maxOf(designation.declaration.resolvePhase, FirResolvePhase.IMPORTS)
+
+        val firProviderInterceptor = resolveWithUnchangedFir.ifFalse {
             FirProviderInterceptorForIDE.createForFirElement(
                 session = designation.firFile.moduleData.session,
                 firFile = designation.firFile,
@@ -307,107 +304,142 @@ internal class FirLazyDeclarationResolver(
             )
         }
 
-        while (currentPhase < toPhase) {
+        while (currentPhase < FirResolvePhase.BODY_RESOLVE) {
             currentPhase = currentPhase.next
             if (currentPhase.pluginPhase) continue
             if (checkPCE) checkCanceled()
-            lazyBodiesCalculator(currentPhase)
-            runLazyResolvePhase(
-                phase = currentPhase,
-                scopeSession = scopeSession,
+
+            currentPhase.createLazyTransformer(
                 designation = designation,
+                scopeSession = scopeSession,
+                declarationPhaseDowngraded = true,
+                moduleFileCache = moduleFileCache,
                 towerDataContextCollector = towerDataContextCollector,
-                firProviderInterceptor = firProviderInterceptor
-            )
-        }
-    }
-
-    private fun runLazyResolvePhase(
-        phase: FirResolvePhase,
-        scopeSession: ScopeSession,
-        designation: FirDeclarationUntypedDesignationWithFile,
-        towerDataContextCollector: FirTowerDataContextCollector?,
-        firProviderInterceptor: FirProviderInterceptor?
-    ) {
-        if (designation.declaration.resolvePhase >= phase) return
-
-        val transformer = phase.createLazyTransformer(
-            designation,
-            scopeSession,
-            towerDataContextCollector,
-            firProviderInterceptor,
-        )
-
-        firFileBuilder.firPhaseRunner.runPhaseWithCustomResolve(phase) {
-            transformer.transformDeclaration()
+                firProviderInterceptor = firProviderInterceptor,
+                checkPCE = checkPCE,
+            ).transformDeclaration(firFileBuilder.firPhaseRunner)
         }
     }
 
     private fun FirResolvePhase.createLazyTransformer(
         designation: FirDeclarationUntypedDesignationWithFile,
         scopeSession: ScopeSession,
+        declarationPhaseDowngraded: Boolean,
+        moduleFileCache: ModuleFileCache,
         towerDataContextCollector: FirTowerDataContextCollector?,
         firProviderInterceptor: FirProviderInterceptor?,
+        checkPCE: Boolean,
     ): FirLazyTransformerForIDE = when (this) {
-        FirResolvePhase.SUPER_TYPES -> FirDesignatedSupertypeResolverTransformerForIDE(
-            designation,
-            designation.firFile.moduleData.session,
-            scopeSession,
-            firProviderInterceptor,
-        )
         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
+            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
+            firProviderInterceptor,
         )
         else -> error("Non-lazy phase $this")
     }
 
     private fun FirDeclaration.getNonLocalDeclarationToResolve(provider: FirProvider, moduleFileCache: ModuleFileCache): FirDeclaration {
         if (this is FirFile) return this
-        val ktDeclaration = psi as? KtDeclaration ?: error("FirDeclaration should have a PSI of type KtDeclaration")
-        if (declarationCanBeLazilyResolved(ktDeclaration)) 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 {
-        private val LAST_NON_LAZY_PHASE = FirResolvePhase.STATUS
+        /**
+         * 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 -> {
-                    if (declaration.hasValOrVar()) declaration.containingClassOrObject?.getClassId() != null
-                    else false
-                }
+                is KtParameter -> declaration.hasValOrVar() && declaration.containingClassOrObject?.getClassId() != null
                 is KtCallableDeclaration, is KtEnumEntry -> {
                     when (val parent = declaration.parent) {
                         is KtFile -> true

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

@@ -15,30 +15,70 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
+import org.jetbrains.kotlin.idea.util.ifTrue
 import org.jetbrains.kotlin.psi.*
 
 internal class RawFirNonLocalDeclarationBuilder private constructor(
     session: FirSession,
     baseScopeProvider: FirScopeProvider,
     private val declarationToBuild: KtDeclaration,
+    private val functionsToRebind: Set<FirFunction<*>>? = null,
     private val replacementApplier: RawFirReplacement.Applier? = null
 ) : RawFirBuilder(session, baseScopeProvider, RawFirBuilderMode.NORMAL) {
 
     companion object {
-        fun build(
+        fun buildWithReplacement(
             session: FirSession,
-            baseScopeProvider: FirScopeProvider,
+            scopeProvider: FirScopeProvider,
             designation: FirDeclarationUntypedDesignation,
             rootNonLocalDeclaration: KtDeclaration,
-            replacement: RawFirReplacement? = null
+            replacement: RawFirReplacement?
         ): FirDeclaration {
+            require(!designation.isLocalDesignation) { "Local declaration is not allowed here" }
+
             val replacementApplier = replacement?.Applier()
-            val builder = RawFirNonLocalDeclarationBuilder(session, baseScopeProvider, rootNonLocalDeclaration, replacementApplier)
+            val builder = RawFirNonLocalDeclarationBuilder(
+                session = session,
+                baseScopeProvider = scopeProvider,
+                declarationToBuild = rootNonLocalDeclaration,
+                replacementApplier = replacementApplier
+            )
             builder.context.packageFqName = rootNonLocalDeclaration.containingKtFile.packageFqName
             return builder.moveNext(designation.path.iterator(), containingClass = null).also {
                 replacementApplier?.ensureApplied()
             }
         }
+
+        fun buildWithRebind(
+            session: FirSession,
+            scopeProvider: FirScopeProvider,
+            designation: FirDeclarationUntypedDesignation,
+            rootNonLocalDeclaration: KtDeclaration,
+        ): FirDeclaration {
+            require(!designation.isLocalDesignation) { "Local declaration is not allowed here" }
+
+            val functionsToRebind = when (val originalDeclaration = designation.declaration) {
+                is FirSimpleFunction -> setOf(originalDeclaration)
+                is FirProperty -> setOfNotNull(originalDeclaration.getter, originalDeclaration.setter)
+                else -> null
+            }
+
+            val builder = RawFirNonLocalDeclarationBuilder(
+                session = session,
+                baseScopeProvider = scopeProvider,
+                declarationToBuild = rootNonLocalDeclaration,
+                functionsToRebind = functionsToRebind,
+            )
+            builder.context.packageFqName = rootNonLocalDeclaration.containingKtFile.packageFqName
+            return builder.moveNext(designation.path.iterator(), containingClass = null)
+        }
+    }
+
+    override fun bindFunctionTarget(target: FirFunctionTarget, function: FirFunction<*>) {
+        val rewrittenTarget = functionsToRebind?.firstOrNull { it.realPsi == function.realPsi } ?: function
+        super.bindFunctionTarget(target, rewrittenTarget)
     }
 
     private inner class VisitorWithReplacement : Visitor() {

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

@@ -6,24 +6,30 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensureTargetPhase
+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.ensurePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhaseForClasses
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.isTargetCallableDeclarationAndInPhase
 
+/**
+ * Transform designation into BODY_RESOLVE declaration. Affects only for target declaration and it's children
+ */
 internal class FirDesignatedBodyResolveTransformerForIDE(
     private val designation: FirDeclarationUntypedDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
-    towerDataContextCollector: FirTowerDataContextCollector? = null,
-    firProviderInterceptor: FirProviderInterceptor? = null,
+    private val declarationPhaseDowngraded: Boolean,
+    towerDataContextCollector: FirTowerDataContextCollector?,
+    firProviderInterceptor: FirProviderInterceptor?,
 ) : FirLazyTransformerForIDE, FirBodyResolveTransformer(
     session,
     phase = FirResolvePhase.BODY_RESOLVE,
@@ -40,23 +46,45 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
 ) {
     private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
-    @Suppress("NAME_SHADOWING")
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration =
         ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) {
             super.transformDeclarationContent(declaration, data)
         }
 
     override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
-        ideDeclarationTransformer.needReplacePhase
+        ideDeclarationTransformer.needReplacePhase && firDeclaration !is FirFile && super.needReplacePhase(firDeclaration)
+
+    override fun transformDeclaration(phaseRunner: FirPhaseRunner) {
+        if (designation.isResolvedForAllDeclarations(FirResolvePhase.BODY_RESOLVE, declarationPhaseDowngraded)) return
+        designation.declaration.updateResolvedPhaseForDeclarationAndChildren(FirResolvePhase.BODY_RESOLVE)
+        if (designation.isTargetCallableDeclarationAndInPhase(FirResolvePhase.BODY_RESOLVE)) return
+
+        (designation.declaration as? FirCallableDeclaration<*>)?.ensurePhase(FirResolvePhase.CONTRACTS)
+        designation.ensurePhaseForClasses(FirResolvePhase.STATUS)
+
+        phaseRunner.runPhaseWithCustomResolve(FirResolvePhase.BODY_RESOLVE) {
+            designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextIndependent)
+        }
 
-    override fun transformDeclaration() {
-        if (designation.declaration.resolvePhase >= FirResolvePhase.BODY_RESOLVE) return
-        designation.ensureTargetPhase(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
-        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextIndependent)
         ideDeclarationTransformer.ensureDesignationPassed()
         //TODO Figure out why the phase is not updated
         (designation.declaration as? FirTypeAlias)?.replaceResolvePhase(FirResolvePhase.BODY_RESOLVE)
-        designation.ensureTargetPhase(FirResolvePhase.BODY_RESOLVE)
+        ensureResolved(designation.declaration)
+        ensureResolvedDeep(designation.declaration)
+    }
+
+    override fun ensureResolved(declaration: FirDeclaration) {
+        when (declaration) {
+            is FirSimpleFunction, is FirConstructor, is FirTypeAlias, is FirField, is FirAnonymousInitializer ->
+                declaration.ensurePhase(FirResolvePhase.BODY_RESOLVE)
+            is FirProperty -> {
+                declaration.ensurePhase(FirResolvePhase.BODY_RESOLVE)
+                declaration.getter?.ensurePhase(FirResolvePhase.BODY_RESOLVE)
+                declaration.setter?.ensurePhase(FirResolvePhase.BODY_RESOLVE)
+            }
+            is FirEnumEntry, is FirClass<*> -> Unit
+            else -> error("Unexpected type: ${declaration::class.simpleName}")
+        }
     }
 }
 

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

@@ -11,44 +11,69 @@ import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveTransformer
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensureTargetPhase
+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.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhase
 
+/**
+ * Transform designation into CONTRACTS declaration. Affects only for target declaration and it's children
+ */
 internal class FirDesignatedContractsResolveTransformerForIDE(
     private val designation: FirDeclarationUntypedDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
+    private val declarationPhaseDowngraded: Boolean,
 ) : FirLazyTransformerForIDE, FirContractResolveTransformer(session, scopeSession) {
+
     private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
     override val declarationsTransformer: FirDeclarationsResolveTransformer = object : FirDeclarationsContractResolveTransformer(this) {
         override fun transformDeclarationContent(firClass: FirClass<*>, data: ResolutionMode) {
             ideDeclarationTransformer.transformDeclarationContent(this, firClass, data) {
                 super.transformDeclarationContent(firClass, data)
+                firClass
             }
         }
     }
 
-
-    @Suppress("NAME_SHADOWING")
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration =
         ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) {
+            designation.declaration.updateResolvedPhaseForDeclarationAndChildren(FirResolvePhase.CONTRACTS)
             super.transformDeclarationContent(declaration, data)
         }
 
-    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
+        ideDeclarationTransformer.needReplacePhase && firDeclaration !is FirFile && super.needReplacePhase(firDeclaration)
 
-    override fun transformDeclaration() {
-        if (designation.declaration.resolvePhase >= FirResolvePhase.CONTRACTS) return
-        val typeAlias = designation.declaration as? FirTypeAlias
-        if (typeAlias != null) {
-            //Nothing to do with typealias to CONTRACTS
-            typeAlias.replaceResolvePhase(FirResolvePhase.CONTRACTS)
-            return
+    override fun transformDeclaration(phaseRunner: FirPhaseRunner) {
+        if (designation.isResolvedForAllDeclarations(FirResolvePhase.CONTRACTS, declarationPhaseDowngraded)) return
+        designation.declaration.updateResolvedPhaseForDeclarationAndChildren(FirResolvePhase.CONTRACTS)
+        if (designation.isTargetCallableDeclarationAndInPhase(FirResolvePhase.CONTRACTS)) return
+        designation.ensureDesignation(FirResolvePhase.STATUS)
+
+        phaseRunner.runPhaseWithCustomResolve(FirResolvePhase.CONTRACTS) {
+            designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextIndependent)
         }
-        designation.ensureTargetPhase(FirResolvePhase.STATUS)
-        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextIndependent)
+
         ideDeclarationTransformer.ensureDesignationPassed()
-        designation.ensureTargetPhase(FirResolvePhase.CONTRACTS)
+        ensureResolved(designation.declaration)
+        ensureResolvedDeep(designation.declaration)
+    }
+
+    override fun ensureResolved(declaration: FirDeclaration) {
+        when (declaration) {
+            is FirSimpleFunction, is FirConstructor, is FirAnonymousInitializer ->
+                declaration.ensurePhase(FirResolvePhase.CONTRACTS)
+            is FirProperty -> {
+                declaration.ensurePhase(FirResolvePhase.CONTRACTS)
+                declaration.getter?.ensurePhase(FirResolvePhase.CONTRACTS)
+                declaration.setter?.ensurePhase(FirResolvePhase.CONTRACTS)
+            }
+            is FirClass<*>, is FirTypeAlias, is FirEnumEntry, is FirField -> Unit
+            else -> error("Unexpected type: ${declaration::class.simpleName}")
+        }
     }
 }

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

@@ -14,15 +14,24 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataCo
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensureTargetPhase
+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.ensurePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhaseForClasses
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.isTargetCallableDeclarationAndInPhase
 
+/**
+ * Transform designation into IMPLICIT_TYPES_BODY_RESOLVE declaration. Affects only for target declaration, it's children and dependents
+ */
 internal class FirDesignatedImplicitTypesTransformerForIDE(
     private val designation: FirDeclarationUntypedDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
-    towerDataContextCollector: FirTowerDataContextCollector? = null,
+    private val declarationPhaseDowngraded: Boolean,
+    towerDataContextCollector: FirTowerDataContextCollector?,
     implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession = ImplicitBodyResolveComputationSession(),
 ) : FirLazyTransformerForIDE, FirImplicitAwareBodyResolveTransformer(
     session,
@@ -40,41 +49,47 @@ internal class FirDesignatedImplicitTypesTransformerForIDE(
 ) {
     private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
-    @Suppress("NAME_SHADOWING")
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration =
         ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) {
             super.transformDeclarationContent(declaration, data)
         }
 
-    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
+        ideDeclarationTransformer.needReplacePhase && firDeclaration !is FirFile && super.needReplacePhase(firDeclaration)
 
-    override fun transformDeclaration() {
-        if (designation.declaration.resolvePhase >= FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) return
-        designation.ensureTargetPhase(FirResolvePhase.CONTRACTS)
-
-        when (val declaration = designation.declaration) {
-            is FirCallableDeclaration<*> -> {
-                //We don't need resolve callable declaration if it is already resolved (for ex. with TYPES)
-                if (declaration.returnTypeRef is FirResolvedTypeRef) {
-                    declaration.replaceResolvePhase(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
-                    return
-                }
-            }
-            is FirTypeAlias -> {
-                //Nothing to do with type alias to this phase
-                declaration.replaceResolvePhase(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
-                return
-            }
-        }
-
-        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextIndependent)
-        ideDeclarationTransformer.ensureDesignationPassed()
+    override fun transformDeclaration(phaseRunner: FirPhaseRunner) {
+        if (designation.isResolvedForAllDeclarations(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE, declarationPhaseDowngraded)) return
+        designation.declaration.updateResolvedPhaseForDeclarationAndChildren(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
+        if (designation.isTargetCallableDeclarationAndInPhase(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)) return
 
         val callableDeclaration = designation.declaration as? FirCallableDeclaration<*>
-        check(callableDeclaration == null || callableDeclaration.returnTypeRef is FirResolvedTypeRef) {
-            "Callable declaration seems to be unresolved after ${FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE} phase"
+        if (callableDeclaration != null) {
+            if (callableDeclaration.returnTypeRef is FirResolvedTypeRef) return
+            callableDeclaration.ensurePhase(FirResolvePhase.CONTRACTS)
+        }
+        designation.ensurePhaseForClasses(FirResolvePhase.STATUS)
+
+        phaseRunner.runPhaseWithCustomResolve(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
+            designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextIndependent)
         }
 
-        designation.ensureTargetPhase(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
+        ideDeclarationTransformer.ensureDesignationPassed()
+        ensureResolved(designation.declaration)
+        ensureResolvedDeep(designation.declaration)
+    }
+
+    override fun ensureResolved(declaration: FirDeclaration) {
+        when (declaration) {
+            is FirSimpleFunction -> check(declaration.returnTypeRef is FirResolvedTypeRef)
+            is FirField -> check(declaration.returnTypeRef is FirResolvedTypeRef)
+            is FirClass<*>, is FirConstructor, is FirTypeAlias, is FirEnumEntry, is FirAnonymousInitializer -> Unit
+            is FirProperty -> {
+                check(declaration.returnTypeRef is FirResolvedTypeRef)
+                //Not resolved for some getters and setters #KT-46995
+//                check(declaration.getter?.returnTypeRef?.let { it is FirResolvedTypeRef } ?: true)
+//                check(declaration.setter?.returnTypeRef?.let { it is FirResolvedTypeRef } ?: true)
+            }
+            else -> error("Unexpected type: ${declaration::class.simpleName}")
+        }
     }
 }

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

@@ -5,72 +5,76 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
 
+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.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.FirDesignatedStatusResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.FirStatusResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.StatusComputationSession
-import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.targetContainingDeclaration
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePathPhase
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensureTargetPhase
+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.ensureDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhase
 
+/**
+ * Transform designation into STATUS phase. Affects only for designation, target declaration, it's children and dependents
+ */
 internal class FirDesignatedStatusResolveTransformerForIDE(
     private val designation: FirDeclarationUntypedDesignationWithFile,
     private val session: FirSession,
     private val scopeSession: ScopeSession,
+    private val declarationPhaseDowngraded: Boolean,
 ) : FirLazyTransformerForIDE {
+    private inner class FirDesignatedStatusResolveTransformerForIDE :
+        FirStatusResolveTransformer(session, scopeSession, StatusComputationSession.Regular()) {
 
-    private fun resolveTopLevelDeclaration(declaration: FirDeclaration) {
-        val transformer = FirStatusResolveTransformer(
-            session = session,
-            scopeSession = scopeSession,
-            statusComputationSession = StatusComputationSession.Regular()
-        )
-        declaration.transformSingle(transformer, null)
-    }
+        val designationTransformer = IDEDeclarationTransformer(designation)
 
-    private fun resolveClassMember(containingClass: FirClass<*>, targetDeclaration: FirDeclaration) {
+        override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
+            firDeclaration !is FirFile && super.needReplacePhase(firDeclaration)
 
-        val transformer = object : FirDesignatedStatusResolveTransformer(
-            session = session,
-            scopeSession = scopeSession,
-            designation = designation.toSequence(includeTarget = true).iterator(),
-            targetClass = if (targetDeclaration is FirRegularClass) targetDeclaration else containingClass,
-            statusComputationSession = StatusComputationSession.Regular(),
-            designationMapForLocalClasses = emptyMap(),
-            scopeForLocalClass = null
-        ) {
-
-            override fun <F : FirClass<F>> transformClass(klass: FirClass<F>, data: FirResolvedDeclarationStatus?): FirStatement {
-                if (klass != containingClass) return super.transformClass(klass, data)
-                val result = storeClass(klass) {
-                    targetDeclaration.transformSingle(this, data)
-                }
-                return result as FirStatement
+        override fun transformDeclarationContent(declaration: FirDeclaration, data: FirResolvedDeclarationStatus?): FirDeclaration =
+            designationTransformer.transformDeclarationContent(this, declaration, data) {
+                super.transformDeclarationContent(declaration, data)
             }
-        }
-
-        val firstItemInDesignation = designation.path.firstOrNull() ?: designation.declaration
-        firstItemInDesignation.transformSingle(transformer, null)
     }
 
-    override fun transformDeclaration() {
-        if (designation.declaration.resolvePhase >= FirResolvePhase.STATUS) return
-        designation.ensurePathPhase(FirResolvePhase.TYPES)
-        designation.ensureTargetPhase(FirResolvePhase.TYPES)
+    override fun transformDeclaration(phaseRunner: FirPhaseRunner) {
+        if (designation.isResolvedForAllDeclarations(FirResolvePhase.STATUS, declarationPhaseDowngraded)) return
+        designation.declaration.updateResolvedPhaseForDeclarationAndChildren(FirResolvePhase.STATUS)
+        designation.ensureDesignation(FirResolvePhase.TYPES)
 
-        val containingClass = designation.targetContainingDeclaration()
-        if (containingClass == null) {
-            resolveTopLevelDeclaration(designation.declaration)
-        } else {
-            check(containingClass is FirClass<*>) { "Invalid designation - the parent is not a class" }
-            resolveClassMember(containingClass, designation.declaration)
+        val transformer = FirDesignatedStatusResolveTransformerForIDE()
+        phaseRunner.runPhaseWithCustomResolve(FirResolvePhase.STATUS) {
+            designation.firFile.transform<FirElement, FirResolvedDeclarationStatus?>(transformer, null)
         }
 
-        designation.ensureTargetPhase(FirResolvePhase.STATUS)
+        transformer.designationTransformer.ensureDesignationPassed()
+        designation.path.forEach(::ensureResolved)
+        ensureResolved(designation.declaration)
+        ensureResolvedDeep(designation.declaration)
+    }
+
+    override fun ensureResolved(declaration: FirDeclaration) {
+        if (declaration !is FirAnonymousInitializer) {
+            declaration.ensurePhase(FirResolvePhase.STATUS)
+        }
+        when (declaration) {
+            is FirSimpleFunction -> check(declaration.status is FirResolvedDeclarationStatus)
+            is FirConstructor -> check(declaration.status is FirResolvedDeclarationStatus)
+            is FirTypeAlias -> check(declaration.status is FirResolvedDeclarationStatus)
+            is FirEnumEntry -> check(declaration.status is FirResolvedDeclarationStatus)
+            is FirField -> check(declaration.status is FirResolvedDeclarationStatus)
+            is FirProperty -> {
+                check(declaration.status is FirResolvedDeclarationStatus)
+                check(declaration.getter?.status?.let { it is FirResolvedDeclarationStatus } ?: true)
+                check(declaration.setter?.status?.let { it is FirResolvedDeclarationStatus } ?: true)
+            }
+            is FirRegularClass -> check(declaration.status is FirResolvedDeclarationStatus)
+            is FirAnonymousInitializer -> Unit
+            else -> error("Unexpected type: ${declaration::class.simpleName}")
+        }
     }
 }

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

@@ -7,40 +7,187 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.FirApplySupertypesTransformer
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
-import org.jetbrains.kotlin.fir.resolve.transformers.FirSupertypeResolverVisitor
-import org.jetbrains.kotlin.fir.resolve.transformers.SupertypeComputationSession
+import org.jetbrains.kotlin.fir.resolve.transformers.*
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.toSymbol
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePathPhase
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensureTargetPhaseIfClass
+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.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
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhase
 
+/**
+ * Transform designation into SUPER_TYPES phase. Affects only for designation, target declaration, it's children and dependents
+ */
 internal class FirDesignatedSupertypeResolverTransformerForIDE(
     private val designation: FirDeclarationUntypedDesignationWithFile,
     private val session: FirSession,
     private val scopeSession: ScopeSession,
-    private val firProviderInterceptor: FirProviderInterceptor?
+    private val declarationPhaseDowngraded: Boolean,
+    private val moduleFileCache: ModuleFileCache,
+    private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
+    private val firProviderInterceptor: FirProviderInterceptor?,
+    private val checkPCE: Boolean,
 ) : FirLazyTransformerForIDE {
 
-    override fun transformDeclaration() {
-        if (designation.declaration.resolvePhase >= FirResolvePhase.SUPER_TYPES) return
-        designation.ensurePathPhase(FirResolvePhase.SUPER_TYPES)
+    private val supertypeComputationSession = SupertypeComputationSession()
 
-        val supertypeComputationSession = SupertypeComputationSession()
-
-        val resolver = FirSupertypeResolverVisitor(
+    private inner class DesignatedFirSupertypeResolverVisitor(classDesignation: FirDeclarationUntypedDesignation) :
+        FirSupertypeResolverVisitor(
             session = session,
             supertypeComputationSession = supertypeComputationSession,
             scopeSession = scopeSession,
             scopeForLocalClass = null,
             localClassesNavigationInfo = null,
             firProviderInterceptor = firProviderInterceptor,
-        )
-        designation.declaration.accept(resolver, null)
-        val applySupertypesTransformer = FirApplySupertypesTransformer(supertypeComputationSession)
-        designation.declaration.transform<FirElement, Void?>(applySupertypesTransformer, null)
-        designation.ensureTargetPhaseIfClass(FirResolvePhase.SUPER_TYPES)
+        ) {
+        val declarationTransformer = IDEDeclarationTransformer(classDesignation)
+
+        override fun visitDeclarationContent(declaration: FirDeclaration, data: Any?) {
+            declarationTransformer.visitDeclarationContent(this, declaration, data) {
+                super.visitDeclarationContent(declaration, data)
+                declaration
+            }
+        }
+    }
+
+    private inner class DesignatedFirApplySupertypesTransformer(classDesignation: FirDeclarationUntypedDesignation) :
+        FirApplySupertypesTransformer(supertypeComputationSession) {
+
+        override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
+            firDeclaration !is FirFile && super.needReplacePhase(firDeclaration)
+
+        override fun transformRegularClass(regularClass: FirRegularClass, data: Any?): FirStatement {
+            return if (regularClass.resolvePhase >= FirResolvePhase.SUPER_TYPES)
+                transformDeclarationContent(regularClass, data) as FirStatement
+            else super.transformRegularClass(regularClass, data)
+        }
+
+        override fun transformAnonymousObject(anonymousObject: FirAnonymousObject, data: Any?): FirStatement {
+            return if (anonymousObject.resolvePhase >= FirResolvePhase.SUPER_TYPES)
+                transformDeclarationContent(anonymousObject, data) as FirStatement
+            else super.transformAnonymousObject(anonymousObject, data)
+        }
+
+        override fun transformTypeAlias(typeAlias: FirTypeAlias, data: Any?): FirDeclaration {
+            return if (typeAlias.resolvePhase >= FirResolvePhase.SUPER_TYPES)
+                transformDeclarationContent(typeAlias, data)
+            else super.transformTypeAlias(typeAlias, data)
+        }
+
+        val declarationTransformer = IDEDeclarationTransformer(classDesignation)
+
+        override fun transformDeclarationContent(declaration: FirDeclaration, data: Any?): FirDeclaration {
+            return declarationTransformer.transformDeclarationContent(this, declaration, data) {
+                super.transformDeclarationContent(declaration, data)
+            }
+        }
+    }
+
+    private fun collect(designation: FirDeclarationUntypedDesignationWithFile): Collection<FirDeclarationUntypedDesignationWithFile> {
+        val visited = mutableMapOf<FirDeclaration, FirDeclarationUntypedDesignationWithFile>()
+        val toVisit = mutableListOf<FirDeclarationUntypedDesignationWithFile>()
+        toVisit.add(designation)
+
+        while (toVisit.isNotEmpty()) {
+            for (nowVisit in toVisit) {
+                if (checkPCE) checkCanceled()
+                val resolver = DesignatedFirSupertypeResolverVisitor(nowVisit)
+                runCustomResolveUnderLock(nowVisit.firFile, moduleFileCache, checkPCE) {
+                    firLazyDeclarationResolver.lazyResolveFileDeclaration(
+                        firFile = nowVisit.firFile,
+                        moduleFileCache = moduleFileCache,
+                        toPhase = FirResolvePhase.IMPORTS,
+                        scopeSession = scopeSession,
+                        checkPCE = false,
+                    )
+                    nowVisit.firFile.accept(resolver, null)
+                }
+                resolver.declarationTransformer.ensureDesignationPassed()
+                visited[nowVisit.declaration] = nowVisit
+            }
+            toVisit.clear()
+
+            for (value in supertypeComputationSession.supertypeStatusMap.values) {
+                if (value !is SupertypeComputationStatus.Computed) continue
+                for (reference in value.supertypeRefs) {
+                    val classLikeDeclaration = reference.type.toSymbol(session)?.fir
+                    if (classLikeDeclaration !is FirClassLikeDeclaration<*>) continue
+                    if (visited.containsKey(classLikeDeclaration)) continue
+                    val containingFile = moduleFileCache.getContainerFirFile(classLikeDeclaration) ?: continue
+                    toVisit.add(classLikeDeclaration.collectDesignation(containingFile))
+                }
+            }
+        }
+        return visited.values
+    }
+
+    private fun apply(visited: Collection<FirDeclarationUntypedDesignationWithFile>) {
+        fun applyToFileSymbols(designations: List<FirDeclarationUntypedDesignationWithFile>) {
+            for (designation in designations) {
+                if (checkPCE) checkCanceled()
+                val applier = DesignatedFirApplySupertypesTransformer(designation)
+                designation.firFile.transform<FirElement, Void?>(applier, null)
+                applier.declarationTransformer.ensureDesignationPassed()
+            }
+        }
+
+        val filesToDesignations = visited.groupBy { it.firFile }
+        for (designationsPerFile in filesToDesignations) {
+            if (checkPCE) checkCanceled()
+            runCustomResolveUnderLock(designationsPerFile.key, moduleFileCache, checkPCE) {
+                applyToFileSymbols(designationsPerFile.value)
+            }
+        }
+    }
+
+    override fun transformDeclaration(phaseRunner: FirPhaseRunner) {
+        check(designation.firFile.resolvePhase >= FirResolvePhase.IMPORTS) {
+            "Invalid resolve phase of file. Should be IMPORTS but found ${designation.firFile.resolvePhase}"
+        }
+
+        val targetDesignation = if (designation.declaration !is FirClassLikeDeclaration<*>) {
+            val resolvableTarget = designation.path.lastOrNull() ?: return
+            check(resolvableTarget is FirClassLikeDeclaration<*>)
+            val targetPath = designation.path.dropLast(1)
+            FirDeclarationUntypedDesignationWithFile(targetPath, resolvableTarget, false, designation.firFile)
+        } else designation
+
+        if (targetDesignation.isResolvedForAllDeclarations(FirResolvePhase.SUPER_TYPES, declarationPhaseDowngraded)) return
+        targetDesignation.declaration.updateResolvedPhaseForDeclarationAndChildren(FirResolvePhase.SUPER_TYPES)
+
+        phaseRunner.runPhaseWithCustomResolve(FirResolvePhase.SUPER_TYPES) {
+            val collected = collect(targetDesignation)
+            supertypeComputationSession.breakLoops(session)
+            apply(collected)
+        }
+
+        designation.path.forEach(::ensureResolved)
+        ensureResolved(designation.declaration)
+        ensureResolvedDeep(designation.declaration)
+    }
+
+    override fun ensureResolved(declaration: FirDeclaration) {
+        when (declaration) {
+            is FirFunction<*>, is FirProperty, is FirEnumEntry, is FirField, is FirAnonymousInitializer -> Unit
+            is FirRegularClass -> {
+                declaration.ensurePhase(FirResolvePhase.SUPER_TYPES)
+                check(declaration.superTypeRefs.all { it is FirResolvedTypeRef })
+            }
+            is FirTypeAlias -> {
+                declaration.ensurePhase(FirResolvePhase.SUPER_TYPES)
+                check(declaration.expandedTypeRef is FirResolvedTypeRef)
+            }
+            else -> error("Unexpected type: ${declaration::class.simpleName}")
+        }
     }
 }

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

@@ -10,34 +10,75 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTypeResolveTransformer
+import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePathPhase
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensureTargetPhase
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensureTargetPhaseIfClass
+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.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.ensurePhaseForClasses
 
+/**
+ * Transform designation into TYPES phase. Affects only for designation, target declaration and it's children
+ */
 internal class FirDesignatedTypeResolverTransformerForIDE(
     private val designation: FirDeclarationUntypedDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
+    private val declarationPhaseDowngraded: Boolean,
 ) : FirLazyTransformerForIDE, FirTypeResolveTransformer(session, scopeSession) {
 
-    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
+    private val declarationTransformer = IDEDeclarationTransformer(designation)
 
     override fun <E : FirElement> transformElement(element: E, data: Any?): E {
-        if (element !is FirRegularClass && element !is FirAnonymousObject && element !is FirFile)
-            return super.transformElement(element, data)
-
-        return ideDeclarationTransformer.transformDeclarationContent(this, element, data) {
+        return if (element is FirDeclaration && (element is FirRegularClass || element is FirFile)) {
+            declarationTransformer.transformDeclarationContent(this, element, data) {
+                super.transformElement(element, data)
+            }
+        } else {
             super.transformElement(element, data)
         }
     }
 
-    override fun transformDeclaration() {
-        if (designation.declaration.resolvePhase >= FirResolvePhase.TYPES) return
-        designation.ensurePathPhase(FirResolvePhase.TYPES)
-        designation.ensureTargetPhaseIfClass(FirResolvePhase.SUPER_TYPES)
-        designation.firFile.transform<FirFile, Any?>(this, null)
-        ideDeclarationTransformer.ensureDesignationPassed()
-        designation.ensureTargetPhase(FirResolvePhase.TYPES)
+    override fun transformDeclaration(phaseRunner: FirPhaseRunner) {
+        if (designation.isResolvedForAllDeclarations(FirResolvePhase.TYPES, declarationPhaseDowngraded)) return
+        designation.declaration.updateResolvedPhaseForDeclarationAndChildren(FirResolvePhase.TYPES)
+        designation.ensurePhaseForClasses(FirResolvePhase.SUPER_TYPES)
+
+        phaseRunner.runPhaseWithCustomResolve(FirResolvePhase.TYPES) {
+            designation.firFile.transform<FirFile, Any?>(this, null)
+        }
+
+        declarationTransformer.ensureDesignationPassed()
+
+        designation.path.forEach(::ensureResolved)
+        ensureResolved(designation.declaration)
+        ensureResolvedDeep(designation.declaration)
+    }
+
+    override fun ensureResolved(declaration: FirDeclaration) {
+        if (declaration !is FirAnonymousInitializer) {
+            declaration.ensurePhase(FirResolvePhase.TYPES)
+        }
+        when (declaration) {
+            is FirFunction<*> -> {
+                check(declaration.returnTypeRef is FirResolvedTypeRef || declaration.returnTypeRef is FirImplicitTypeRef)
+                check(declaration.receiverTypeRef?.let { it is FirResolvedTypeRef } ?: true)
+                declaration.valueParameters.forEach {
+                    check(it.returnTypeRef is FirResolvedTypeRef || it.returnTypeRef is FirImplicitTypeRef)
+                }
+            }
+            is FirProperty -> {
+                check(declaration.returnTypeRef is FirResolvedTypeRef || declaration.returnTypeRef is FirImplicitTypeRef)
+                check(declaration.receiverTypeRef?.let { it is FirResolvedTypeRef } ?: true)
+                declaration.getter?.run(::ensureResolved)
+                declaration.setter?.run(::ensureResolved)
+            }
+            is FirField -> check(declaration.returnTypeRef is FirResolvedTypeRef || declaration.returnTypeRef is FirImplicitTypeRef)
+            is FirClass<*>, is FirTypeAlias, is FirAnonymousInitializer -> Unit
+            is FirEnumEntry -> check(declaration.returnTypeRef is FirResolvedTypeRef)
+            else -> error("Unexpected type: ${declaration::class.simpleName}")
+        }
     }
 }

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolve
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
 
 internal class FirFileAnnotationsResolveTransformer(
@@ -45,9 +46,10 @@ internal class FirFileAnnotationsResolveTransformer(
         return declaration
     }
 
-    override fun transformDeclaration() {
+    override fun transformDeclaration(phaseRunner: FirPhaseRunner) {
         if (annotations.all { it.resolveStatus == FirAnnotationResolveStatus.Resolved }) return
         check(firFile.resolvePhase >= FirResolvePhase.IMPORTS) { "Invalid file resolve phase ${firFile.resolvePhase}" }
+
         firFile.accept(this, ResolutionMode.ContextDependent)
         check(annotations.all { it.resolveStatus == FirAnnotationResolveStatus.Resolved }) { "Annotation was not resolved" }
     }

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

@@ -5,12 +5,69 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
 
+import org.jetbrains.annotations.TestOnly
+import org.jetbrains.kotlin.asJava.classes.lazyPub
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignation
+
 internal interface FirLazyTransformerForIDE {
-    fun transformDeclaration()
+    fun transformDeclaration(phaseRunner: FirPhaseRunner)
+    fun ensureResolved(declaration: FirDeclaration): Unit = error("Not implemented")
+    fun ensureResolvedDeep(declaration: FirDeclaration) {
+        if (!enableDeepEnsure) return
+        ensureResolved(declaration)
+        if (declaration is FirRegularClass) {
+            declaration.declarations.forEach(::ensureResolvedDeep)
+        }
+    }
 
     companion object {
+        internal var enableDeepEnsure: Boolean = false
+            @TestOnly set
+
+        private object ResolvePhaseWithForAllDeclarationsKey : FirDeclarationDataKey()
+
+        private var FirDeclaration.resolvePhaseForDeclarationAndChildrenAttr: FirResolvePhase?
+                by FirDeclarationDataRegistry.data(ResolvePhaseWithForAllDeclarationsKey)
+
+        /**
+         * This resolve phase is used to check if current declaration and it's children were resolved for phase
+         */
+        var FirDeclaration.resolvePhaseForDeclarationAndChildren: FirResolvePhase
+            get() = resolvePhaseForDeclarationAndChildrenAttr ?: FirResolvePhase.RAW_FIR
+            set(value) {
+                resolvePhaseForDeclarationAndChildrenAttr = value
+            }
+
+        fun FirDeclaration.updateResolvedPhaseForDeclarationAndChildren(phase: FirResolvePhase) {
+            val allDeclaration = resolvePhaseForDeclarationAndChildren
+            if (allDeclaration < phase) {
+                resolvePhaseForDeclarationAndChildren = phase
+            }
+        }
+
+        fun FirDeclarationUntypedDesignation.resolvePhaseForAllDeclarations(includeDeclarationPhase: Boolean): FirResolvePhase {
+            //resolvePhaseWithForAllDeclarations for these origins are derived from original declaration
+            val includeTarget = when (declaration.origin) {
+                is FirDeclarationOrigin.SubstitutionOverride,
+                is FirDeclarationOrigin.IntersectionOverride,
+                is FirDeclarationOrigin.Delegated -> false
+                else -> true
+            }
+
+            val allContaining = toSequence(includeTarget = includeTarget)
+                .maxByOrNull { it.resolvePhaseForDeclarationAndChildren }
+                ?.resolvePhaseForDeclarationAndChildren
+                ?: FirResolvePhase.RAW_FIR
+            return if (includeDeclarationPhase) minOf(declaration.resolvePhase, allContaining) else allContaining
+        }
+
+        fun FirDeclarationUntypedDesignation.isResolvedForAllDeclarations(phase: FirResolvePhase, includeDeclarationPhase: Boolean) =
+            resolvePhaseForAllDeclarations(includeDeclarationPhase) >= phase
+
         val DUMMY = object : FirLazyTransformerForIDE {
-            override fun transformDeclaration() = Unit
+            override fun transformDeclaration(phaseRunner: FirPhaseRunner) = Unit
         }
     }
 }

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

@@ -6,9 +6,12 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.transformers.FirAbstractPhaseTransformer
+import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignation
 
 internal class IDEDeclarationTransformer(private val designation: FirDeclarationUntypedDesignation) {
@@ -16,15 +19,36 @@ internal class IDEDeclarationTransformer(private val designation: FirDeclaration
     private var isInsideTargetDeclaration: Boolean = false
     private var designationPassed: Boolean = false
 
-    inline fun <K, D> transformDeclarationContent(
-        transformer: FirAbstractPhaseTransformer<D>,
+    inline fun <K : FirDeclaration, D> visitDeclarationContent(
+        visitor: FirVisitor<Unit, D>,
         declaration: K,
         data: D,
-        defaultCallTransform: () -> K
+        default: () -> K
+    ) = processDeclarationContent(declaration, default) {
+        it.accept(visitor, data)
+    }
+
+    inline fun <K : FirDeclaration, D> transformDeclarationContent(
+        transformer: FirDefaultTransformer<D>,
+        declaration: K,
+        data: D,
+        default: () -> K
+    ): K = processDeclarationContent(declaration, default) { toTransform ->
+        toTransform.transform<FirElement, D>(transformer, data).also { transformed ->
+            check(transformed === toTransform) {
+                "become $transformed `${transformed.render()}`, was ${toTransform}: `${toTransform.render()}`"
+            }
+        }
+    }
+
+    inline fun <K : FirDeclaration> processDeclarationContent(
+        declaration: K,
+        default: () -> K,
+        applyToDesignated: (FirDeclaration) -> Unit,
     ): K {
         //It means that we are inside the target declaration
         if (isInsideTargetDeclaration) {
-            return defaultCallTransform()
+            return default()
         }
 
         //It means that we already transform target declaration and now can skip all others
@@ -33,12 +57,12 @@ internal class IDEDeclarationTransformer(private val designation: FirDeclaration
         }
 
         if (designationWithoutTargetIterator.hasNext()) {
-            designationWithoutTargetIterator.next().visitNoTransform(transformer, data)
+            applyToDesignated(designationWithoutTargetIterator.next())
         } else {
             try {
                 isInsideTargetDeclaration = true
                 designationPassed = true
-                designation.declaration.visitNoTransform(transformer, data)
+                applyToDesignated(designation.declaration)
             } finally {
                 isInsideTargetDeclaration = false
             }
@@ -52,9 +76,4 @@ internal class IDEDeclarationTransformer(private val designation: FirDeclaration
     fun ensureDesignationPassed() {
         check(designationPassed) { "Designation not passed for declaration ${designation.declaration::class.simpleName}" }
     }
-}
-
-private fun <D> FirElement.visitNoTransform(transformer: FirTransformer<D>, data: D) {
-    val result = this.transform<FirElement, D>(transformer, data)
-    require(result === this) { "become $result `${result.render()}`, was ${this}: `${this.render()}`" }
-}
+}

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.firProvider
 
 fun FirDeclaration.getContainingFile(): FirFile? {
@@ -19,4 +20,4 @@ fun FirDeclaration.getContainingFile(): FirFile? {
         is FirClassLikeDeclaration<*> -> provider.getFirClassifierContainerFile(symbol)
         else -> error("Unsupported declaration ${this::class.java}")
     }
-}
+}

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

@@ -6,23 +6,32 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.util
 
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.transformers.ensureResolved
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationUntypedDesignation
 
-internal fun FirDeclarationUntypedDesignation.ensurePathPhase(firResolvePhase: FirResolvePhase) {
-    toSequence(includeTarget = false).forEach { firDeclaration ->
-        check(firDeclaration.resolvePhase >= firResolvePhase) {
-            "Designation element phase required to be $firResolvePhase but element resolved to ${firDeclaration.resolvePhase}"
+internal fun FirDeclaration.ensurePhase(firResolvePhase: FirResolvePhase) =
+    check(resolvePhase >= firResolvePhase) {
+        "Element phase required to be $firResolvePhase but element resolved to $resolvePhase"
+    }
+
+internal fun FirDeclarationUntypedDesignation.ensurePathPhase(firResolvePhase: FirResolvePhase) =
+    path.forEach { it.ensurePhase(firResolvePhase) }
+
+internal fun FirDeclarationUntypedDesignation.ensureDesignation(firResolvePhase: FirResolvePhase) {
+    ensurePathPhase(firResolvePhase)
+    declaration.ensurePhase(firResolvePhase)
+}
+internal fun FirDeclarationUntypedDesignation.ensurePhaseForClasses(firResolvePhase: FirResolvePhase) {
+    ensurePathPhase(firResolvePhase)
+    if (declaration is FirClassLikeDeclaration<*>) {
+        check(declaration.resolvePhase >= firResolvePhase) {
+            "Expected $firResolvePhase but found ${declaration.resolvePhase}"
         }
     }
 }
 
-internal fun FirDeclarationUntypedDesignation.ensureTargetPhase(firResolvePhase: FirResolvePhase) =
-    check(declaration.resolvePhase >= firResolvePhase) { "Expected $firResolvePhase but found ${declaration.resolvePhase}" }
-
-internal fun FirDeclarationUntypedDesignation.ensureTargetPhaseIfClass(firResolvePhase: FirResolvePhase) = when (declaration) {
-    is FirProperty, is FirSimpleFunction -> Unit
-    is FirClass<*>, is FirTypeAlias -> ensureTargetPhase(firResolvePhase)
-    else -> error("Unexpected target")
-}
+internal fun FirDeclarationUntypedDesignation.isTargetCallableDeclarationAndInPhase(firResolvePhase: FirResolvePhase): Boolean =
+    (declaration as? FirCallableDeclaration<*>)?.let { it.resolvePhase >= firResolvePhase } ?: false
 
 internal fun FirDeclarationUntypedDesignation.targetContainingDeclaration(): FirDeclaration? = path.lastOrNull()

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/onAirTypesResolve.txt

@@ -3,15 +3,15 @@ public final [BODY_RESOLVE] class Y : R|kotlin/Any| {
         super<R|kotlin/Any|>()
     }
 
-    public? open [BODY_RESOLVE] class BASE : R|X| {
-        public? [BODY_RESOLVE] constructor(): R|X.Y.BASE| {
+    public open [BODY_RESOLVE] class BASE : R|X| {
+        public [BODY_RESOLVE] constructor(): R|X.Y.BASE| {
             super<R|X|>()
         }
 
     }
 
-    public? final? [BODY_RESOLVE] class DERIVED : R|X.Y.BASE| {
-        public? [BODY_RESOLVE] constructor(): R|X.Y.DERIVED| {
+    public final [BODY_RESOLVE] class DERIVED : R|X.Y.BASE| {
+        public [BODY_RESOLVE] constructor(): R|X.Y.DERIVED| {
             super<R|X.Y.BASE|>()
         }
 

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/replacementInsidePropertyBody.kt

@@ -5,4 +5,8 @@ val x: Int get() {
     2
 }
 
-/*ONAIR*/val z = qqq
+val y: String get() {
+    /*ONAIR*/val z = qqq
+    ""
+}
+

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/replacementInsidePropertyBody.txt

@@ -1,2 +1 @@
-public? final? [BODY_RESOLVE] val z: R|kotlin/String| = R|<local>/qqq|
-    [BODY_RESOLVE] public? get(): R|kotlin/String|
+[BODY_RESOLVE] lval z: R|kotlin/String| = R|<local>/qqq|

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/replacementInsidePropertyBody2.kt

@@ -3,4 +3,7 @@ var x = 2
         /*PLACE*/val inside = ""
     }
 
-/*ONAIR*/val inSetter = value
+var y = 2
+    set(value) {
+        /*ONAIR*/val inSetter = value
+    }

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/onAirResolve/replacementInsidePropertyBody2.txt

@@ -1,2 +1 @@
-public? final? [BODY_RESOLVE] val inSetter: R|kotlin/Int| = R|<local>/value|
-    [BODY_RESOLVE] public? get(): R|kotlin/Int|
+[BODY_RESOLVE] lval inSetter: R|kotlin/Int| = R|<local>/value|

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

@@ -20,6 +20,7 @@ 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.FirDeclarationUntypedDesignation
 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
@@ -132,11 +133,12 @@ abstract class AbstractPartialRawFirBuilderTestCase : KtParsingTestCase(
         val designation = designationBuilder.resultDesignation
         TestCase.assertTrue(designation != null)
 
-        val firElement = RawFirNonLocalDeclarationBuilder.build(
-            session,
-            scopeProvider,
+        val firElement = RawFirNonLocalDeclarationBuilder.buildWithReplacement(
+            session = session,
+            scopeProvider = scopeProvider,
             designation!!,
-            elementToBuild
+            elementToBuild,
+            null
         )
 
         val firDump = firElement.render(FirRenderer.RenderMode.WithFqNames)

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/compiler/based/AbstractDiagnosisCompilerTestDataTest.kt

@@ -8,12 +8,13 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostic.compiler.based
 import org.jetbrains.kotlin.idea.fir.low.level.api.compiler.based.AbstractCompilerBasedTest
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.runners.baseFirDiagnosticTestConfiguration
-import org.jetbrains.kotlin.idea.fir.low.level.api.compiler.based.IdeTestIgnoreHandler
 import org.jetbrains.kotlin.idea.fir.low.level.api.compiler.based.addIdeTestIgnoreHandler
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE
 
 abstract class AbstractDiagnosisCompilerTestDataTest : AbstractCompilerBasedTest() {
     override fun TestConfigurationBuilder.configureTest() {
         baseFirDiagnosticTestConfiguration()
+        FirLazyTransformerForIDE.enableDeepEnsure = true
         addIdeTestIgnoreHandler()
     }
 }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.transformers.resolveSupertypesInTheAir
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.annotations.KtFirAnnotationCall
@@ -28,7 +29,7 @@ internal class KtFirTypeAndAnnotations<T : FirDeclaration>(
 
     override val token: ValidityToken get() = containingDeclaration.token
 
-    override val type: KtType by containingDeclaration.withFirAndCache(typeResolvePhase) { fir ->
+    override val type: KtType by containingDeclaration.withFirAndCache(ResolveType.CallableReturnType) { fir ->
         builder.typeBuilder.buildKtType(typeRef(fir))
     }
 

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.frontend.api.fir.utils.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRefWithValidityCheck
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.convertConstantExpression
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.mapAnnotationParameters
@@ -33,8 +34,10 @@ internal class KtFirAnnotationCall(
         fir.findPsi(fir.moduleData.session) as? KtCallElement
     }
 
-    override val classId: ClassId? by containingDeclaration.withFirAndCache(AnnotationPhases.PHASE_FOR_ANNOTATION_CLASS_ID) { fir ->
-        annotationCallRef.getClassId(fir.moduleData.session)
+    override val classId: ClassId? by cached {
+        containingDeclaration.withFirWithPossibleResolveInside(AnnotationPhases.PHASE_FOR_ANNOTATION_CLASS_ID) { fir ->
+            annotationCallRef.getClassId(fir.moduleData.session)
+        }
     }
 
     override val useSiteTarget: AnnotationUseSiteTarget? get() = annotationCallRef.useSiteTarget

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

@@ -8,16 +8,21 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols.annotations
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.expandedConeType
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 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.frontend.api.fir.symbols.expandTypeAliasIfNeeded
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRefWithValidityCheck
 import org.jetbrains.kotlin.name.ClassId
 
 
 internal fun FirAnnotationCall.getClassId(session: FirSession): ClassId? =
-    coneClassLikeType?.expandTypeAliasIfNeeded(session)?.classId
+    coneClassLikeType?.fullyExpandedType(session) { alias ->
+        alias.ensureResolved(FirResolvePhase.SUPER_TYPES, session)
+        alias.expandedConeType
+    }?.classId
 
 internal fun FirRefWithValidityCheck<FirAnnotatedDeclaration>.toAnnotationsList() = withFir { fir ->
     fir.annotations.map { KtFirAnnotationCall(this, it) }

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

@@ -12,11 +12,14 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.renderWithType
+import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
+import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.psi.KtDeclaration
@@ -48,10 +51,3 @@ internal fun <F : FirMemberDeclaration> KtFirSymbol<F>.getVisibility(phase: FirR
 
 internal fun KtFirSymbol<FirCallableDeclaration<*>>.getCallableIdIfNonLocal(): CallableId? =
     firRef.withFir { fir -> fir.symbol.callableId.takeUnless { it.isLocal } }
-
-internal fun ConeClassLikeType.expandTypeAliasIfNeeded(session: FirSession): ConeClassLikeType {
-    val firTypeAlias = lookupTag.toSymbol(session) as? FirTypeAliasSymbol ?: return this
-    val expandedType = firTypeAlias.fir.expandedTypeRef.coneType
-    return expandedType.fullyExpandedType(session) as? ConeClassLikeType
-        ?: return this
-}

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclarationInWriteLock
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.tokens.assertIsValidAndAccessible
@@ -86,6 +87,29 @@ internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveSt
         ValidityAwareCachedValue(token) {
             withFir(phase) { fir -> createValue(fir) }
         }
+
+    inline fun <R> withFir(type: ResolveType, crossinline action: (fir: D) -> R): R {
+        token.assertIsValidAndAccessible()
+        val fir = firWeakRef.get()
+            ?: throw EntityWasGarbageCollectedException("FirElement")
+        val resolveState = resolveStateWeakRef.get()
+            ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
+        return when (type) {
+            ResolveType.CallableBodyResolve -> {
+                /*
+                 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
+                 */
+                action(fir.withFirDeclaration(type, resolveState) { it })
+            }
+            else -> fir.withFirDeclaration(type, resolveState) { action(it) }
+        }
+    }
+
+    inline fun <R> withFirAndCache(type: ResolveType, crossinline createValue: (fir: D) -> R) =
+        ValidityAwareCachedValue(token) {
+            withFir(type) { fir -> createValue(fir) }
+        }
 }
 
 @Suppress("NOTHING_TO_INLINE")