Analysis API: get rid of weak refs scheme for working with FIR internals for KtScope

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirScopeProvider.kt

@@ -50,7 +50,6 @@ internal class KtFirScopeProvider(
     override val analysisSession: KtFirAnalysisSession by weakRef(analysisSession)
     private val builder by weakRef(builder)
     private val firResolveState by weakRef(firResolveState)
-    private val firScopeStorage = FirScopeRegistry()
 
     private val memberScopeCache = IdentityHashMap<KtSymbolWithMembers, KtMemberScope>()
     private val declaredMemberScopeCache = IdentityHashMap<KtSymbolWithMembers, KtDeclaredMemberScope>()
@@ -80,7 +79,6 @@ internal class KtFirScopeProvider(
                 )
             } ?: return@getOrPut KtFirEmptyMemberScope(classSymbol)
 
-            firScopeStorage.register(firScope)
             KtFirMemberScope(classSymbol, firScope, token, builder)
         }
     }
@@ -89,7 +87,6 @@ internal class KtFirScopeProvider(
         val firScope = symbol.withFirForScope { fir ->
             fir.scopeProvider.getStaticScope(fir, analysisSession.rootModuleSession, ScopeSession())
         } ?: return KtFirEmptyMemberScope(symbol)
-        firScopeStorage.register(firScope)
         check(firScope is FirContainingNamesAwareScope)
         return KtFirDelegatingScopeImpl(firScope, builder, token)
     }
@@ -100,8 +97,6 @@ internal class KtFirScopeProvider(
                 analysisSession.rootModuleSession.declaredMemberScope(it)
             } ?: return@getOrPut KtFirEmptyMemberScope(classSymbol)
 
-            firScopeStorage.register(firScope)
-
             KtFirDeclaredMemberScope(classSymbol, firScope, token, builder)
         }
     }
@@ -119,7 +114,6 @@ internal class KtFirScopeProvider(
                 packageSymbol.fqName,
                 project,
                 builder,
-                this,
                 token,
                 analysisSession.searchScope,
                 analysisSession.targetPlatform
@@ -127,9 +121,6 @@ internal class KtFirScopeProvider(
         }
     }
 
-    fun registerScope(scope: FirScope) {
-        firScopeStorage.register(scope)
-    }
 
     override fun getCompositeScope(subScopes: List<KtScope>): KtCompositeScope = withValidityAssertion {
         KtFirCompositeScope(subScopes, token)
@@ -190,7 +181,6 @@ internal class KtFirScopeProvider(
     }
 
     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)
@@ -198,7 +188,6 @@ internal class KtFirScopeProvider(
                 firScope.fqName,
                 project,
                 builder,
-                this,
                 token,
                 analysisSession.searchScope,
                 analysisSession.targetPlatform
@@ -211,21 +200,7 @@ internal class KtFirScopeProvider(
 }
 
 private class KtFirDelegatingScopeImpl<S>(
-    firScope: S, builder: KtSymbolByFirBuilder,
+    override val firScope: S,
+    builder: KtSymbolByFirBuilder,
     token: ValidityToken
-) : KtFirDelegatingScope<S>(builder, token), ValidityTokenOwner where S : FirContainingNamesAwareScope, S : FirScope {
-    override val firScope: S by weakRef(firScope)
-}
-
-/**
- * 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] & [KtAnalysisSessionComponent]
- * Entities which needs storing [FirScope] instances will store them as weak references via [org.jetbrains.kotlin.analysis.api.fir.utils.weakRef]
- */
-internal class FirScopeRegistry {
-    private val scopes = mutableListOf<FirScope>()
-
-    fun register(scope: FirScope) {
-        scopes += scope
-    }
-}
+) : KtFirDelegatingScope<S>(builder, token), ValidityTokenOwner where S : FirContainingNamesAwareScope, S : FirScope

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

@@ -15,12 +15,9 @@ import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
 
 internal class KtFirDeclaredMemberScope(
     override val owner: KtSymbolWithMembers,
-    firScope: FirClassDeclaredMemberScope,
+    override val firScope: FirClassDeclaredMemberScope,
     token: ValidityToken,
     builder: KtSymbolByFirBuilder
 ) : KtFirDelegatingScope<FirClassDeclaredMemberScope>(builder, token),
     KtDeclaredMemberScope,
-    ValidityTokenOwner {
-    override val firScope: FirClassDeclaredMemberScope by weakRef(firScope)
-}
-
+    ValidityTokenOwner

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

@@ -15,10 +15,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
 
 internal class KtFirMemberScope(
     override val owner: KtSymbolWithMembers,
-    firScope: FirTypeScope,
+    override val firScope: FirTypeScope,
     token: ValidityToken,
     builder: KtSymbolByFirBuilder
-) : KtFirDelegatingScope<FirTypeScope>(builder, token), KtMemberScope, ValidityTokenOwner {
-    override val firScope: FirTypeScope by weakRef(firScope)
-}
-
+) : KtFirDelegatingScope<FirTypeScope>(builder, token), KtMemberScope, ValidityTokenOwner

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirNonStarImportingScope.kt

@@ -22,11 +22,10 @@ import org.jetbrains.kotlin.analysis.api.withValidityAssertion
 import org.jetbrains.kotlin.name.Name
 
 internal class KtFirNonStarImportingScope(
-    firScope: FirAbstractSimpleImportingScope,
+    private val firScope : FirAbstractSimpleImportingScope,
     private val builder: KtSymbolByFirBuilder,
     override val token: ValidityToken
 ) : KtNonStarImportingScope, ValidityTokenOwner {
-    private val firScope: FirAbstractSimpleImportingScope by weakRef(firScope)
 
     @OptIn(ExperimentalStdlibApi::class)
     override val imports: List<NonStarImport> by cached {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirPackageScope.kt

@@ -14,8 +14,6 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirPackageMemberScope
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.analysis.api.fir.components.KtFirScopeProvider
-import org.jetbrains.kotlin.analysis.api.fir.utils.lazyThreadUnsafeWeakRef
-import org.jetbrains.kotlin.analysis.api.fir.utils.weakRef
 import org.jetbrains.kotlin.analysis.api.scopes.KtPackageScope
 import org.jetbrains.kotlin.analysis.api.scopes.KtScopeNameFilter
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
@@ -31,19 +29,15 @@ internal class KtFirPackageScope(
     override val fqName: FqName,
     private val project: Project,
     private val builder: KtSymbolByFirBuilder,
-    _scopeProvider: KtFirScopeProvider,
     override val token: ValidityToken,
     private val searchScope: GlobalSearchScope,
     private val targetPlatform: TargetPlatform,
 ) : KtPackageScope {
-    private val scopeProvider by weakRef(_scopeProvider)
     private val declarationsProvider = project.createDeclarationProvider(searchScope)
     private val packageProvider = project.createPackageProvider(searchScope)
 
-    private val firScope: FirPackageMemberScope by lazyThreadUnsafeWeakRef {
-        val scope = FirPackageMemberScope(fqName, builder.rootSession)
-        scopeProvider.registerScope(scope)
-        scope
+    private val firScope: FirPackageMemberScope by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        FirPackageMemberScope(fqName, builder.rootSession)
     }
 
     override fun getPossibleCallableNames() = withValidityAssertion {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirStarImportingScope.kt

@@ -28,12 +28,11 @@ import org.jetbrains.kotlin.analysis.api.withValidityAssertion
 import org.jetbrains.kotlin.name.Name
 
 internal class KtFirStarImportingScope(
-    firScope: FirAbstractStarImportingScope,
+    private val firScope: FirAbstractStarImportingScope,
     private val builder: KtSymbolByFirBuilder,
     project: Project,
     override val token: ValidityToken,
 ) : KtStarImportingScope, ValidityTokenOwner {
-    private val firScope: FirAbstractStarImportingScope by weakRef(firScope)
     override val isDefaultImportingScope: Boolean = withValidityAssertion { firScope is FirDefaultStarImportingScope }
 
     //todo use more concrete scope

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/utils/ReadOnlyWeakRef.kt

@@ -1,32 +0,0 @@
-/*
- * 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.analysis.api.fir.utils
-
-import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
-import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
-import java.lang.ref.WeakReference
-import kotlin.reflect.KProperty
-
-internal class LazyThreadUnsafeWeakRef<V : Any>
-constructor(_createValue: () -> V, val token: ValidityToken) {
-    private var createValue: (() -> V)? = _createValue
-    private var weakRef: WeakReference<V>? = null
-
-    @Suppress("NOTHING_TO_INLINE")
-    inline operator fun getValue(thisRef: Any, property: KProperty<*>): V {
-        if (weakRef == null) {
-            weakRef = WeakReference(createValue!!())
-            createValue = null
-        }
-        return weakRef!!.get() ?: if (token.isValid()) {
-            throw EntityWasGarbageCollectedException(property.toString())
-        } else {
-            error("Accessing the invalid value of $property")
-        }
-    }
-}
-
-internal fun <V : Any> ValidityTokenOwner.lazyThreadUnsafeWeakRef(createValue: () -> V) = LazyThreadUnsafeWeakRef(createValue, token)