[FIR] Cleanup caching symbol in ConeClassLikeLookupTagImpl

Type of `boundSymbol` replaced to OneElementWeakMap
`FirSymbolProvider.getSymbolByLookupTag` moved to extensions

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -10,21 +10,19 @@ import org.jetbrains.kotlin.builtins.functions.FunctionClassDescriptor
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.*
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
-import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.FirThisReference
-import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByLookupTag
 import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByTypeRef
 import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SamResolution.kt

@@ -15,8 +15,8 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.diagnostics.ConeIntermediateDiagnostic
-import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticFunctionSymbol
+import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByLookupTag
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
 import org.jetbrains.kotlin.fir.symbols.CallableId
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
+import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/FirSymbolProvider.kt

@@ -23,26 +23,9 @@ import org.jetbrains.kotlin.name.Name
 @RequiresOptIn
 annotation class FirSymbolProviderInternals
 
-abstract class FirSymbolProvider(protected val session: FirSession) : FirSessionComponent {
+abstract class FirSymbolProvider(val session: FirSession) : FirSessionComponent {
     abstract fun getClassLikeSymbolByFqName(classId: ClassId): FirClassLikeSymbol<*>?
 
-    fun getSymbolByLookupTag(lookupTag: ConeClassLikeLookupTag): FirClassLikeSymbol<*>? {
-        (lookupTag as? ConeClassLikeLookupTagImpl)
-            ?.boundSymbol?.takeIf { it.first === this }?.let { return it.second }
-
-        return getClassLikeSymbolByFqName(lookupTag.classId).also {
-            (lookupTag as? ConeClassLikeLookupTagImpl)?.bindSymbolToLookupTag(this, it)
-        }
-    }
-
-    fun getSymbolByLookupTag(lookupTag: ConeClassifierLookupTag): FirClassifierSymbol<*>? {
-        return when (lookupTag) {
-            is ConeClassLikeLookupTag -> getSymbolByLookupTag(lookupTag)
-            is ConeClassifierLookupTagWithFixedSymbol -> lookupTag.symbol
-            else -> error("Unknown lookupTag type: ${lookupTag::class}")
-        }
-    }
-
     @OptIn(ExperimentalStdlibApi::class, FirSymbolProviderInternals::class)
     open fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
         return buildList { getTopLevelCallableSymbolsTo(this, packageFqName, name) }
@@ -70,6 +53,26 @@ inline fun <reified T : AbstractFirBasedSymbol<*>> FirSymbolProvider.getSymbolBy
     return getSymbolByLookupTag(lookupTag) as? T
 }
 
-fun ConeClassLikeLookupTagImpl.bindSymbolToLookupTag(provider: FirSymbolProvider, symbol: FirClassLikeSymbol<*>?) {
+@OptIn(LookupTagInternals::class)
-    boundSymbol = Pair(provider, symbol)
+fun ConeClassLikeLookupTagImpl.bindSymbolToLookupTag(session: FirSession, symbol: FirClassLikeSymbol<*>?) {
+    boundSymbol = OneElementWeakMap(session, symbol)
+}
+
+fun FirSymbolProvider.getSymbolByLookupTag(lookupTag: ConeClassifierLookupTag): FirClassifierSymbol<*>? {
+    return when (lookupTag) {
+        is ConeClassLikeLookupTag -> getSymbolByLookupTag(lookupTag)
+        is ConeClassifierLookupTagWithFixedSymbol -> lookupTag.symbol
+        else -> error("Unknown lookupTag type: ${lookupTag::class}")
+    }
+}
+
+@OptIn(LookupTagInternals::class)
+fun FirSymbolProvider.getSymbolByLookupTag(lookupTag: ConeClassLikeLookupTag): FirClassLikeSymbol<*>? {
+    (lookupTag as? ConeClassLikeLookupTagImpl)
+        ?.boundSymbol?.takeIf { it.key === this.session }?.let { return it.value }
+
+    return getClassLikeSymbolByFqName(lookupTag.classId)
+        .also {
+            (lookupTag as? ConeClassLikeLookupTagImpl)?.bindSymbolToLookupTag(session, it)
+        }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirTypeResolverImpl.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeWrongNumberOfTypeArgumentsError
 import org.jetbrains.kotlin.fir.resolve.providers.bindSymbolToLookupTag
+import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByLookupTag
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
@@ -136,7 +137,7 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver {
             .also {
                 val lookupTag = it.lookupTag
                 if (lookupTag is ConeClassLikeLookupTagImpl && symbol is FirClassLikeSymbol<*>) {
-                    lookupTag.bindSymbolToLookupTag(session.firSymbolProvider, symbol)
+                    lookupTag.bindSymbolToLookupTag(session, symbol)
                 }
             }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSealedClassInheritorsTransformer.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
@@ -99,4 +100,4 @@ class FirSealedClassInheritorsTransformer : FirTransformer<Nothing?>() {
 
 object SealedClassInheritorsKey : FirDeclarationDataKey()
 
-var FirRegularClass.sealedInheritors: List<ClassId>? by FirDeclarationDataRegistry.data(SealedClassInheritorsKey)
+var FirRegularClass.sealedInheritors: List<ClassId>? by FirDeclarationDataRegistry.data(SealedClassInheritorsKey)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirWhenExhaustivenessTransformer.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByLookupTag
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol

compiler/fir/tree/src/org/jetbrains/kotlin/fir/symbols/impl/ConeClassLikeLookupTagImpl.kt

@@ -5,11 +5,16 @@
 
 package org.jetbrains.kotlin.fir.symbols.impl
 
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.name.ClassId
 
+@RequiresOptIn
+annotation class LookupTagInternals
+
 class ConeClassLikeLookupTagImpl(override val classId: ClassId) : ConeClassLikeLookupTag() {
-    var boundSymbol: Pair<*, FirClassLikeSymbol<*>?>? = null
+    @LookupTagInternals
+    var boundSymbol: OneElementWeakMap<FirSession, FirClassLikeSymbol<*>?>? = null
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
@@ -25,4 +30,4 @@ class ConeClassLikeLookupTagImpl(override val classId: ClassId) : ConeClassLikeL
     override fun hashCode(): Int {
         return classId.hashCode()
     }
-}
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/symbols/impl/OneElementWeakMap.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.fir.symbols.impl
+
+import java.lang.ref.WeakReference
+
+class OneElementWeakMap<K, V> private constructor(
+    private val keyReference: WeakReference<K>,
+    private val valueReference: WeakReference<V>
+) {
+    constructor(key: K, value: V) : this(WeakReference(key), WeakReference(value))
+
+    val key: K?
+        get() = keyReference.get()
+
+    val value: V?
+        get() = valueReference.get()
+}

compiler/visualizer/render-fir/src/org/jetbrains/kotlin/compiler/visualizer/FirVisualizer.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.directExpansionType
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByLookupTag
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.firUnsafe
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
@@ -19,19 +20,8 @@ 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.*
-import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.types.*
-import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirLocalVariableSymbol
-import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirPropertySymbol
-import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirFunctionSymbol
-import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirFunctionValueParameterSymbol
-import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirClassType
-import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirErrorType
-import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirFlexibleType
-import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirIntersectionType
-import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirTypeArgumentWithVariance
-import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirTypeParameterType
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.threadLocal
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
@@ -221,4 +211,4 @@ internal fun FirElement.buildSymbol(builder: KtSymbolByFirBuilder) =
     (this as? FirDeclaration)?.let(builder::buildSymbol)
 
 internal fun FirDeclaration.buildSymbol(builder: KtSymbolByFirBuilder) =
-    builder.buildSymbol(this)
+    builder.buildSymbol(this)