FIR IDE: fix recursive supertypes resolve in light classes

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

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2010-2021 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.api
+
+enum class DeclarationLockType {
+    READ_LOCK,
+    WRITE_LOCK,
+}

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

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
 import org.jetbrains.annotations.TestOnly
 import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
@@ -16,6 +15,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
@@ -41,6 +41,22 @@ abstract class FirModuleResolveState {
 
     internal abstract fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<FirPsiDiagnostic<*>>
 
+    internal inline fun <D : FirDeclaration, R> withLock(declaration: D, declarationLockType: DeclarationLockType, action: (D) -> R): R {
+        val originalDeclaration = (declaration as? FirCallableDeclaration<*>)?.unwrapFakeOverrides() ?: declaration
+        val session = originalDeclaration.session
+        return when {
+            originalDeclaration.origin == FirDeclarationOrigin.Source
+                    && session is FirIdeSourcesSession
+            -> {
+                val cache = session.cache
+                val file = getFirFile(declaration, cache)
+                    ?: error("Fir file was not found for\n${declaration.render()}\n${(declaration.psi as? KtElement)?.getElementTextInContext()}")
+                cache.firFileLockProvider.withLock(file, declarationLockType) { action(declaration) }
+            }
+            else -> action(declaration)
+        }
+    }
+
     @TestOnly
     internal abstract fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile?
 

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

@@ -92,7 +92,7 @@ inline fun <reified F : FirDeclaration, R> KtLambdaExpression.withFirDeclaration
 }
 
 /**
- * Executes [action] with given [FirDeclaration]
+ * Executes [action] with given [FirDeclaration] under read action, so resolve **is not possible** inside [action]
  * [FirDeclaration] passed to [action] will be resolved at least to [phase] when executing [action] on it
  */
 fun <D : FirDeclaration, R> D.withFirDeclaration(
@@ -101,21 +101,17 @@ fun <D : FirDeclaration, R> D.withFirDeclaration(
     action: (D) -> R,
 ): R {
     resolvedFirToPhase(phase, resolveState)
-    val originalDeclaration = (this as? FirCallableDeclaration<*>)?.unwrapFakeOverrides() ?: this
-    val session = originalDeclaration.session
-    return when {
-        originalDeclaration.origin == FirDeclarationOrigin.Source
-                && session is FirIdeSourcesSession
-        -> {
-            val cache = session.cache
-            val file = resolveState.getFirFile(this, cache)
-                ?: error("Fir file was not found for\n${render()}\n${(psi as? KtElement)?.getElementTextInContext()}")
-            cache.firFileLockProvider.withReadLock(file) { action(this) }
-        }
-        else -> action(this)
-    }
+    return resolveState.withLock(this, DeclarationLockType.READ_LOCK, action)
 }
 
+/**
+ * Executes [action] with given [FirDeclaration] under write lock, so resolve **is possible** inside [action]
+ */
+fun <D : FirDeclaration, R> D.withFirDeclarationInWriteLock(
+    resolveState: FirModuleResolveState,
+    action: (D) -> R,
+): R = resolveState.withLock(this, DeclarationLockType.WRITE_LOCK, action)
+
 /**
  * Returns a list of Diagnostics compiler finds for given [KtElement]
  */

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/LockProvider.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.file.builder
 
 import com.google.common.collect.MapMaker
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.PrivateForInline
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.DeclarationLockType
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.lockWithPCECheck
 import java.util.concurrent.ConcurrentMap
 import java.util.concurrent.locks.ReadWriteLock
@@ -39,6 +40,11 @@ internal class LockProvider<KEY> {
         val writeLock = getLockFor(key).writeLock()
         return deadLockGuard.guardWriteLock { writeLock.lockWithPCECheck(lockingIntervalMs, action) }
     }
+
+    inline fun <R> withLock(declaration: KEY, declarationLockType: DeclarationLockType, action: () -> R): R = when (declarationLockType) {
+        DeclarationLockType.READ_LOCK -> withReadLock(declaration, action)
+        DeclarationLockType.WRITE_LOCK -> withWriteLock(declaration, action)
+    }
 }
 
 @PrivateForInline

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirNamedClassOrObjectSymbol.kt

@@ -65,7 +65,7 @@ internal class KtFirNamedClassOrObjectSymbol(
     }
 
     override val superTypes: List<KtTypeAndAnnotations> by cached {
-        firRef.superTypesAndAnnotationsList(builder)
+        firRef.superTypesAndAnnotationsListForRegularClass(builder)
     }
 
     override val typeParameters by firRef.withFirAndCache { fir ->

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirTypeAndAnnotations.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.transformers.resolveSupertypesInTheAir
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
@@ -57,14 +58,29 @@ internal class KtSimpleFirTypeAndAnnotations(
 
 internal fun FirRefWithValidityCheck<FirClass<*>>.superTypesAndAnnotationsList(builder: KtSymbolByFirBuilder): List<KtTypeAndAnnotations> =
     withFir(FirResolvePhase.SUPER_TYPES) { fir ->
-        fir.superTypeRefs.map { typeRef ->
-            val annotations = typeRef.annotations.map { annotation ->
-                KtFirAnnotationCall(this, annotation)
-            }
-            KtSimpleFirTypeAndAnnotations(typeRef.coneType, annotations, builder, token)
-        }
+        fir.superTypeRefs.mapToTypeAndAnnotations(this, builder)
     }
 
+internal fun FirRefWithValidityCheck<FirRegularClass>.superTypesAndAnnotationsListForRegularClass(builder: KtSymbolByFirBuilder): List<KtTypeAndAnnotations> {
+    return withFir { fir ->
+        if (fir.resolvePhase >= FirResolvePhase.SUPER_TYPES) {
+            fir.superTypeRefs.mapToTypeAndAnnotations(this, builder)
+        } else null
+    } ?: withFirWithPossibleResolveInside { fir ->
+        fir.resolveSupertypesInTheAir(builder.rootSession).mapToTypeAndAnnotations(this, builder)
+    }
+}
+
+private fun List<FirTypeRef>.mapToTypeAndAnnotations(
+    containingDeclaration: FirRefWithValidityCheck<FirClass<*>>,
+    builder: KtSymbolByFirBuilder,
+) = map { typeRef ->
+    val annotations = typeRef.annotations.map { annotation ->
+        KtFirAnnotationCall(containingDeclaration, annotation)
+    }
+    KtSimpleFirTypeAndAnnotations(typeRef.coneType, annotations, builder, containingDeclaration.token)
+}
+
 internal fun FirRefWithValidityCheck<FirTypedDeclaration>.returnTypeAndAnnotations(
     typeResolvePhase: FirResolvePhase,
     builder: KtSymbolByFirBuilder

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/FirRefWithValidityCheck.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclaration
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclarationInWriteLock
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.tokens.assertIsValidAndAccessible
@@ -53,6 +54,19 @@ internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveSt
         return action(fir)
     }
 
+    /**
+     * Runs [action] with fir element with write action hold
+     * Consider using this then [action] may call some resolve
+     */
+    inline fun <R> withFirWithPossibleResolveInside(crossinline action: (fir: D) -> R): R {
+        token.assertIsValidAndAccessible()
+        val fir = firWeakRef.get()
+            ?: throw EntityWasGarbageCollectedException("FirElement")
+        val resolveState = resolveStateWeakRef.get()
+            ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
+        return fir.withFirDeclarationInWriteLock(resolveState) { action(it) }
+    }
+
     val resolveState
         get() = resolveStateWeakRef.get() ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")