From 4e67afcabae1097f207dff1af3b2f3297ebd61c2 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Tue, 20 Jul 2021 16:41:33 +0300
Subject: [PATCH] FIR IDE: Use write lock in `KtFirOverrideInfoProvider`
 implementation

`withFir` takes read lock, but checking visibility or implementation
status actually requires resolving declaration to `STATUS` phase. But
resolve cannot be done in `withFir` function, since resolve requires
write lock, and you cannot take both read and write lock

Now we use `withFirWithPossibleResolveInside`, which takes write lock
and allows to perform resolve in it

This should fix failing `FirOverrideImplementTest` tests
in Kotlin IDE plugin
---
 .../fir/symbols/KtFirOverrideInfoProvider.kt  | 32 +++++++++++++++----
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirOverrideInfoProvider.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirOverrideInfoProvider.kt
index a9e880693fd..1d97d0df4bd 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirOverrideInfoProvider.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirOverrideInfoProvider.kt
@@ -10,6 +10,8 @@ import org.jetbrains.kotlin.fir.analysis.checkers.isVisibleInClass
 import org.jetbrains.kotlin.fir.containingClass
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.originalForIntersectionOverrideAttr
 import org.jetbrains.kotlin.fir.originalForSubstitutionOverride
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -19,6 +21,7 @@ import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.components.KtOverrideInfoProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRefWithValidityCheck
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
@@ -34,10 +37,14 @@ class KtFirOverrideInfoProvider(
     override fun isVisible(memberSymbol: KtCallableSymbol, classSymbol: KtClassOrObjectSymbol): Boolean {
         require(memberSymbol is KtFirSymbol<*>)
         require(classSymbol is KtFirSymbol<*>)
-        return memberSymbol.firRef.withFir { memberFir ->
-            if (memberFir !is FirCallableDeclaration) return@withFir false
-            classSymbol.firRef.withFir inner@{ parentClassFir ->
+
+        // Inspecting visibility requires resolving to status
+        return memberSymbol.firRef.withFirWithResolveAllowed outer@{ memberFir ->
+            if (memberFir !is FirCallableDeclaration) return@outer false
+
+            classSymbol.firRef.withFirWithResolveAllowed inner@{ parentClassFir ->
                 if (parentClassFir !is FirClass) return@inner false
+
                 memberFir.isVisibleInClass(parentClassFir)
             }
         }
@@ -46,10 +53,14 @@ class KtFirOverrideInfoProvider(
     override fun getImplementationStatus(memberSymbol: KtCallableSymbol, parentClassSymbol: KtClassOrObjectSymbol): ImplementationStatus? {
         require(memberSymbol is KtFirSymbol<*>)
         require(parentClassSymbol is KtFirSymbol<*>)
-        return memberSymbol.firRef.withFir { memberFir ->
-            if (memberFir !is FirCallableDeclaration) return@withFir null
-            parentClassSymbol.firRef.withFir inner@{ parentClassFir ->
+
+        // Inspecting implementation status requires resolving to status
+        return memberSymbol.firRef.withFirWithResolveAllowed outer@{ memberFir ->
+            if (memberFir !is FirCallableDeclaration) return@outer null
+
+            parentClassSymbol.firRef.withFirWithResolveAllowed inner@{ parentClassFir ->
                 if (parentClassFir !is FirClass) return@inner null
+
                 memberFir.symbol.getImplementationStatus(
                     SessionHolderImpl(firAnalysisSession.rootModuleSession, ScopeSession()),
                     parentClassFir.symbol
@@ -92,3 +103,12 @@ class KtFirOverrideInfoProvider(
         return member
     }
 }
+
+/**
+ * Convenience function.
+ *
+ * We use [FirRefWithValidityCheck.withFirWithPossibleResolveInside] instead of [FirRefWithValidityCheck.withFir] because we want to be able
+ * to call this function inside itself - and it requires write lock, since each layer might want to call resolve and to use write lock.
+ */
+private inline fun <D : FirDeclaration, R> FirRefWithValidityCheck<D>.withFirWithResolveAllowed(crossinline action: (D) -> R): R =
+    withFirWithPossibleResolveInside(phase = FirResolvePhase.RAW_FIR, action)