From a172439f2a81a8b2a04e57deb143220cd60e6463 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Wed, 4 Jan 2023 15:51:51 +0100
Subject: [PATCH] [LL FIR] move update declaration resolve phase out of
 updatePhaseForDeclarationInternals

needed for further safe publication of phase update

^KT-56543
---
 .../api/fir/lazy/resolve/LLFirPhaseUpdater.kt | 21 +++++++++++--------
 .../LLFirLazyTransformerExecutor.kt           |  2 ++
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirPhaseUpdater.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirPhaseUpdater.kt
index 21219909007..f93821c5206 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirPhaseUpdater.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirPhaseUpdater.kt
@@ -18,37 +18,40 @@ internal object LLFirPhaseUpdater {
         updateForLocalDeclarations: Boolean,
     ) {
         if (updateForLocalDeclarations) {
-            PhaseUpdatingTransformer.visitElement(target, newPhase)
+            target.acceptChildren(PhaseUpdatingTransformer, newPhase)
         } else {
-            updatePhaseForNonLocals(target, newPhase)
+            updatePhaseForNonLocals(target, newPhase, isTargetDeclaration = true)
         }
     }
 
 
-    private fun updatePhaseForNonLocals(element: FirElementWithResolveState, newPhase: FirResolvePhase) {
+    private fun updatePhaseForNonLocals(element: FirElementWithResolveState, newPhase: FirResolvePhase, isTargetDeclaration: Boolean) {
         if (element.resolvePhase >= newPhase) return
-        element.replaceResolveState(newPhase.asResolveState())
+        if (!isTargetDeclaration) {
+            // phase update for target declaration happens as a declaration publication event after resolve is finished
+            element.replaceResolveState(newPhase.asResolveState())
+        }
 
         if (element is FirTypeParameterRefsOwner) {
             element.typeParameters.forEach { typeParameter ->
                 // if it is not a type parameter of outer declaration
                 if (typeParameter is FirTypeParameter) {
-                    updatePhaseForNonLocals(typeParameter, newPhase)
+                    updatePhaseForNonLocals(typeParameter, newPhase, isTargetDeclaration = false)
                 }
             }
         }
 
         when (element) {
             is FirFunction -> {
-                element.valueParameters.forEach { updatePhaseForNonLocals(it, newPhase) }
+                element.valueParameters.forEach { updatePhaseForNonLocals(it, newPhase, isTargetDeclaration = false) }
             }
             is FirProperty -> {
-                element.getter?.let { updatePhaseForNonLocals(it, newPhase) }
-                element.setter?.let { updatePhaseForNonLocals(it, newPhase) }
+                element.getter?.let { updatePhaseForNonLocals(it, newPhase, isTargetDeclaration = false) }
+                element.setter?.let { updatePhaseForNonLocals(it, newPhase, isTargetDeclaration = false) }
             }
             is FirClass -> {
                 element.declarations.forEach {
-                    updatePhaseForNonLocals(it, newPhase)
+                    updatePhaseForNonLocals(it, newPhase, isTargetDeclaration = false)
                 }
             }
             else -> Unit
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyTransformerExecutor.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyTransformerExecutor.kt
index 7d3d9818a94..9afe9961b0c 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyTransformerExecutor.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyTransformerExecutor.kt
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirPhaseRunner
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.asResolveState
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 
@@ -33,6 +34,7 @@ internal class LLFirLazyTransformerExecutor {
             lockProvider.withLock(designation.firFile) {
                 lazyTransformer.transformDeclaration(phaseRunner)
                 lazyTransformer.updatePhaseForDeclarationInternals(designation.target)
+                designation.target.replaceResolveState(phase.asResolveState())
                 lazyTransformer.checkIsResolved(designation.target)
             }
         }