[FIR] Store cachedExpandedType in WeakPair instead of usual Pair

WeakPair is a renamed OneElementWeakMap

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/impl/Impl.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.types.ConeAttributes
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.ConeTypeProjection
 import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.utils.WeakPair
 
 class ConeClassLikeTypeImpl(
     override val lookupTag: ConeClassLikeLookupTag,
@@ -20,7 +21,7 @@ class ConeClassLikeTypeImpl(
     override val nullability: ConeNullability = ConeNullability.create(isNullable)
 
     // Cached expanded type and the relevant session
-    var cachedExpandedType: Pair<*, ConeClassLikeType>? = null
+    var cachedExpandedType: WeakPair<*, ConeClassLikeType>? = null
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true

compiler/fir/cones/src/org/jetbrains/kotlin/fir/utils/WeakPair.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2021 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.utils
+
+import java.lang.ref.WeakReference
+
+class WeakPair<K, V>(first: K, second: V) {
+    private val firstReference: WeakReference<K> = WeakReference(first)
+    private val secondReference: WeakReference<V> = WeakReference(second)
+
+    val first: K?
+        get() = firstReference.get()
+
+    val second: V?
+        get() = secondReference.get()
+}
+
+operator fun <K, V> WeakPair<K, V>?.component1(): K? {
+    return this?.first
+}
+
+operator fun <K, V> WeakPair<K, V>?.component2(): V? {
+    return this?.second
+}

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTagWithFixedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.utils.WeakPair
 
 fun ConeClassifierLookupTag.toSymbol(useSiteSession: FirSession): FirClassifierSymbol<*>? =
     when (this) {
@@ -26,7 +27,7 @@ fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): FirClassLikeSym
         return this.symbol
     }
     val firSymbolProvider = useSiteSession.symbolProvider
-    (this as? ConeClassLikeLookupTagImpl)?.boundSymbol?.takeIf { it.key === useSiteSession }?.let { return it.value }
+    (this as? ConeClassLikeLookupTagImpl)?.boundSymbol?.takeIf { it.first === useSiteSession }?.let { return it.second }
 
     return firSymbolProvider.getClassLikeSymbolByFqName(classId).also {
         (this as? ConeClassLikeLookupTagImpl)?.bindSymbolToLookupTag(useSiteSession, it)
@@ -44,7 +45,7 @@ fun ConeClassLikeLookupTag.toFirRegularClass(session: FirSession): FirRegularCla
 
 @OptIn(LookupTagInternals::class)
 fun ConeClassLikeLookupTagImpl.bindSymbolToLookupTag(session: FirSession, symbol: FirClassLikeSymbol<*>?) {
-    boundSymbol = OneElementWeakMap(session, symbol)
+    boundSymbol = WeakPair(session, symbol)
 }
 
 fun FirSymbolProvider.getSymbolByLookupTag(lookupTag: ConeClassifierLookupTag): FirClassifierSymbol<*>? {

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

@@ -12,19 +12,22 @@ import org.jetbrains.kotlin.fir.resolve.substitution.AbstractConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.utils.WeakPair
+import org.jetbrains.kotlin.fir.utils.component1
+import org.jetbrains.kotlin.fir.utils.component2
 
 fun ConeClassLikeType.fullyExpandedType(
     useSiteSession: FirSession,
     expandedConeType: (FirTypeAlias) -> ConeClassLikeType? = FirTypeAlias::expandedConeType,
 ): ConeClassLikeType {
     if (this is ConeClassLikeTypeImpl) {
-        val expandedTypeAndSession = cachedExpandedType
-        if (expandedTypeAndSession != null && expandedTypeAndSession.first === useSiteSession) {
-            return expandedTypeAndSession.second
+        val (cachedSession, cachedExpandedType) = cachedExpandedType
+        if (cachedSession === useSiteSession && cachedExpandedType != null) {
+            return cachedExpandedType
         }
 
         val computedExpandedType = fullyExpandedTypeNoCache(useSiteSession, expandedConeType)
-        cachedExpandedType = Pair(useSiteSession, computedExpandedType)
+        this.cachedExpandedType = WeakPair(useSiteSession, computedExpandedType)
         return computedExpandedType
     }
 

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.symbols.impl
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
+import org.jetbrains.kotlin.fir.utils.WeakPair
 import org.jetbrains.kotlin.name.ClassId
 
 @RequiresOptIn
@@ -19,7 +20,7 @@ class ConeClassLikeLookupTagImpl(override val classId: ClassId) : ConeClassLikeL
     }
 
     @LookupTagInternals
-    var boundSymbol: OneElementWeakMap<FirSession, FirClassLikeSymbol<*>?>? = null
+    var boundSymbol: WeakPair<FirSession, FirClassLikeSymbol<*>?>? = null
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true

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

@@ -1,21 +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.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()
-}