FIR IDE: store FirScope strong references only in KtFirAnalysisSession

compiler/fir/tree/src/org/jetbrains/kotlin/fir/scopes/FirCompositeScope.kt

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.name.Name
 
-class FirCompositeScope(private val scopes: Iterable<FirScope>) : FirScope() {
+class FirCompositeScope(val scopes: Iterable<FirScope>) : FirScope() {
 
     override fun processClassifiersByNameWithSubstitution(
         name: Name,

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

@@ -15,10 +15,10 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.firClassLike
+import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.fir.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
@@ -59,7 +59,10 @@ private constructor(
             firSymbolBuilder
         )
 
-    override val scopeProvider: KtScopeProvider = KtFirScopeProvider(token, firSymbolBuilder, project, firSession, firResolveState)
+    private val firScopeStorage = FirScopeRegistry()
+
+    override val scopeProvider: KtScopeProvider =
+        KtFirScopeProvider(token, firSymbolBuilder, project, firSession, firResolveState, firScopeStorage)
 
     init {
         assertIsValid()
@@ -221,4 +224,17 @@ private constructor(
             )
         }
     }
+}
+
+/**
+ * Stores strong references to all instances of [FirScope] used
+ * Needed as the only entity which may have a strong references to FIR internals is [KtFirAnalysisSession]
+ * Entities which needs storing [FirScope] instances will store them as weak references via [org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef]
+ */
+internal class FirScopeRegistry {
+    private val scopes = mutableListOf<FirScope>()
+
+    fun register(scope: FirScope) {
+        scopes += scope
+    }
 }

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

@@ -17,12 +17,13 @@ import org.jetbrains.kotlin.idea.frontend.api.scopes.KtUnsubstitutedScope
 
 internal class KtFirDeclaredMemberScope(
     override val owner: KtFirClassOrObjectSymbol,
+    firScope: FirScope,
     token: ValidityToken,
     builder: KtSymbolByFirBuilder
 ) : KtFirDelegatingScope(builder, token),
     KtDeclaredMemberScope,
     KtUnsubstitutedScope<KtFirDelegatingScope>,
     ValidityTokenOwner {
-    override val firScope: FirScope by weakRef(declaredMemberScope(owner.fir))
+    override val firScope: FirScope by weakRef(firScope)
 }
 

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

@@ -5,8 +5,6 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.scopes
 
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.buildUseSiteMemberScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
@@ -18,11 +16,10 @@ import org.jetbrains.kotlin.idea.frontend.api.scopes.KtUnsubstitutedScope
 
 internal class KtFirMemberScope(
     override val owner: KtFirClassOrObjectSymbol,
+    firScope: FirScope,
     token: ValidityToken,
     builder: KtSymbolByFirBuilder
 ) : KtFirDelegatingScope(builder, token), KtMemberScope, KtUnsubstitutedScope<KtMemberScope>, ValidityTokenOwner {
-    override val firScope: FirScope by weakRef {
-        owner.fir.buildUseSiteMemberScope(owner.fir.session, ScopeSession())
-    }
+    override val firScope: FirScope by weakRef(firScope)
 }
 

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

@@ -9,15 +9,18 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.buildUseSiteMemberScope
 import org.jetbrains.kotlin.fir.resolve.scope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractSimpleImportingScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractStarImportingScope
+import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
 import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
+import org.jetbrains.kotlin.idea.frontend.api.fir.FirScopeRegistry
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirPackageSymbol
@@ -40,9 +43,12 @@ internal class KtFirScopeProvider(
     private val builder: KtSymbolByFirBuilder,
     private val project: Project,
     session: FirSession,
-    val firResolveState: FirModuleResolveState
+    firResolveState: FirModuleResolveState,
+    firScopeRegistry: FirScopeRegistry,
 ) : KtScopeProvider(), ValidityTokenOwner {
     private val session by weakRef(session)
+    private val firResolveState by weakRef(firResolveState)
+    private val firScopeStorage by weakRef(firScopeRegistry)
 
     private val memberScopeCache = IdentityHashMap<KtClassOrObjectSymbol, KtMemberScope>()
     private val declaredMemberScopeCache = IdentityHashMap<KtClassOrObjectSymbol, KtDeclaredMemberScope>()
@@ -51,14 +57,16 @@ internal class KtFirScopeProvider(
     override fun getMemberScope(classSymbol: KtClassOrObjectSymbol): KtMemberScope = withValidityAssertion {
         memberScopeCache.getOrPut(classSymbol) {
             check(classSymbol is KtFirClassOrObjectSymbol)
-            KtFirMemberScope(classSymbol, token, builder)
+            val firScope = classSymbol.fir.buildUseSiteMemberScope(classSymbol.fir.session, ScopeSession()).also(firScopeStorage::register)
+            KtFirMemberScope(classSymbol, firScope, token, builder)
         }
     }
 
     override fun getDeclaredMemberScope(classSymbol: KtClassOrObjectSymbol): KtDeclaredMemberScope = withValidityAssertion {
         declaredMemberScopeCache.getOrPut(classSymbol) {
             check(classSymbol is KtFirClassOrObjectSymbol)
-            KtFirDeclaredMemberScope(classSymbol, token, builder)
+            val firScope = declaredMemberScope(classSymbol.fir).also(firScopeStorage::register)
+            KtFirDeclaredMemberScope(classSymbol, firScope, token, builder)
         }
     }
 
@@ -110,17 +118,20 @@ internal class KtFirScopeProvider(
         KtScopeContext(getCompositeScope(allKtScopes), implicitReceiversTypes)
     }
 
-    private fun convertToKtScope(firScope: FirScope): KtScope = when (firScope) {
-        is FirAbstractSimpleImportingScope -> KtFirNonStarImportingScope(firScope, builder, token)
-        is FirAbstractStarImportingScope -> KtFirStarImportingScope(firScope, builder, project, token)
-        else -> {
-            // todo create concrete KtScope here instead of a generic one
-            KtFirDelegatingScopeImpl(firScope, builder, token)
+    private fun convertToKtScope(firScope: FirScope): KtScope {
+        firScopeStorage.register(firScope)
+        return when (firScope) {
+            is FirAbstractSimpleImportingScope -> KtFirNonStarImportingScope(firScope, builder, token)
+            is FirAbstractStarImportingScope -> KtFirStarImportingScope(firScope, builder, project, token)
+            else -> {
+                // todo create concrete KtScope here instead of a generic one
+                KtFirDelegatingScopeImpl(firScope, builder, token)
+            }
         }
     }
 }
 
 private class KtFirDelegatingScopeImpl(firScope: FirScope, builder: KtSymbolByFirBuilder, token: ValidityToken) :
     KtFirDelegatingScope(builder, token), ValidityTokenOwner {
-    override val firScope: FirScope = firScope
-}
+    override val firScope: FirScope by weakRef(firScope)
+}