[FIR IDE] Fix invalid body lazy resolve

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

@@ -6,13 +6,11 @@
 package org.jetbrains.kotlin.fir.resolve.transformers
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.lookupSuperTypes
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -83,7 +81,7 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
 
             // Note that annotations are still visited here
             // again, although there's no need in it
-            transformElement(firClass, data)
+            transformDeclarationContent(firClass, data) as FirClass<*>
         }
     }
 
@@ -92,6 +90,10 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
             scopes.add(FirMemberTypeParameterScope(this))
         }
     }
+
+    open fun transformDeclarationContent(declaration: FirDeclaration, data: Any?): FirDeclaration {
+        return transformElement(declaration, data)
+    }
 }
 
 fun createSubstitutionForSupertype(superType: ConeLookupTagBasedType, session: FirSession): ConeSubstitutor {

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirDelegatedConstructorCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeCyclicTypeBound
 import org.jetbrains.kotlin.fir.scopes.FirScope
@@ -38,7 +39,7 @@ fun <F : FirClassLikeDeclaration<F>> F.runTypeResolvePhaseForLocalClass(
 }
 
 open class FirTypeResolveTransformer(
-    override val session: FirSession,
+    final override val session: FirSession,
     scopeSession: ScopeSession,
     initialScopes: List<FirScope> = emptyList()
 ) : FirAbstractTreeTransformerWithSuperTypes(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformer.kt

@@ -52,7 +52,7 @@ open class FirContractResolveTransformer(
         return annotationCall
     }
 
-    private class FirDeclarationsContractResolveTransformer(transformer: FirBodyResolveTransformer) : FirDeclarationsResolveTransformer(transformer) {
+    protected open class FirDeclarationsContractResolveTransformer(transformer: FirBodyResolveTransformer) : FirDeclarationsResolveTransformer(transformer) {
         override fun transformSimpleFunction(
             simpleFunction: FirSimpleFunction,
             data: ResolutionMode
@@ -224,11 +224,15 @@ open class FirContractResolveTransformer(
             }
         }
 
+        open fun transformDeclarationContent(firClass: FirClass<*>, data: ResolutionMode) {
+            firClass.transformDeclarations(this, data)
+        }
+
         override fun transformRegularClass(regularClass: FirRegularClass, data: ResolutionMode): FirStatement {
             regularClass.updatePhase()
             regularClass.transformCompanionObject(this, data)
             context.withRegularClass(regularClass, components, forContracts = true) {
-                regularClass.transformDeclarations(this, data)
+                transformDeclarationContent(regularClass, data)
             }
             return regularClass
         }
@@ -239,7 +243,7 @@ open class FirContractResolveTransformer(
         ): FirStatement {
             anonymousObject.updatePhase()
             context.withAnonymousObject(anonymousObject, components) {
-                anonymousObject.transformDeclarations(this, data)
+                transformDeclarationContent(anonymousObject, data)
             }
             return anonymousObject
         }

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

@@ -5,6 +5,7 @@
 
 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.builder.RawFirFragmentForLazyBodiesBuilder
 import org.jetbrains.kotlin.fir.builder.RawFirReplacement
@@ -14,8 +15,6 @@ 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.providers.firIdeProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 
 object DeclarationCopyBuilder {
@@ -59,6 +58,12 @@ object DeclarationCopyBuilder {
         val originalFunction = rootNonLocalDeclaration.getOrBuildFirOfType<FirSimpleFunction>(state)
         val builtFunction = createCopy(rootNonLocalDeclaration, originalFunction, replacement)
 
+        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?
+        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
@@ -77,6 +82,12 @@ object DeclarationCopyBuilder {
         val originalFirClass = rootNonLocalDeclaration.getOrBuildFirOfType<FirRegularClass>(state)
         val builtClass = createCopy(rootNonLocalDeclaration, originalFirClass, replacement)
 
+        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?
+        val classBody = rootNonLocalDeclaration.body
+        if (classBody == null || !PsiTreeUtil.isAncestor(classBody, replacement.from, true)) {
+            return builtClass
+        }
+
         return buildRegularClassCopy(originalFirClass) {
             declarations.clear()
             declarations.addAll(builtClass.declarations)
@@ -94,12 +105,13 @@ object DeclarationCopyBuilder {
 
         val originalFirTypeAlias = rootNonLocalDeclaration.getOrBuildFirOfType<FirTypeAlias>(state)
         val builtTypeAlias = createCopy(rootNonLocalDeclaration, originalFirTypeAlias, replacement)
+        return builtTypeAlias
 
-        return buildTypeAliasCopy(originalFirTypeAlias) {
+//        return buildTypeAliasCopy(originalFirTypeAlias) {
-            expandedTypeRef = builtTypeAlias.expandedTypeRef
+//            expandedTypeRef = builtTypeAlias.expandedTypeRef
-            symbol = builtTypeAlias.symbol
+//            symbol = builtTypeAlias.symbol
-            initDeclaration(originalFirTypeAlias, builtTypeAlias, state)
+//            initDeclaration(originalFirTypeAlias, builtTypeAlias, state)
-        }
+//        }
     }
 
     private fun <T : KtElement> createPropertyCopy(
@@ -110,6 +122,18 @@ object DeclarationCopyBuilder {
         val originalProperty = rootNonLocalDeclaration.getOrBuildFirOfType<FirProperty>(state)
         val builtProperty = createCopy(rootNonLocalDeclaration, originalProperty, replacement)
 
+        //TODO Check do we really need this optimisation? Maybe it is not optimisation at all because of Kt travers?
+        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
 

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

@@ -69,7 +69,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         require(state is FirModuleResolveStateImpl)
         require(place.isPhysical)
 
-        val declaration = runResolveBodyResolveOnAir(
+        val declaration = runBodyResolveOnAir(
             state = state,
             replacement = RawFirReplacement(place, elementToResolve),
             collector = null,
@@ -102,7 +102,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             FileTowerProvider(place, onAirGetTowerContextForFile(state, place))
         } else {
             FirTowerDataContextAllElementsCollector().also {
-                runResolveBodyResolveOnAir(state, collector = it, replacement = RawFirReplacement(place, place))
+                runBodyResolveOnAir(state, collector = it, replacement = RawFirReplacement(place, place))
             }
         }
     }
@@ -152,7 +152,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         )
 
         val collector = FirTowerDataContextAllElementsCollector()
-        val copiedFirDeclaration = runResolveBodyResolveOnAir(
+        val copiedFirDeclaration = runBodyResolveOnAir(
             originalState,
             collector = collector,
             replacement = RawFirReplacement(sameDeclarationInOriginalFile, dependencyNonLocalDeclaration),
@@ -163,7 +163,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         return FirModuleResolveStateDepended(originalState, collector, recordedMap)
     }
 
-    private fun <T : KtElement> runResolveBodyResolveOnAir(
+    private fun <T : KtElement> runBodyResolveOnAir(
         state: FirModuleResolveStateImpl,
         replacement: RawFirReplacement<T>,
         collector: FirTowerDataContextCollector? = null,

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

@@ -173,9 +173,7 @@ internal class FirLazyDeclarationResolver(
         firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
         designation: FirDeclarationDesignationWithFile
     ) {
-        if (designation.toSequence(includeTarget = true).all { it.resolvePhase >= phase }) {
+        if (designation.declaration.resolvePhase >= phase) return
-            return
-        }
 
         val transformer = phase.createLazyTransformer(
             designation,
@@ -220,7 +218,8 @@ internal class FirLazyDeclarationResolver(
         FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
             designation,
             designation.firFile.moduleData.session,
-            scopeSession
+            scopeSession,
+            towerDataContextCollector
         )
         FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
             designation,

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

@@ -9,13 +9,14 @@ 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.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase
 
 internal class FirDesignatedBodyResolveTransformerForIDE(
     private val designation: FirDeclarationDesignationWithFile,
@@ -49,7 +50,13 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
         ideDeclarationTransformer.needReplacePhase
 
     override fun transformDeclaration() {
-        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
+        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)
     }
 }
 

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

@@ -6,12 +6,13 @@
 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.*
-import org.jetbrains.kotlin.fir.declarations.FirFile
 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.api.FirDeclarationDesignationWithFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase
 
 internal class FirDesignatedContractsResolveTransformerForIDE(
     private val designation: FirDeclarationDesignationWithFile,
@@ -20,16 +21,34 @@ internal class FirDesignatedContractsResolveTransformerForIDE(
 ) : 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)
+            }
+        }
+    }
+
+
     @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 transformDeclaration() {
-        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
+        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
+        }
+        designation.ensureTargetPhase(FirResolvePhase.STATUS)
+        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextIndependent)
+        ideDeclarationTransformer.ensureDesignationPassed()
+        designation.ensureTargetPhase(FirResolvePhase.CONTRACTS)
     }
 }

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

@@ -6,21 +6,23 @@
 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.*
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitAwareBodyResolveTransformer
+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.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase
 
 internal class FirDesignatedImplicitTypesTransformerForIDE(
     private val designation: FirDeclarationDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
+    towerDataContextCollector: FirTowerDataContextCollector? = null,
     implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession = ImplicitBodyResolveComputationSession(),
 ) : FirLazyTransformerForIDE, FirImplicitAwareBodyResolveTransformer(
     session,
@@ -28,6 +30,7 @@ internal class FirDesignatedImplicitTypesTransformerForIDE(
     phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
     implicitTypeOnly = true,
     scopeSession = scopeSession,
+    firTowerDataContextCollector = towerDataContextCollector,
     returnTypeCalculator = createReturnTypeCalculatorForIDE(
         session,
         scopeSession,
@@ -46,6 +49,32 @@ internal class FirDesignatedImplicitTypesTransformerForIDE(
     override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase
 
     override fun transformDeclaration() {
-        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
+        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()
+
+        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"
+        }
+
+        designation.ensureTargetPhase(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
     }
 }

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

@@ -14,13 +14,14 @@ 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.api.FirDeclarationDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePathPhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase
 
 class FirDesignatedStatusResolveTransformerForIDE(
     private val designation: FirDeclarationDesignationWithFile,
     private val session: FirSession,
     private val scopeSession: ScopeSession,
-): FirLazyTransformerForIDE {
+) : FirLazyTransformerForIDE {
 
     private val firstItemInDesignation = designation.path.firstOrNull() ?: designation.declaration
 
@@ -52,34 +53,33 @@ class FirDesignatedStatusResolveTransformerForIDE(
         firstItemInDesignation.transformSingle(transformer, null)
     }
 
-    private fun resolveTopLevelMethod(targetCallable: FirCallableDeclaration<*>) {
+    private fun resolveTopLevelDeclaration(declaration: FirDeclaration) {
         val transformer = FirStatusResolveTransformer(
             session = session,
             scopeSession = scopeSession,
             statusComputationSession = StatusComputationSession.Regular()
         )
-        targetCallable.transformSingle(transformer, null)
+        declaration.transformSingle(transformer, null)
     }
 
-    private fun resolveClassMember(containingClass: FirClass<*>, targetCallable: FirCallableDeclaration<*>) {
+    private fun resolveClassMember(containingClass: FirClass<*>, targetCallable: FirDeclaration) {
 
         val transformer = object : FirDesignatedStatusResolveTransformer(
             session = session,
             scopeSession = scopeSession,
             designation = designation.toSequence(includeTarget = true).iterator(),
-            targetClass = containingClass,
+            targetClass = if (targetCallable is FirRegularClass) targetCallable else containingClass,
             statusComputationSession = StatusComputationSession.Regular(),
             designationMapForLocalClasses = emptyMap(),
             scopeForLocalClass = null
         ) {
 
-            override fun <F : FirClass<F>> transformClass(
+            override fun <F : FirClass<F>> transformClass(klass: FirClass<F>, data: FirResolvedDeclarationStatus?): FirStatement {
-                klass: FirClass<F>,
+                if (klass != containingClass) return super.transformClass(klass, data)
-                data: FirResolvedDeclarationStatus?
+                val result = storeClass(klass) {
-            ): FirStatement {
-                return storeClass(klass) {
                     targetCallable.transformSingle(this, data)
-                } as FirStatement
+                }
+                return result as FirStatement
             }
         }
 
@@ -87,21 +87,18 @@ class FirDesignatedStatusResolveTransformerForIDE(
     }
 
     override fun transformDeclaration() {
-        designation.ensurePhase(FirResolvePhase.STATUS, exceptTarget = true)
+        if (designation.declaration.resolvePhase >= FirResolvePhase.STATUS) return
-        when (val resolveTarget = designation.declaration) {
+        designation.ensurePathPhase(FirResolvePhase.TYPES)
-            is FirClass<*> -> resolveClass(resolveTarget)
+        designation.ensureTargetPhase(FirResolvePhase.TYPES)
-            is FirTypeAlias -> resolveTypeAlias(resolveTarget)
+
-            is FirCallableDeclaration<*> -> {
+        val containingClass = designation.path.lastOrNull()
-                val containingClass = designation.path.lastOrNull()
+        if (containingClass == null) {
-                if (containingClass == null) {
+            resolveTopLevelDeclaration(designation.declaration)
-                    check(designation.path.isEmpty()) { "Invalid designation - should be single element designation for top level declaration" }
+        } else {
-                    resolveTopLevelMethod(resolveTarget)
+            check(containingClass is FirClass<*>) { "Invalid designation - the parent is a class" }
-                } else {
+            resolveClassMember(containingClass, designation.declaration)
-                    check(containingClass is FirClass<*>) { "Invalid designation - the parent of callable is not a class" }
-                    resolveClassMember(containingClass, resolveTarget)
-                }
-            }
-            else -> error("Unexpected declaration to resolve ${resolveTarget::class.simpleName}")
         }
+
+        designation.ensureTargetPhase(FirResolvePhase.STATUS)
     }
 }

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

@@ -13,9 +13,10 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirApplySupertypesTransform
 import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSupertypeResolverVisitor
 import org.jetbrains.kotlin.fir.resolve.transformers.SupertypeComputationSession
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePathPhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhaseIfClass
 
 internal class FirDesignatedSupertypeResolverTransformerForIDE(
     private val designation: FirDeclarationDesignationWithFile,
@@ -27,7 +28,8 @@ internal class FirDesignatedSupertypeResolverTransformerForIDE(
     private val supertypeComputationSession = SupertypeComputationSession()
 
     override fun transformDeclaration() {
-        designation.ensurePhase(FirResolvePhase.SUPER_TYPES, exceptTarget = true)
+        if (designation.declaration.resolvePhase >= FirResolvePhase.SUPER_TYPES) return
+        designation.ensurePathPhase(FirResolvePhase.SUPER_TYPES)
 
         val resolver = FirSupertypeResolverVisitor(
             session = session,
@@ -40,5 +42,6 @@ internal class FirDesignatedSupertypeResolverTransformerForIDE(
         designation.declaration.accept(resolver, null)
         val applySupertypesTransformer = FirApplySupertypesTransformer(supertypeComputationSession)
         designation.declaration.transform<FirElement, Void?>(applySupertypesTransformer, null)
+        designation.ensureTargetPhaseIfClass(FirResolvePhase.SUPER_TYPES)
     }
 }

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

@@ -5,13 +5,17 @@
 
 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.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTypeResolveTransformer
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePathPhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensureTargetPhaseIfClass
 
 class FirDesignatedTypeResolverTransformerForIDE(
     private val designation: FirDeclarationDesignationWithFile,
@@ -21,22 +25,21 @@ class FirDesignatedTypeResolverTransformerForIDE(
 
     private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
-    @Suppress("NAME_SHADOWING")
+    override fun <E : FirElement> transformElement(element: E, data: Any?): E {
-    override fun transformRegularClass(regularClass: FirRegularClass, data: Any?): FirStatement {
+        if (element !is FirRegularClass && element !is FirAnonymousObject && element !is FirFile)
-        return ideDeclarationTransformer.transformDeclarationContent(this, regularClass, data) {
+            return super.transformElement(element, data)
-            super.transformRegularClass(regularClass, data) as FirRegularClass
-        }
-    }
 
-    @Suppress("NAME_SHADOWING")
+        return ideDeclarationTransformer.transformDeclarationContent(this, element, data) {
-    override fun transformAnonymousObject(anonymousObject: FirAnonymousObject, data: Any?): FirStatement {
+            super.transformElement(element, data)
-        return ideDeclarationTransformer.transformDeclarationContent(this, anonymousObject, data) {
-            super.transformAnonymousObject(anonymousObject, data) as FirAnonymousObject
         }
     }
 
     override fun transformDeclaration() {
-        designation.ensurePhase(FirResolvePhase.TYPES, exceptTarget = true)
+        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)
     }
 }

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

@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
 
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 
 internal interface FirLazyTransformerForIDE {
@@ -16,13 +16,27 @@ internal interface FirLazyTransformerForIDE {
             override fun transformDeclaration() = Unit
         }
 
-        fun FirDeclarationDesignationWithFile.ensurePhase(firResolvePhase: FirResolvePhase, exceptTarget: Boolean) {
+        fun FirDeclarationDesignationWithFile.ensurePathPhase(firResolvePhase: FirResolvePhase) {
-            val designationIterator = toSequence(includeTarget = !exceptTarget)
+            toSequence(includeTarget = false).forEach { firDeclaration ->
-            designationIterator.forEach { firDeclaration ->
                 check(firDeclaration.resolvePhase >= firResolvePhase) {
                     "Designation element phase required to be $firResolvePhase but element resolved to ${firDeclaration.resolvePhase}"
                 }
             }
         }
+
+        fun FirDeclarationDesignationWithFile.ensureTargetPhase(firResolvePhase: FirResolvePhase) =
+            check(declaration.resolvePhase >= firResolvePhase) { "Expected $firResolvePhase but found ${declaration.resolvePhase}" }
+
+        fun FirDeclarationDesignationWithFile.ensureTargetPhaseIfClass(firResolvePhase: FirResolvePhase) = when (declaration) {
+            is FirProperty, is FirSimpleFunction -> Unit
+            is FirClass<*>, is FirTypeAlias -> ensureTargetPhase(firResolvePhase)
+            else -> error("Unexpected target")
+        }
+
+        fun FirDeclarationDesignationWithFile.ensureTargetPhaseIfMember(firResolvePhase: FirResolvePhase) = when (declaration) {
+            is FirProperty, is FirSimpleFunction -> ensureTargetPhase(firResolvePhase)
+            is FirClass<*>, is FirTypeAlias -> Unit
+            else -> error("Unexpected target")
+        }
     }
 }

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

@@ -26,11 +26,16 @@ internal class IDEDeclarationTransformer(private val designation: FirDeclaration
             return defaultCallTransform()
         }
 
+        if (designationPassed) {
+            return declaration
+        }
+
         if (designationWithoutTargetIterator.hasNext()) {
             designationWithoutTargetIterator.next().visitNoTransform(transformer, data)
         } else {
             try {
                 isInsideTargetDeclaration = true
+                designationPassed = true
                 designation.declaration.visitNoTransform(transformer, data)
             } finally {
                 isInsideTargetDeclaration = false
@@ -41,6 +46,11 @@ internal class IDEDeclarationTransformer(private val designation: FirDeclaration
     }
 
     val needReplacePhase: Boolean get() = isInsideTargetDeclaration
+    private var designationPassed: Boolean = false
+
+    fun ensureDesignationPassed() {
+        check(designationPassed) { "Designation not passed for declaration ${designation.declaration::class.simpleName}" }
+    }
 }
 
 private fun <D> FirElement.visitNoTransform(transformer: FirTransformer<D>, data: D) {