FIR IDE: resolve kt symbols only to the phase they actually need

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/highlighter/KotlinFirPsiChecker.kt

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.fir.highlighter.visitors.FirAfterResolveHighlightingVisitor
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.analyze
-import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.highlighter.AbstractKotlinPsiChecker
 import org.jetbrains.kotlin.idea.highlighter.Diagnostic2Annotation
 import org.jetbrains.kotlin.idea.highlighter.IdeErrorMessages

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

@@ -50,5 +50,5 @@ object DebugSymbolRenderer {
         else -> value::class.simpleName!!
     }
 
-    private val ignoredPropertyNames = setOf("fir", "psi", "token")
+    private val ignoredPropertyNames = setOf("firRef", "psi", "token")
 }

idea/idea-frontend-fir/idea-fir-low-level-api/build.gradle.kts

@@ -9,6 +9,8 @@ dependencies {
     compile(project(":idea:idea-frontend-api"))
     compile(project(":idea:idea-core"))
     compile(project(":compiler:fir:fir2ir"))
+    compile(project(":compiler:fir:fir2ir:jvm-backend"))
+    compile(project(":compiler:ir.serialization.common"))
     compile(project(":compiler:fir:resolve"))
     compile(project(":compiler:fir:checkers"))
     compile(project(":compiler:fir:java"))

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

@@ -41,6 +41,8 @@ abstract class FirModuleResolveState {
     // todo temporary, used only in completion
     abstract fun getCachedMappingForCompletion(element: KtElement): FirElement?
 
+    abstract fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D
+
     // todo temporary, used only in completion
     internal abstract fun lazyResolveFunctionForCompletion(
         firFunction: FirFunction<*>,
@@ -77,6 +79,11 @@ internal open class FirModuleResolveStateImpl(
         psiToFirCache.recordElementsForCompletionFrom(fir, firFile, ktFile)
     }
 
+    override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
+        elementBuilder.lazyResolveDeclaration(declaration, fileCache, toPhase)
+        return declaration
+    }
+
     override fun getCachedMappingForCompletion(element: KtElement): FirElement? =
         psiToFirCache.getCachedMapping(element)
 

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

@@ -51,6 +51,10 @@ internal class FirModuleResolveStateForCompletion(
         return null
     }
 
+    override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
+        return originalState.resolvedFirToPhase(declaration, toPhase)
+    }
+
     override fun lazyResolveFunctionForCompletion(
         firFunction: FirFunction<*>,
         containerFirFile: FirFile,

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

@@ -0,0 +1,20 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.backend.jvm.FirJvmKotlinMangler
+import org.jetbrains.kotlin.fir.signaturer.FirBasedSignatureComposer
+
+data class IdeSessionComponents(val signatureComposer: FirBasedSignatureComposer): FirSessionComponent {
+    companion object {
+        fun create(session: FirSession) = IdeSessionComponents(
+            signatureComposer = FirBasedSignatureComposer(FirJvmKotlinMangler(session))
+        )
+    }
+}
+val FirSession.ideSessionComponents: IdeSessionComponents by FirSession.sessionComponentAccessor()

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt

@@ -12,6 +12,8 @@ import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.firProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
@@ -21,6 +23,7 @@ import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
 
 /**
  * Maps [KtElement] to [FirElement]
@@ -55,13 +58,26 @@ internal class FirElementBuilder(
         return psiToFirCache.getFir(element, containerFir, firFile)
     }
 
+    fun lazyResolveDeclaration(
+        declaration: FirDeclaration,
+        moduleFileCache: ModuleFileCache,
+        toPhase: FirResolvePhase
+    ) {
+        if (declaration.resolvePhase < toPhase) {
+            val firFile = declaration.getContainingFile() ?: error("FirFile was not found for\n${declaration.render()}")
+            firFileBuilder.runCustomResolve(firFile, moduleFileCache) {
+                runLazyResolveWithoutLock(declaration, firFile, declaration.session.firProvider, toPhase)
+            }
+        }
+    }
+
     /**
      * Should be invoked under lock
      */
     private fun runLazyResolveWithoutLock(
         firDeclarationToResolve: FirDeclaration,
         containerFirFile: FirFile,
-        firIdeProvider: FirIdeProvider,
+        provider: FirProvider,
         toPhase: FirResolvePhase,
         towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
     ) {
@@ -72,6 +88,7 @@ internal class FirElementBuilder(
         if (toPhase <= nonLazyPhase) return
         val designation = mutableListOf<FirDeclaration>(containerFirFile)
         if (firDeclarationToResolve !is FirFile) {
+
             val id = when (firDeclarationToResolve) {
                 is FirCallableDeclaration<*> -> {
                     firDeclarationToResolve.symbol.callableId.classId
@@ -83,7 +100,7 @@ internal class FirElementBuilder(
             }
             val outerClasses = generateSequence(id) { classId ->
                 classId.outerClassId
-            }.mapTo(mutableListOf()) { firIdeProvider.getFirClassifierByFqName(it)!! }
+            }.mapTo(mutableListOf()) { provider.getFirClassifierByFqName(it)!! }
             designation += outerClasses.asReversed()
             if (firDeclarationToResolve is FirCallableDeclaration<*>) {
                 designation += firDeclarationToResolve
@@ -115,7 +132,6 @@ internal class FirElementBuilder(
 }
 
 
-
 private fun KtElement.getNonLocalContainingDeclarationWithFqName(): KtDeclaration? {
     var container = parent
     while (container != null && container !is KtFile) {

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

@@ -7,7 +7,9 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.file.builder
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.name.ClassId
@@ -44,6 +46,8 @@ internal abstract class ModuleFileCache {
      */
     abstract fun fileCached(file: KtFile, createValue: () -> FirFile): FirFile
 
+    abstract fun getContainerFirFile(declaration: FirDeclaration): FirFile?
+
     abstract fun getCachedFirFile(ktFile: KtFile): FirFile?
 
     // todo make it ReadWriteLock and allow access fir elements only under read lock
@@ -65,6 +69,10 @@ internal class ModuleFileCacheImpl(override val session: FirSession) : ModuleFil
 
     override fun getCachedFirFile(ktFile: KtFile): FirFile? = ktFileToFirFile[ktFile]
 
+    override fun getContainerFirFile(declaration: FirDeclaration): FirFile? {
+        val ktFile = declaration.psi?.containingFile as? KtFile ?: return null
+        return getCachedFirFile(ktFile)
+    }
 
     override val firFileLockProvider: LockProvider<FirFile, ReentrantLock> = LockProvider { ReentrantLock() }
 }

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

@@ -76,7 +76,7 @@ internal class FirIdeProvider(
 
     override fun getFirClassifierContainerFileIfAny(fqName: ClassId): FirFile? {
         val fir = getFirClassifierByFqName(fqName) ?: return null // Necessary to ensure cacheProvider contains this classifier
-        return providerHelper.getContainingFirFile(fir.symbol)
+        return cache.getContainerFirFile(fir)
     }
 
     override fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile {
@@ -85,7 +85,7 @@ internal class FirIdeProvider(
     }
 
     override fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? =
-        providerHelper.getContainingFirFile(symbol)
+        cache.getContainerFirFile(symbol.fir)
 
 
     override fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile? {
@@ -98,7 +98,7 @@ internal class FirIdeProvider(
                 return getFirCallableContainerFile(fir.getter.delegate.symbol)
             }
         }
-        return providerHelper.getContainingFirFile(symbol)
+        return cache.getContainerFirFile(symbol.fir)
     }
 
     override fun getFirFilesByPackage(fqName: FqName): List<FirFile> = error("Should not be called in FIR IDE")

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

@@ -72,11 +72,6 @@ internal class FirProviderHelper(
         }
     }
 
-    fun getContainingFirFile(symbol: FirBasedSymbol<*>): FirFile? {
-        val ktFile = symbol.fir.psi?.containingFile as? KtFile ?: return null
-        return cache.getCachedFirFile(ktFile)
-    }
-
     private fun FirFile.collectCallableDeclarationsTo(list: MutableList<FirCallableSymbol<*>>, name: Name) {
         declarations.mapNotNullTo(list) { declaration ->
             if (declaration is FirCallableDeclaration<*> && declaration.symbol.callableId.callableName == name) {

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

@@ -0,0 +1,21 @@
+/*
+ * 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.util
+
+import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.firProvider
+
+fun FirDeclaration.getContainingFile(): FirFile? {
+    val provider = session.firProvider
+    return when (this) {
+        is FirCallableDeclaration<*> -> provider.getFirCallableContainerFile(symbol)
+        is FirClassLikeDeclaration<*> -> provider.getFirClassifierContainerFile(symbol)
+        else -> error("Unsupported declaration ${this::class.java}")
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractFirLazyResolveTest.kt

@@ -102,7 +102,7 @@ abstract class AbstractFirLazyResolveTest : KotlinLightCodeInsightFixtureTestCas
         val files = FileTypeIndex.getFiles(KotlinFileType.INSTANCE, contentScope)
         for (file in files) {
             val psiFile = psiManager.findFile(file) as KtFile
-            val session = resolveState.firSession
+            val session = resolveState.firIdeSourcesSession
             val firProvider = session.firIdeProvider
             val firFile = firProvider.cache.getCachedFirFile(psiFile) ?: continue
             KotlinTestUtils.assertEqualsToFile(File(expectedTxtPath(file)), firFile.render())

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

@@ -39,7 +39,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 
-class KtFirAnalysisSession
+internal class KtFirAnalysisSession
 private constructor(
     private val element: KtElement,
     val firSession: FirSession,
@@ -132,7 +132,7 @@ private constructor(
             element,
             firSession,
             contextResolveState,
-            firSymbolBuilder.createReadOnlyCopy(),
+            firSymbolBuilder.createReadOnlyCopy(contextResolveState),
             typeContext,
             token,
             isContextSession = true
@@ -211,6 +211,7 @@ private constructor(
             val firSymbolBuilder = KtSymbolByFirBuilder(
                 firSession.firSymbolProvider,
                 typeContext,
+                firResolveState,
                 project,
                 token
             )

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.*
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.*
@@ -45,6 +46,7 @@ internal class KtSymbolByFirBuilder private constructor(
     firProvider: FirSymbolProvider,
     typeCheckerContext: ConeTypeCheckerContext,
     private val project: Project,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     val withReadOnlyCaching: Boolean,
     private val symbolsCache: BuilderCache<FirDeclaration, KtSymbol>,
@@ -54,6 +56,7 @@ internal class KtSymbolByFirBuilder private constructor(
     constructor(
         firProvider: FirSymbolProvider,
         typeCheckerContext: ConeTypeCheckerContext,
+        resolveState: FirModuleResolveState,
         project: Project,
         token: ValidityToken
     ) : this(
@@ -61,24 +64,27 @@ internal class KtSymbolByFirBuilder private constructor(
         typeCheckerContext = typeCheckerContext,
         project = project,
         token = token,
+        resolveState = resolveState,
         withReadOnlyCaching = false,
         symbolsCache = BuilderCache(),
         typesCache = BuilderCache()
     )
 
     private val firProvider by weakRef(firProvider)
+    private val resolveState by weakRef(resolveState)
     private val typeCheckerContext by weakRef(typeCheckerContext)
 
-    fun createReadOnlyCopy(): KtSymbolByFirBuilder {
+    fun createReadOnlyCopy(newResolveState: FirModuleResolveState): KtSymbolByFirBuilder {
         check(!withReadOnlyCaching) { "Cannot create readOnly KtSymbolByFirBuilder from a readonly one" }
         return KtSymbolByFirBuilder(
             firProvider,
             typeCheckerContext,
             project,
-            token,
+            token = token,
+            resolveState = newResolveState,
             withReadOnlyCaching = true,
-            symbolsCache.createReadOnlyCopy(),
-            typesCache.createReadOnlyCopy()
+            symbolsCache = symbolsCache.createReadOnlyCopy(),
+            typesCache = typesCache.createReadOnlyCopy()
         )
     }
 
@@ -110,30 +116,33 @@ internal class KtSymbolByFirBuilder private constructor(
             TODO(fir::class.toString())
     }
 
-    fun buildClassSymbol(fir: FirRegularClass) = symbolsCache.cache(fir) { KtFirClassOrObjectSymbol(fir, token, this) }
+    fun buildClassSymbol(fir: FirRegularClass) = symbolsCache.cache(fir) { KtFirClassOrObjectSymbol(fir, resolveState, token, this) }
 
     // TODO it can be a constructor parameter, which may be split into parameter & property
     // we should handle them both
-    fun buildParameterSymbol(fir: FirValueParameterImpl) = symbolsCache.cache(fir) { KtFirFunctionValueParameterSymbol(fir, token, this) }
+    fun buildParameterSymbol(fir: FirValueParameterImpl) =
+        symbolsCache.cache(fir) { KtFirFunctionValueParameterSymbol(fir, resolveState, token, this) }
+
     fun buildFirConstructorParameter(fir: FirValueParameterImpl) =
-        symbolsCache.cache(fir) { KtFirConstructorValueParameterSymbol(fir, token, this) }
+        symbolsCache.cache(fir) { KtFirConstructorValueParameterSymbol(fir, resolveState, token, this) }
 
     fun buildFunctionSymbol(fir: FirSimpleFunction, forcedOrigin: FirDeclarationOrigin? = null) = symbolsCache.cache(fir) {
-        KtFirFunctionSymbol(fir, token, this, forcedOrigin)
+        KtFirFunctionSymbol(fir, resolveState, token, this, forcedOrigin)
     }
 
-    fun buildConstructorSymbol(fir: FirConstructor) = symbolsCache.cache(fir) { KtFirConstructorSymbol(fir, token, this) }
-    fun buildTypeParameterSymbol(fir: FirTypeParameter) = symbolsCache.cache(fir) { KtFirTypeParameterSymbol(fir, token) }
+    fun buildConstructorSymbol(fir: FirConstructor) = symbolsCache.cache(fir) { KtFirConstructorSymbol(fir, resolveState, token, this) }
+    fun buildTypeParameterSymbol(fir: FirTypeParameter) = symbolsCache.cache(fir) { KtFirTypeParameterSymbol(fir, resolveState, token) }
 
-    fun buildTypeAliasSymbol(fir: FirTypeAlias) = symbolsCache.cache(fir) { KtFirTypeAliasSymbol(fir, token) }
-    fun buildEnumEntrySymbol(fir: FirEnumEntry) = symbolsCache.cache(fir) { KtFirEnumEntrySymbol(fir, token, this) }
-    fun buildFieldSymbol(fir: FirField) = symbolsCache.cache(fir) { KtFirJavaFieldSymbol(fir, token, this) }
-    fun buildAnonymousFunctionSymbol(fir: FirAnonymousFunction) = symbolsCache.cache(fir) { KtFirAnonymousFunctionSymbol(fir, token, this) }
+    fun buildTypeAliasSymbol(fir: FirTypeAlias) = symbolsCache.cache(fir) { KtFirTypeAliasSymbol(fir, resolveState, token) }
+    fun buildEnumEntrySymbol(fir: FirEnumEntry) = symbolsCache.cache(fir) { KtFirEnumEntrySymbol(fir, resolveState, token, this) }
+    fun buildFieldSymbol(fir: FirField) = symbolsCache.cache(fir) { KtFirJavaFieldSymbol(fir, resolveState, token, this) }
+    fun buildAnonymousFunctionSymbol(fir: FirAnonymousFunction) =
+        symbolsCache.cache(fir) { KtFirAnonymousFunctionSymbol(fir, resolveState, token, this) }
 
     fun buildVariableSymbol(fir: FirProperty, forcedOrigin: FirDeclarationOrigin? = null): KtVariableSymbol = symbolsCache.cache(fir) {
         when {
-            fir.isLocal -> KtFirLocalVariableSymbol(fir, token, this)
-            else -> KtFirPropertySymbol(fir, token, this, forcedOrigin)
+            fir.isLocal -> KtFirLocalVariableSymbol(fir, resolveState, token, this)
+            else -> KtFirPropertySymbol(fir, resolveState, token, this, forcedOrigin)
         }
     }
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirScopeProvider.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.scopes
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.scope
 import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
@@ -58,7 +59,10 @@ internal class KtFirScopeProvider(
     override fun getMemberScope(classSymbol: KtClassOrObjectSymbol): KtMemberScope = withValidityAssertion {
         memberScopeCache.getOrPut(classSymbol) {
             check(classSymbol is KtFirClassOrObjectSymbol)
-            val firScope = classSymbol.fir.unsubstitutedScope(classSymbol.fir.session, ScopeSession()).also(firScopeStorage::register)
+            val firScope =
+                classSymbol.firRef.withFir(FirResolvePhase.SUPER_TYPES) { fir ->
+                    fir.unsubstitutedScope(fir.session, ScopeSession())
+                }.also(firScopeStorage::register)
             KtFirMemberScope(classSymbol, firScope, token, builder)
         }
     }
@@ -66,7 +70,8 @@ internal class KtFirScopeProvider(
     override fun getDeclaredMemberScope(classSymbol: KtClassOrObjectSymbol): KtDeclaredMemberScope = withValidityAssertion {
         declaredMemberScopeCache.getOrPut(classSymbol) {
             check(classSymbol is KtFirClassOrObjectSymbol)
-            val firScope = declaredMemberScope(classSymbol.fir).also(firScopeStorage::register)
+            val firScope = classSymbol.firRef.withFir(FirResolvePhase.SUPER_TYPES) { declaredMemberScope(it) }
+                .also(firScopeStorage::register)
             KtFirDeclaredMemberScope(classSymbol, firScope, token, builder)
         }
     }

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

@@ -7,26 +7,31 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 
 internal class KtFirAnonymousFunctionSymbol(
     fir: FirAnonymousFunction,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder
 ) : KtAnonymousFunctionSymbol(), KtFirSymbol<FirAnonymousFunction> {
-    override val fir: FirAnonymousFunction by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
 
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
-    override val valueParameters: List<KtParameterSymbol> by cached {
+    override val type: KtType by firRef.withFirAndCache(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) { builder.buildKtType(it.returnTypeRef) }
+    override val valueParameters: List<KtParameterSymbol> by firRef.withFirAndCache {
         fir.valueParameters.map { valueParameter ->
             check(valueParameter is FirValueParameterImpl)
             builder.buildParameterSymbol(valueParameter)

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

@@ -9,38 +9,37 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
-import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.ReadOnlyWeakRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
-import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 internal class KtFirClassOrObjectSymbol(
     fir: FirRegularClass,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder
 ) : KtClassOrObjectSymbol(), KtFirSymbol<FirRegularClass> {
-    override val fir: FirRegularClass by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
-    override val name: Name get() = withValidityAssertion { fir.classId.shortClassName }
-    override val classId: ClassId get() = withValidityAssertion { fir.classId }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
+    override val name: Name get() = firRef.withFir { it.symbol.classId.shortClassName }
+    override val classId: ClassId get() = firRef.withFir { it.symbol.classId }
 
     override val modality: KtSymbolModality
-        get() = withValidityAssertion { fir.modality.getSymbolModality() }
+        get() = firRef.withFir { it.modality.getSymbolModality() }
 
-    override val typeParameters by cached {
+    override val typeParameters by firRef.withFirAndCache {
         fir.typeParameters.map { typeParameter ->
             builder.buildTypeParameterSymbol(typeParameter.symbol.fir)
         }
     }
 
     override val classKind: KtClassKind
-        get() = withValidityAssertion {
+        get() = firRef.withFir { fir ->
             when (fir.classKind) {
                 ClassKind.INTERFACE -> KtClassKind.INTERFACE
                 ClassKind.ENUM_CLASS -> KtClassKind.ENUM_CLASS
@@ -51,7 +50,7 @@ internal class KtFirClassOrObjectSymbol(
             }
         }
     override val symbolKind: KtSymbolKind
-        get() = withValidityAssertion {
+        get() = firRef.withFir { fir ->
             when {
                 fir.isLocal -> KtSymbolKind.LOCAL
                 fir.symbol.classId.isNestedClass -> KtSymbolKind.MEMBER

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

@@ -8,15 +8,18 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.getPrimaryConstructorIfAny
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
@@ -24,29 +27,30 @@ import org.jetbrains.kotlin.name.ClassId
 
 internal class KtFirConstructorSymbol(
     fir: FirConstructor,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder
 ) : KtConstructorSymbol(), KtFirSymbol<FirConstructor> {
-    override val fir: FirConstructor by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
 
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
-    override val valueParameters: List<KtConstructorParameterSymbol> by cached {
+    override val type: KtType by firRef.withFirAndCache(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) { builder.buildKtType(it.returnTypeRef) }
+    override val valueParameters: List<KtConstructorParameterSymbol> by firRef.withFirAndCache { fir ->
         fir.valueParameters.map { valueParameter ->
             check(valueParameter is FirValueParameterImpl)
             builder.buildFirConstructorParameter(valueParameter)
         }
     }
 
-    override val isPrimary: Boolean get() = withValidityAssertion { fir.isPrimary }
+    override val isPrimary: Boolean get() = firRef.withFir { it.isPrimary }
     override val symbolKind: KtSymbolKind get() = KtSymbolKind.MEMBER
 
     override val ownerClassId: ClassId
-        get() = withValidityAssertion {
-            fir.symbol.callableId.classId ?: error("ClassID should present for constructor")
+        get() = firRef.withFir {
+            it.symbol.callableId.classId ?: error("ClassID should present for constructor")
         }
 
-    override val owner: KtClassOrObjectSymbol by cached {
+    override val owner: KtClassOrObjectSymbol by firRef.withFirAndCache { fir ->
         val session = fir.session
         val classId = ownerClassId
         val firClass = session.firSymbolProvider.getClassLikeSymbolByFqName(classId)?.fir
@@ -64,7 +68,8 @@ internal class KtFirConstructorSymbol(
         return symbolPointer { session ->
             val ownerSymbol = session.symbolProvider.getClassOrObjectSymbolByClassId(ownerClassId) ?: return@symbolPointer null
             check(ownerSymbol is KtFirSymbol<*>)
-            val classFir = (ownerSymbol.fir as? FirRegularClass) ?: error("FirRegularClass expected but ${ownerSymbol.fir::class} found")
+            val classFir = ownerSymbol.firRef.withFir { (it as? FirRegularClass) }
+                ?: error("FirRegularClass expected")
             classFir.getPrimaryConstructorIfAny()?.let(builder::buildConstructorSymbol)
         }
     }

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

@@ -8,9 +8,11 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtConstructorParameterSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtConstructorParameterSymbolKind
@@ -21,23 +23,24 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirConstructorValueParameterSymbol(
     fir: FirValueParameterImpl,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder
 ) : KtConstructorParameterSymbol(), KtFirSymbol<FirValueParameterImpl> {
-    override val fir: FirValueParameterImpl by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
 
-    override val name: Name get() = withValidityAssertion { fir.name }
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
+    override val name: Name get() = firRef.withFir { it.name }
+    override val type: KtType by firRef.withFirAndCache { builder.buildKtType(it.returnTypeRef) }
     override val symbolKind: KtSymbolKind
-        get() = withValidityAssertion {
+        get() = firRef.withFir { fir ->
             when {
                 fir.isVal || fir.isVal -> KtSymbolKind.MEMBER
                 else -> KtSymbolKind.LOCAL
             }
         }
     override val constructorParameterKind: KtConstructorParameterSymbolKind
-        get() = withValidityAssertion {
+        get() = firRef.withFir { fir ->
             when {
                 fir.isVal -> KtConstructorParameterSymbolKind.VAL_PROPERTY
                 fir.isVar -> KtConstructorParameterSymbolKind.VAR_PROPERTY

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

@@ -7,12 +7,15 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirEnumEntry
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtEnumEntrySymbol
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
@@ -20,12 +23,13 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirEnumEntrySymbol(
     fir: FirEnumEntry,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder
 ) : KtEnumEntrySymbol(), KtFirSymbol<FirEnumEntry> {
-    override val fir: FirEnumEntry by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
 
-    override val name: Name get() = withValidityAssertion { fir.name }
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
+    override val name: Name get() = firRef.withFir { it.name }
+    override val type: KtType by firRef.withFirAndCache(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) { fir -> builder.buildKtType(fir.returnTypeRef) }
 }

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

@@ -9,11 +9,13 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCommonSymbolModality
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionSymbol
@@ -25,38 +27,40 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirFunctionSymbol(
     fir: FirSimpleFunction,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder,
     private val forcedOrigin: FirDeclarationOrigin? = null
 ) : KtFunctionSymbol(), KtFirSymbol<FirSimpleFunction> {
-    override val fir: FirSimpleFunction by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
-    override val name: Name get() = withValidityAssertion { fir.name }
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
-    override val valueParameters: List<KtFirFunctionValueParameterSymbol> by cached {
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
+    override val name: Name get() = firRef.withFir { it.name }
+    override val type: KtType by firRef.withFirAndCache(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) { fir -> builder.buildKtType(fir.returnTypeRef) }
+    override val valueParameters: List<KtFirFunctionValueParameterSymbol> by firRef.withFirAndCache { fir ->
         fir.valueParameters.map { valueParameter ->
             check(valueParameter is FirValueParameterImpl)
             builder.buildParameterSymbol(valueParameter)
         }
     }
-    override val typeParameters by cached {
+    override val typeParameters by firRef.withFirAndCache { fir ->
         fir.typeParameters.map { typeParameter ->
             builder.buildTypeParameterSymbol(typeParameter.symbol.fir)
         }
     }
+
     override val origin: KtSymbolOrigin get() = withValidityAssertion { forcedOrigin?.asKtSymbolOrigin() ?: super.origin }
-    override val isSuspend: Boolean get() = withValidityAssertion { fir.isSuspend }
-    override val receiverType: KtType? by cached { fir.receiverTypeRef?.let(builder::buildKtType) }
-    override val isOperator: Boolean get() = withValidityAssertion { fir.isOperator }
-    override val isExtension: Boolean get() = withValidityAssertion { fir.receiverTypeRef != null }
-    override val fqName: FqName? get() = withValidityAssertion { fir.symbol.callableId.asFqNameForDebugInfo() }
+    override val isSuspend: Boolean get() = firRef.withFir { it.isSuspend }
+    override val receiverType: KtType? by firRef.withFirAndCache(FirResolvePhase.TYPES) { fir -> fir.receiverTypeRef?.let(builder::buildKtType) }
+    override val isOperator: Boolean get() = firRef.withFir { it.isOperator }
+    override val isExtension: Boolean get() = firRef.withFir { it.receiverTypeRef != null }
+    override val fqName: FqName? get() = firRef.withFir { it.symbol.callableId.asFqNameForDebugInfo() }
     override val symbolKind: KtSymbolKind
-        get() = withValidityAssertion {
+        get() = firRef.withFir { fir ->
             when {
                 fir.isLocal -> KtSymbolKind.LOCAL
                 fir.symbol.callableId.classId == null -> KtSymbolKind.TOP_LEVEL
                 else -> KtSymbolKind.MEMBER
             }
         }
-    override val modality: KtCommonSymbolModality get() = withValidityAssertion { fir.modality.getSymbolModality() }
+    override val modality: KtCommonSymbolModality get() = firRef.withFir { it.modality.getSymbolModality() }
 }

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

@@ -6,11 +6,14 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionParameterSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolKind
@@ -20,15 +23,16 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirFunctionValueParameterSymbol(
     fir: FirValueParameterImpl,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder
 ) : KtFunctionParameterSymbol(), KtFirSymbol<FirValueParameterImpl> {
-    override val fir: FirValueParameterImpl by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
 
-    override val name: Name get() = withValidityAssertion { fir.name }
-    override val isVararg: Boolean get() = withValidityAssertion { fir.isVararg }
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
+    override val name: Name get() = firRef.withFir { it.name }
+    override val isVararg: Boolean get() = firRef.withFir { it.isVararg }
+    override val type: KtType by firRef.withFirAndCache(FirResolvePhase.TYPES) { fir -> builder.buildKtType(fir.returnTypeRef) }
     override val symbolKind: KtSymbolKind get() = KtSymbolKind.LOCAL
     override val hasDefaultValue: Boolean get() = withValidityAssertion { fir.defaultValue != null }
 }

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

@@ -7,13 +7,16 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirField
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.modality
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCommonSymbolModality
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtJavaFieldSymbol
@@ -22,15 +25,16 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirJavaFieldSymbol(
     fir: FirField,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder
 ) : KtJavaFieldSymbol(), KtFirSymbol<FirField> {
-    override val fir: FirField by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
 
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
-    override val isVal: Boolean get() = withValidityAssertion { fir.isVal }
-    override val name: Name get() = withValidityAssertion { fir.name }
+    override val type: KtType by firRef.withFirAndCache(FirResolvePhase.TYPES) { fir -> builder.buildKtType(fir.returnTypeRef) }
+    override val isVal: Boolean get() = firRef.withFir { it.isVal }
+    override val name: Name get() = firRef.withFir { it.name }
 
-    override val modality: KtCommonSymbolModality get() = withValidityAssertion { fir.modality.getSymbolModality() }
+    override val modality: KtCommonSymbolModality get() = firRef.withFir { it.modality.getSymbolModality() }
 }

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

@@ -7,12 +7,15 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtLocalVariableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolKind
@@ -21,6 +24,7 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirLocalVariableSymbol(
     fir: FirProperty,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder
 ) : KtLocalVariableSymbol(),
@@ -29,11 +33,11 @@ internal class KtFirLocalVariableSymbol(
         assert(fir.isLocal)
     }
 
-    override val fir: FirProperty by weakRef(fir)
+    override val firRef = firRef(fir, resolveState)
     override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
 
-    override val isVal: Boolean get() = withValidityAssertion { fir.isVal }
-    override val name: Name get() = withValidityAssertion { fir.name }
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
+    override val isVal: Boolean get() = firRef.withFir { it.isVal }
+    override val name: Name get() = firRef.withFir { it.name }
+    override val type: KtType by firRef.withFirAndCache(FirResolvePhase.BODY_RESOLVE) { fir -> builder.buildKtType(fir.returnTypeRef) }
     override val symbolKind: KtSymbolKind get() = KtSymbolKind.LOCAL
 }

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

@@ -8,13 +8,16 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.modality
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCommonSymbolModality
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPropertySymbol
@@ -26,6 +29,7 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirPropertySymbol(
     fir: FirProperty,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken,
     private val builder: KtSymbolByFirBuilder,
     private val forcedOrigin: FirDeclarationOrigin?
@@ -34,23 +38,23 @@ internal class KtFirPropertySymbol(
         assert(!fir.isLocal)
     }
 
-    override val fir: FirProperty by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
 
     override val origin: KtSymbolOrigin get() = withValidityAssertion { forcedOrigin?.asKtSymbolOrigin() ?: super.origin }
 
-    override val fqName: FqName get() = withValidityAssertion { fir.symbol.callableId.asFqNameForDebugInfo() }
-    override val isVal: Boolean get() = withValidityAssertion { fir.isVal }
-    override val name: Name get() = withValidityAssertion { fir.name }
-    override val type: KtType by cached { builder.buildKtType(fir.returnTypeRef) }
-    override val receiverType: KtType? by cached { builder.buildKtType(fir.returnTypeRef) }
-    override val isExtension: Boolean get() = withValidityAssertion { fir.receiverTypeRef != null }
+    override val fqName: FqName get() = firRef.withFir { it.symbol.callableId.asFqNameForDebugInfo() }
+    override val isVal: Boolean get() = firRef.withFir { it.isVal }
+    override val name: Name get() = firRef.withFir { it.name }
+    override val type: KtType by firRef.withFirAndCache(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) { fir -> builder.buildKtType(fir.returnTypeRef) }
+    override val receiverType: KtType? by firRef.withFirAndCache(FirResolvePhase.TYPES) { fir -> builder.buildKtType(fir.returnTypeRef) }
+    override val isExtension: Boolean get() = firRef.withFir { it.receiverTypeRef != null }
     override val symbolKind: KtSymbolKind
-        get() = withValidityAssertion {
+        get() = firRef.withFir { fir ->
             when (fir.symbol.callableId.classId) {
                 null -> KtSymbolKind.TOP_LEVEL
                 else -> KtSymbolKind.MEMBER
             }
         }
-    override val modality: KtCommonSymbolModality get() = withValidityAssertion { fir.modality.getSymbolModality() }
+    override val modality: KtCommonSymbolModality get() = firRef.withFir { it.modality.getSymbolModality() }
 }

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

@@ -7,14 +7,17 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 
 internal interface KtFirSymbol<F : FirDeclaration> : KtSymbol, ValidityTokenOwner {
-    val fir: F
-    override val origin: KtSymbolOrigin get() = withValidityAssertion { fir.origin.asKtSymbolOrigin() }
+    val firRef: FirRef<F>
+
+    override val origin: KtSymbolOrigin get() = firRef.withFir { it.origin.asKtSymbolOrigin() }
 }
 
 internal fun FirDeclarationOrigin.asKtSymbolOrigin() = when (this) {

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

@@ -8,10 +8,12 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.ReadOnlyWeakRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtTypeAliasSymbol
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
@@ -20,10 +22,11 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirTypeAliasSymbol(
     fir: FirTypeAlias,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken
 ) : KtTypeAliasSymbol(), KtFirSymbol<FirTypeAlias> {
-    override val fir: FirTypeAlias by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
-    override val name: Name get() = withValidityAssertion { fir.name }
-    override val classId: ClassId get() = withValidityAssertion { fir.symbol.classId }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
+    override val name: Name get() = firRef.withFir { it.name }
+    override val classId: ClassId get() = firRef.withFir { it.symbol.classId }
 }

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

@@ -8,9 +8,11 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
 import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.cached
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtTypeParameterSymbol
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
@@ -18,10 +20,11 @@ import org.jetbrains.kotlin.name.Name
 
 internal class KtFirTypeParameterSymbol(
     fir: FirTypeParameter,
+    resolveState: FirModuleResolveState,
     override val token: ValidityToken
 ) : KtTypeParameterSymbol(), KtFirSymbol<FirTypeParameter> {
-    override val fir: FirTypeParameter by weakRef(fir)
-    override val psi: PsiElement? by cached { fir.findPsi(fir.session) }
+    override val firRef = firRef(fir, resolveState)
+    override val psi: PsiElement? by firRef.withFirAndCache { it.findPsi(fir.session) }
 
-    override val name: Name get() = withValidityAssertion { fir.name }
+    override val name: Name get() = firRef.withFir { it.name }
 }

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

@@ -0,0 +1,36 @@
+/*
+ * 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.frontend.api.fir.utils
+
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
+import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
+import org.jetbrains.kotlin.idea.frontend.api.assertIsValid
+import java.lang.ref.WeakReference
+
+internal class FirRef<D : FirDeclaration>(fir: D, resolveState: FirModuleResolveState, val token: ValidityToken) {
+    private val firWeakRef = WeakReference(fir)
+    private val resolveStateWeakRef = WeakReference(resolveState)
+
+    inline fun <R> withFir(phase: FirResolvePhase = FirResolvePhase.RAW_FIR, action: (fir: D) -> R): R {
+        token.assertIsValid()
+        val fir = firWeakRef.get() ?: error("FirElement was garbage collected while analysis session is still valid")
+        val resolveState =
+            resolveStateWeakRef.get() ?: error("FirModuleResolveState was garbage collected while analysis session is still valid")
+        return action(resolveState.resolvedFirToPhase(fir, phase))
+    }
+
+    inline fun <R> withFirAndCache(phase: FirResolvePhase = FirResolvePhase.RAW_FIR, crossinline createValue: (fir: D) -> R) =
+        ValidityAwareCachedValue(token) {
+            withFir(phase) { fir -> createValue(fir) }
+        }
+}
+
+@Suppress("NOTHING_TO_INLINE")
+internal inline fun <D : FirDeclaration> ValidityTokenOwner.firRef(fir: D, resolveState: FirModuleResolveState) =
+    FirRef(fir, resolveState, token)

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

@@ -8,8 +8,7 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.utils
 import kotlin.reflect.KProperty
 
 
-@Suppress("EXPERIMENTAL_FEATURE_WARNING")
-internal inline class ThreadLocalValue<V>(private val threadLocal: ThreadLocal<V>) {
+internal class ThreadLocalValue<V>(private val threadLocal: ThreadLocal<V>) {
     @Suppress("NOTHING_TO_INLINE")
     inline operator fun getValue(thisRef: Any?, property: KProperty<*>): V = threadLocal.get()
 }