FIR IDE: do not mark declaration with some lazy resolve phase if only some children are resolved to that phase

Otherwise, we will not be able to resolve parent one

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

@@ -60,6 +60,14 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         context.addLocalScope(localScope)
     }
 
+    protected open fun needReplacePhase(firDeclaration: FirDeclaration) = true
+
+    fun replaceDeclarationResolvePhaseIfNeeded(firDeclaration: FirDeclaration, newResolvePhase: FirResolvePhase) {
+        if (needReplacePhase(firDeclaration)) {
+            firDeclaration.replaceResolvePhase(newResolvePhase)
+        }
+    }
+
     @OptIn(PrivateForInline::class)
     internal inline fun <T> withFullBodyResolve(crossinline l: () -> T): T {
         val shouldSwitchMode = implicitTypeOnly

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

@@ -57,7 +57,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     ): CompositeTransformResult<FirDeclaration> {
         transformer.onBeforeDeclarationContentResolve(declaration)
         return context.withContainer(declaration) {
-            declaration.replaceResolvePhase(transformerPhase)
+            transformer.replaceDeclarationResolvePhaseIfNeeded(declaration, transformerPhase)
             transformer.transformDeclarationContent(declaration, data)
         }
     }
@@ -84,7 +84,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         if (declaration.typeParameters.isEmpty()) return null
 
         for (typeParameter in declaration.typeParameters) {
-            (typeParameter as? FirTypeParameter)?.replaceResolvePhase(FirResolvePhase.STATUS)
+            (typeParameter as? FirTypeParameter)?.let { transformer.replaceDeclarationResolvePhaseIfNeeded(it, FirResolvePhase.STATUS) }
             typeParameter.transformChildren(transformer, ResolutionMode.ContextIndependent)
         }
 
@@ -115,7 +115,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             if (returnTypeRef !is FirImplicitTypeRef && implicitTypeOnly) return@withTypeParametersOf property.compose()
             if (property.resolvePhase == transformerPhase) return@withTypeParametersOf property.compose()
             if (property.resolvePhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE && transformerPhase == FirResolvePhase.BODY_RESOLVE) {
-                property.replaceResolvePhase(transformerPhase)
+                transformer.replaceDeclarationResolvePhaseIfNeeded(property, transformerPhase)
                 return@withTypeParametersOf property.compose()
             }
             dataFlowAnalyzer.enterProperty(property)
@@ -142,7 +142,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                             }
                         }
                     }
-                    property.replaceResolvePhase(transformerPhase)
+                    transformer.replaceDeclarationResolvePhaseIfNeeded(property, transformerPhase)
                     dataFlowAnalyzer.exitProperty(property)?.let {
                         property.replaceControlFlowGraphReference(FirControlFlowGraphReferenceImpl(it))
                     }
@@ -267,7 +267,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             variable.transformAccessors()
         }
         context.storeVariable(variable)
-        variable.replaceResolvePhase(transformerPhase)
+        transformer.replaceDeclarationResolvePhaseIfNeeded(variable, transformerPhase)
         dataFlowAnalyzer.exitLocalVariableDeclaration(variable)
         return variable.compose()
     }
@@ -449,7 +449,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     ): CompositeTransformResult<FirSimpleFunction> {
         if (simpleFunction.resolvePhase == transformerPhase) return simpleFunction.compose()
         if (simpleFunction.resolvePhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE && transformerPhase == FirResolvePhase.BODY_RESOLVE) {
-            simpleFunction.replaceResolvePhase(transformerPhase)
+            transformer.replaceDeclarationResolvePhaseIfNeeded(simpleFunction, transformerPhase)
             return simpleFunction.compose()
         }
         val returnTypeRef = simpleFunction.returnTypeRef
@@ -546,7 +546,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
     private fun doTransformConstructor(constructor: FirConstructor, data: ResolutionMode): CompositeTransformResult<FirConstructor> {
         return context.withContainer(constructor) {
-            constructor.replaceResolvePhase(transformerPhase)
+            transformer.replaceDeclarationResolvePhaseIfNeeded(constructor, transformerPhase)
             dataFlowAnalyzer.enterFunction(constructor)
 
             constructor.transformTypeParameters(transformer, data)
@@ -640,7 +640,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     override fun transformValueParameter(valueParameter: FirValueParameter, data: ResolutionMode): CompositeTransformResult<FirStatement> {
         context.storeVariable(valueParameter)
         if (valueParameter.returnTypeRef is FirImplicitTypeRef) {
-            valueParameter.replaceResolvePhase(transformerPhase)
+            transformer.replaceDeclarationResolvePhaseIfNeeded(valueParameter, transformerPhase)
             return valueParameter.compose()
         }
 

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

@@ -275,7 +275,7 @@ open class FirContractResolveTransformer(
             get() = contractDescription is FirLegacyRawContractDescription || contractDescription is FirRawContractDescription
 
         private fun FirDeclaration.updatePhase() {
-            replaceResolvePhase(FirResolvePhase.CONTRACTS)
+            replaceDeclarationResolvePhaseIfNeeded(this, FirResolvePhase.CONTRACTS)
         }
     }
 }

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

@@ -155,7 +155,13 @@ internal class FirLazyDeclarationResolver(
             return
         }
 
-        val transformer = phase.createLazyTransformer(designation, containerFirFile, scopeSession, towerDataContextCollector)
+        val transformer = phase.createLazyTransformer(
+            designation,
+            firDeclarationToResolve,
+            containerFirFile,
+            scopeSession,
+            towerDataContextCollector
+        )
 
         firFileBuilder.firPhaseRunner.runPhaseWithCustomResolve(phase) {
             containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
@@ -164,22 +170,26 @@ internal class FirLazyDeclarationResolver(
 
     private fun FirResolvePhase.createLazyTransformer(
         designation: List<FirDeclaration>,
+        targetDeclaration: FirDeclaration,
         containerFirFile: FirFile,
         scopeSession: ScopeSession,
         towerDataContextCollector: FirTowerDataContextCollector?
     ) = when (this) {
         FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
             designation.iterator(),
+            targetDeclaration,
             containerFirFile.session,
             scopeSession,
         )
         FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
             designation.iterator(),
+            targetDeclaration,
             containerFirFile.session,
             scopeSession,
         )
         FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
             designation.iterator(),
+            targetDeclaration,
             containerFirFile.session,
             scopeSession,
             towerDataContextCollector

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

@@ -21,7 +21,8 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesigna
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 
 internal class FirDesignatedBodyResolveTransformerForIDE(
-    private val designation: Iterator<FirElement>,
+    private val designation: Iterator<FirDeclaration>,
+    targetDeclaration: FirDeclaration,
     session: FirSession,
     scopeSession: ScopeSession,
     private val towerDataContextCollector: FirTowerDataContextCollector? = null
@@ -37,10 +38,11 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     )
 ) {
+    private val phaseReplaceOracle = PhaseReplaceOracle(targetDeclaration)
 
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
-        if (designation.hasNext()) {
+        if (designation.hasNext()) phaseReplaceOracle.transformDeclarationInside(designation.next()) {
-            designation.next().visitNoTransform(this, data)
+            it.visitNoTransform(this, data)
             return declaration.compose()
         }
 
@@ -54,4 +56,6 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
     override fun onBeforeStatementResolution(statement: FirStatement) {
         towerDataContextCollector?.addStatementContext(statement, context.towerDataContext)
     }
+
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = phaseReplaceOracle.needReplacePhase(firDeclaration)
 }

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

@@ -15,16 +15,20 @@ import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
 
 internal class FirDesignatedContractsResolveTransformerForIDE(
-    private val designation: Iterator<FirElement>,
+    private val designation: Iterator<FirDeclaration>,
+    targetDeclaration: FirDeclaration,
     session: FirSession,
     scopeSession: ScopeSession,
 ) : FirContractResolveTransformer(session, scopeSession) {
+    private val phaseReplaceOracle = PhaseReplaceOracle(targetDeclaration)
 
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
-        if (designation.hasNext()) {
+        if (designation.hasNext()) phaseReplaceOracle.transformDeclarationInside(designation.next()) {
-            designation.next().visitNoTransform(this, data)
+            it.visitNoTransform(this, data)
             return declaration.compose()
         }
         return super.transformDeclarationContent(declaration, data)
     }
+
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = phaseReplaceOracle.needReplacePhase(firDeclaration)
 }

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

@@ -19,7 +19,8 @@ import org.jetbrains.kotlin.fir.visitors.compose
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
 
 internal class FirDesignatedImplicitTypesTransformerForIDE(
-    private val designation: Iterator<FirElement>,
+    private val designation: Iterator<FirDeclaration>,
+    targetDeclaration: FirDeclaration,
     session: FirSession,
     scopeSession: ScopeSession,
     implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession = ImplicitBodyResolveComputationSession(),
@@ -36,12 +37,15 @@ internal class FirDesignatedImplicitTypesTransformerForIDE(
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     )
 ) {
+    private val phaseReplaceOracle = PhaseReplaceOracle(targetDeclaration)
 
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
-        if (designation.hasNext()) {
+        if (designation.hasNext()) phaseReplaceOracle.transformDeclarationInside(designation.next()) {
-            designation.next().visitNoTransform(this, data)
+            it.visitNoTransform(this, data)
             return declaration.compose()
         }
         return super.transformDeclarationContent(declaration, data)
     }
+
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = phaseReplaceOracle.needReplacePhase(firDeclaration)
 }

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

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.trasformers
+
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+
+internal class PhaseReplaceOracle(private val targetDeclaration: FirDeclaration) {
+    private var isInsideCurrentDeclaration = false
+
+    fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
+        isInsideCurrentDeclaration
+
+    inline fun <T> transformDeclarationInside(declaration: FirDeclaration, transform: (FirDeclaration) -> T): T {
+        if (declaration == targetDeclaration) isInsideCurrentDeclaration = true
+        return try {
+            transform(declaration)
+        } finally {
+            if (declaration == targetDeclaration) isInsideCurrentDeclaration = false
+        }
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/classMembers.txt

@@ -1,5 +1,5 @@
 FILE: classMembers.kt
-    public final [BODY_RESOLVE] class A : R|kotlin/Any| {
+    public final [CONTRACTS] class A : R|kotlin/Any| {
         public [CONTRACTS] constructor(): R|A| {
             super<R|kotlin/Any|>()
         }

idea/testData/fir/lazyResolve/elvis/main.txt

@@ -1,5 +1,5 @@
 FILE: main.kt
-    public final [BODY_RESOLVE] class WithElvis : R|kotlin/Any| {
+    public final [CONTRACTS] class WithElvis : R|kotlin/Any| {
         public [CONTRACTS] constructor(value: R|kotlin/String?|): R|elvis/WithElvis| {
             super<R|kotlin/Any|>()
         }

idea/testData/fir/lazyResolve/inSecondary/main.txt

@@ -1,5 +1,5 @@
 FILE: main.kt
-    public final [BODY_RESOLVE] class Foo : R|kotlin/Any| {
+    public final [CONTRACTS] class Foo : R|kotlin/Any| {
         public final [CONTRACTS] val x: R|kotlin/Int|
             public get(): R|kotlin/Int|