From 65a7ee501248959e7f6cbb63ddf9b27ad2062417 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Thu, 26 Nov 2020 10:57:33 +0100
Subject: [PATCH] FIR IDE: remove duplicating withFirResolvedToBodyResolve

---
 .../KtFirCompletionCandidateChecker.kt        |  2 +-
 .../api/fir/utils/FirRefWithValidityCheck.kt  | 21 +++++++++----------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt
index 4fe03329dfda8..8767bd0a20b37 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt
@@ -61,7 +61,7 @@ internal class KtFirCompletionCandidateChecker(
 
     private inline fun <reified T : KtFirSymbol<F>, F : FirDeclaration, R> KtCallableSymbol.withResolvedFirOfType(
         noinline action: (F) -> R,
-    ): R? = this.safeAs<T>()?.firRef?.withFirResolvedToBodyResolve(action)
+    ): R? = this.safeAs<T>()?.firRef?.withFir(phase = FirResolvePhase.BODY_RESOLVE, action)
 
     private fun checkExtension(
         candidateSymbol: FirCallableDeclaration<*>,
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/FirRefWithValidityCheck.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/FirRefWithValidityCheck.kt
index bf646649ee76f..dbb13f8cba553 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/FirRefWithValidityCheck.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/FirRefWithValidityCheck.kt
@@ -25,17 +25,16 @@ internal class FirRefWithValidityCheck<D : FirDeclaration>(fir: D, resolveState:
         val resolveState = resolveStateWeakRef.get()
             ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
         fir.resolvedFirToPhase(phase, resolveState)
-        return resolveState.withFirDeclaration(fir) { action(it) }
-    }
-
-    inline fun <R> withFirResolvedToBodyResolve(action: (fir: D) -> R): R {
-        token.assertIsValid()
-        val fir = firWeakRef.get()
-            ?: throw EntityWasGarbageCollectedException("FirElement")
-        val resolveState = resolveStateWeakRef.get()
-            ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
-        fir.resolvedFirToPhase(FirResolvePhase.BODY_RESOLVE, resolveState)
-        return action(resolveState.withFirDeclaration(fir) { it })
+        return when (phase) {
+            FirResolvePhase.BODY_RESOLVE -> {
+                /*
+                 The BODY_RESOLVE phase is the maximum possible phase we can resolve our declaration to
+                 So there is not need to run whole `action` under read lock
+                 */
+                action(resolveState.withFirDeclaration(fir) { it })
+            }
+            else -> resolveState.withFirDeclaration(fir) { action(it) }
+        }
     }
 
     val resolveState