Analysis API: get rid of weak refs scheme for working with FIR internals for KtSymbol and KtType

2021-09-22 17:50:12 +02:00
parent 2de4a6f924
commit d348c48e4f
21 changed files with 61 additions and 121 deletions
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.api.fir
 import com.intellij.openapi.project.Project
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.kotlin.analysis.api.KtStarProjectionTypeArgument
 import org.jetbrains.kotlin.analysis.api.KtTypeArgument
 import org.jetbrains.kotlin.analysis.api.KtTypeArgumentWithVariance
@@ -411,7 +412,7 @@ private class BuilderCache<From, To: Any> private constructor(
    private val cache: ConcurrentMap<From, To>,
    private val isReadOnly: Boolean
 ) {
-    constructor() : this(ConcurrentHashMap<From, To>(), isReadOnly = false)
+    constructor() : this(ContainerUtil.createConcurrentSoftMap(), isReadOnly = false)

    fun createReadOnlyCopy(): BuilderCache<From, To> {
        check(!isReadOnly) { "Cannot create readOnly BuilderCache from a readonly one" }
@@ -26,9 +26,8 @@ internal class KtFirAnonymousFunctionSymbol(
    fir: FirAnonymousFunction,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtAnonymousFunctionSymbol(), KtFirSymbol<FirAnonymousFunction> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -29,9 +29,8 @@ internal class KtFirAnonymousObjectSymbol(
    fir: FirAnonymousObject,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtAnonymousObjectSymbol(), KtFirSymbol<FirAnonymousObject> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -23,9 +23,8 @@ internal class KtFirBackingFieldSymbol(
    propertyFir: FirProperty,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
-) : KtBackingFieldSymbol(){
-    private val builder by weakRef(_builder)
+    private val builder: KtSymbolByFirBuilder
+) : KtBackingFieldSymbol() {
    private val propertyFirRef = firRef(propertyFir, resolveState)

    override val annotatedType: KtTypeAndAnnotations by cached {
@@ -39,9 +39,8 @@ internal class KtFirConstructorSymbol(
    fir: FirConstructor,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtConstructorSymbol(), KtFirSymbol<FirConstructor> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -29,9 +29,8 @@ internal class KtFirEnumEntrySymbol(
    fir: FirEnumEntry,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtEnumEntrySymbol(), KtFirSymbol<FirEnumEntry> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)

    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }
@@ -40,9 +40,8 @@ internal class KtFirFunctionSymbol(
    fir: FirSimpleFunction,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtFunctionSymbol(), KtFirSymbol<FirSimpleFunction> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }
    override val name: Name get() = firRef.withFir { it.name }
@@ -30,9 +30,8 @@ internal class KtFirJavaFieldSymbol(
    fir: FirField,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtJavaFieldSymbol(), KtFirSymbol<FirField> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -44,14 +44,13 @@ internal class KtFirKotlinPropertySymbol(
    fir: FirProperty,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtKotlinPropertySymbol(), KtFirSymbol<FirProperty> {
    init {
        assert(!fir.isLocal)
        check(fir !is FirSyntheticProperty)
    }

-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -27,14 +27,13 @@ internal class KtFirLocalVariableSymbol(
    fir: FirProperty,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtLocalVariableSymbol(),
    KtFirSymbol<FirProperty> {
    init {
        assert(fir.isLocal)
    }

-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -38,9 +38,8 @@ internal class KtFirNamedClassOrObjectSymbol(
    fir: FirRegularClass,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtNamedClassOrObjectSymbol(), KtFirSymbol<FirRegularClass> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }
    override val name: Name get() = firRef.withFir { it.name }
@@ -39,13 +39,12 @@ internal class KtFirPropertyGetterSymbol(
    fir: FirPropertyAccessor,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtPropertyGetterSymbol(), KtFirSymbol<FirPropertyAccessor> {
    init {
        require(fir.isGetter)
    }

-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -39,13 +39,12 @@ internal class KtFirPropertySetterSymbol(
    fir: FirPropertyAccessor,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtPropertySetterSymbol(), KtFirSymbol<FirPropertyAccessor> {
    init {
        require(fir.isSetter)
    }

-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -29,9 +29,8 @@ internal class KtFirSamConstructorSymbol(
    fir: FirSimpleFunction,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtSamConstructorSymbol(), KtFirSymbol<FirSimpleFunction> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }
    override val name: Name get() = firRef.withFir { it.name }
@@ -40,9 +40,8 @@ internal class KtFirSyntheticJavaPropertySymbol(
    fir: FirSyntheticProperty,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtSyntheticJavaPropertySymbol(), KtFirSymbol<FirSyntheticProperty> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -29,9 +29,8 @@ internal class KtFirTypeAliasSymbol(
    fir: FirTypeAlias,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtTypeAliasSymbol(), KtFirSymbol<FirTypeAlias> {
-    private val builder by weakRef(_builder)

    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }
@@ -20,14 +20,14 @@ import org.jetbrains.kotlin.analysis.api.symbols.markers.KtAnnotationCall
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtTypeAndAnnotations
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.analysis.api.withValidityAssertion

 internal class KtFirTypeAndAnnotations<T : FirDeclaration>(
    private val containingDeclaration: FirRefWithValidityCheck<T>,
    @Suppress("UNUSED_PARAMETER") typeResolvePhase: FirResolvePhase,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
    private val typeRef: (T) -> FirTypeRef,
 ) : KtTypeAndAnnotations() {
-    private val builder by weakRef(_builder)

    override val token: ValidityToken get() = containingDeclaration.token

@@ -43,20 +43,16 @@ internal class KtFirTypeAndAnnotations<T : FirDeclaration>(
 }

 internal class KtSimpleFirTypeAndAnnotations(
-    coneType: ConeKotlinType,
-    annotationsList: List<KtAnnotationCall>,
-    builder: KtSymbolByFirBuilder,
+    private val coneType: ConeKotlinType,
+    private val _annotations: List<KtAnnotationCall>,
+    private val builder: KtSymbolByFirBuilder,
    override val token: ValidityToken
 ) : KtTypeAndAnnotations() {
-
-    private val coneTypeRef by weakRef(coneType)
-    private val annotationsListRef by weakRef(annotationsList)
-
    override val type: KtType by cached {
-        builder.typeBuilder.buildKtType(coneTypeRef)
+        builder.typeBuilder.buildKtType(coneType)
    }

-    override val annotations: List<KtAnnotationCall> get() = annotationsListRef
+    override val annotations: List<KtAnnotationCall> get() = withValidityAssertion { _annotations }
 }

 internal fun FirRefWithValidityCheck<FirClass>.superTypesAndAnnotationsList(builder: KtSymbolByFirBuilder): List<KtTypeAndAnnotations> =
@@ -66,7 +62,7 @@ internal fun FirRefWithValidityCheck<FirClass>.superTypesAndAnnotationsList(buil

 internal fun FirRefWithValidityCheck<FirRegularClass>.superTypesAndAnnotationsListForRegularClass(builder: KtSymbolByFirBuilder): List<KtTypeAndAnnotations> {
    return withFir { fir ->
-        if(fir.resolvePhase >= FirResolvePhase.SUPER_TYPES) {
+        if (fir.resolvePhase >= FirResolvePhase.SUPER_TYPES) {
            fir.superTypeRefs.mapToTypeAndAnnotations(this, builder)
        } else null
    } ?: withFirByType(ResolveType.NoResolve) { fir ->
@@ -25,9 +25,8 @@ internal class KtFirTypeParameterSymbol(
    fir: FirTypeParameter,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtTypeParameterSymbol(), KtFirSymbol<FirTypeParameter> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -34,9 +34,8 @@ internal class KtFirValueParameterSymbol(
    fir: FirValueParameter,
    resolveState: FirModuleResolveState,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder
+    private val builder: KtSymbolByFirBuilder
 ) : KtValueParameterSymbol(), KtFirSymbol<FirValueParameter> {
-    private val builder by weakRef(_builder)
    override val firRef = firRef(fir, resolveState)
    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.moduleData.session) }

@@ -5,16 +5,11 @@

 package org.jetbrains.kotlin.analysis.api.fir.types

-import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
-import org.jetbrains.kotlin.fir.resolve.inference.receiverType
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.analysis.api.fir.getCandidateSymbols
 import org.jetbrains.kotlin.analysis.api.KtTypeArgument
 import org.jetbrains.kotlin.analysis.api.KtTypeArgumentWithVariance
 import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
 import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.analysis.api.fir.getCandidateSymbols
 import org.jetbrains.kotlin.analysis.api.fir.utils.cached
 import org.jetbrains.kotlin.analysis.api.fir.utils.weakRef
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
@@ -22,6 +17,11 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtTypeParameterSymbol
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.types.*
 import org.jetbrains.kotlin.analysis.api.withValidityAssertion
+import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
+import org.jetbrains.kotlin.fir.resolve.inference.receiverType
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name

@@ -38,13 +38,10 @@ private fun KtFirType.typeEquals(other: Any?): Boolean {
 private fun KtFirType.typeHashcode(): Int = token.hashCode() * 31 + coneType.hashCode()

 internal class KtFirUsualClassType(
-    _coneType: ConeClassLikeTypeImpl,
+    override val coneType: ConeClassLikeTypeImpl,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtUsualClassType(), KtFirType {
-    override val coneType by weakRef(_coneType)
-    private val builder by weakRef(_builder)
-
    override val classId: ClassId get() = withValidityAssertion { coneType.lookupTag.classId }
    override val classSymbol: KtClassLikeSymbol by cached {
        builder.classifierBuilder.buildClassLikeSymbolByLookupTag(coneType.lookupTag)
@@ -63,13 +60,10 @@ internal class KtFirUsualClassType(
 }

 internal class KtFirFunctionalType(
-    _coneType: ConeClassLikeTypeImpl,
+    override val coneType: ConeClassLikeTypeImpl,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtFunctionalType(), KtFirType {
-    override val coneType by weakRef(_coneType)
-    private val builder by weakRef(_builder)
-
    override val classId: ClassId get() = withValidityAssertion { coneType.lookupTag.classId }
    override val classSymbol: KtClassLikeSymbol by cached {
        builder.classifierBuilder.buildClassLikeSymbolByLookupTag(coneType.lookupTag)
@@ -116,12 +110,10 @@ internal class KtFirFunctionalType(
 }

 internal class KtFirClassErrorType(
-    _coneType: ConeClassErrorType,
+    override val coneType: ConeClassErrorType,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtClassErrorType(), KtFirType {
-    override val coneType by weakRef(_coneType)
-    private val builder by weakRef(_builder)

    override val error: String get() = withValidityAssertion { coneType.diagnostic.reason }
    override val nullability: KtTypeNullability get() = withValidityAssertion { coneType.nullability.asKtNullability() }
@@ -137,10 +129,9 @@ internal class KtFirClassErrorType(
 }

 internal class KtFirCapturedType(
-    _coneType: ConeCapturedType,
+    override val coneType: ConeCapturedType,
    override val token: ValidityToken,
 ) : KtCapturedType(), KtFirType {
-    override val coneType by weakRef(_coneType)
    override val nullability: KtTypeNullability get() = withValidityAssertion { coneType.nullability.asKtNullability() }
    override fun asStringForDebugging(): String = withValidityAssertion { coneType.render() }
    override fun equals(other: Any?) = typeEquals(other)
@@ -148,13 +139,10 @@ internal class KtFirCapturedType(
 }

 internal class KtFirDefinitelyNotNullType(
-    _coneType: ConeDefinitelyNotNullType,
+    override val coneType: ConeDefinitelyNotNullType,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtDefinitelyNotNullType(), KtFirType {
-    override val coneType by weakRef(_coneType)
-    private val builder by weakRef(_builder)
-
    override val original: KtType by cached { builder.typeBuilder.buildKtType(this.coneType.original) }

    override fun asStringForDebugging(): String = withValidityAssertion { coneType.render() }
@@ -163,13 +151,10 @@ internal class KtFirDefinitelyNotNullType(
 }

 internal class KtFirTypeParameterType(
-    _coneType: ConeTypeParameterType,
+    override val coneType: ConeTypeParameterType,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtTypeParameterType(), KtFirType {
-    override val coneType by weakRef(_coneType)
-    private val builder by weakRef(_builder)
-
    override val name: Name get() = withValidityAssertion { coneType.lookupTag.name }
    override val symbol: KtTypeParameterSymbol by cached {
        builder.classifierBuilder.buildTypeParameterSymbolByLookupTag(coneType.lookupTag)
@@ -184,12 +169,10 @@ internal class KtFirTypeParameterType(
 }

 internal class KtFirFlexibleType(
-    _coneType: ConeFlexibleType,
+    override val coneType: ConeFlexibleType,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtFlexibleType(), KtFirType {
-    override val coneType by weakRef(_coneType)
-    private val builder by weakRef(_builder)

    override val lowerBound: KtType by cached { builder.typeBuilder.buildKtType(coneType.lowerBound) }
    override val upperBound: KtType by cached { builder.typeBuilder.buildKtType(coneType.upperBound) }
@@ -202,13 +185,10 @@ internal class KtFirFlexibleType(
 }

 internal class KtFirIntersectionType(
-    _coneType: ConeIntersectionType,
+    override val coneType: ConeIntersectionType,
    override val token: ValidityToken,
-    _builder: KtSymbolByFirBuilder,
+    private val builder: KtSymbolByFirBuilder,
 ) : KtIntersectionType(), KtFirType {
-    override val coneType by weakRef(_coneType)
-    private val builder by weakRef(_builder)
-
    override val conjuncts: List<KtType> by cached {
        coneType.intersectedTypes.map { conjunct -> builder.typeBuilder.buildKtType(conjunct) }
    }
@@ -5,31 +5,23 @@

 package org.jetbrains.kotlin.analysis.api.fir.utils

-import org.jetbrains.annotations.TestOnly
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDeclaration
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.tokens.assertIsValidAndAccessible
-import java.lang.ref.WeakReference
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveType
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase

-internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveState: FirModuleResolveState, val token: ValidityToken) {
-    private val firWeakRef = WeakReference(fir)
-    private val resolveStateWeakRef = WeakReference(resolveState)
-
-    @TestOnly
-    internal fun isCollected(): Boolean =
-        firWeakRef.get() == null && resolveStateWeakRef.get() == null
+internal class FirRefWithValidityCheck<out D : FirDeclaration>(
+    private val fir: D,
+    val resolveState: FirModuleResolveState,
+    val token: ValidityToken
+) {

    inline fun <R> withFir(phase: FirResolvePhase = FirResolvePhase.RAW_FIR, crossinline action: (fir: D) -> R): R {
        token.assertIsValidAndAccessible()
-        val fir = firWeakRef.get()
-            ?: throw EntityWasGarbageCollectedException("FirElement")
-        val resolveState = resolveStateWeakRef.get()
-            ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
        return fir.withFirDeclaration(resolveState, phase) { action(it) }
    }

@@ -40,14 +32,9 @@ internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveSt
     */
    inline fun <R> withFirUnsafe(action: (fir: D) -> R): R {
        token.assertIsValidAndAccessible()
-        val fir = firWeakRef.get()
-            ?: throw EntityWasGarbageCollectedException("FirElement")
        return action(fir)
    }

-    val resolveState
-        get() = resolveStateWeakRef.get() ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
-
    inline fun <R> withFirAndCache(phase: FirResolvePhase = FirResolvePhase.RAW_FIR, crossinline createValue: (fir: D) -> R) =
        ValidityAwareCachedValue(token) {
            withFir(phase) { fir -> createValue(fir) }
@@ -55,10 +42,6 @@ internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveSt

    inline fun <R> withFirByType(type: ResolveType, crossinline action: (fir: D) -> R): R {
        token.assertIsValidAndAccessible()
-        val fir = firWeakRef.get()
-            ?: throw EntityWasGarbageCollectedException("FirElement")
-        val resolveState = resolveStateWeakRef.get()
-            ?: throw EntityWasGarbageCollectedException("FirModuleResolveState")
        return fir.withFirDeclaration(type, resolveState) { action(it) }
    }

@@ -69,16 +52,13 @@ internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveSt

    override fun equals(other: Any?): Boolean {
        if (other !is FirRefWithValidityCheck<*>) return false
-        return getRefFir() == other.getRefFir() && this.token == other.token
+        return fir == other.fir && this.token == other.token
    }

    override fun hashCode(): Int {
-        return getRefFir().hashCode() * 31 + token.hashCode()
+        return fir.hashCode() * 31 + token.hashCode()
    }

-    private fun getRefFir(): FirDeclaration {
-        return firWeakRef.get() ?: throw EntityWasGarbageCollectedException("FirElement")
-    }
 }

@Suppress("NOTHING_TO_INLINE")