FIR: Add cache for expanded types

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

@@ -16,6 +16,10 @@ open class ConeClassLikeTypeImpl(
     isNullable: Boolean
 ) : ConeClassLikeType() {
     override val nullability: ConeNullability = ConeNullability.create(isNullable)
+
+    // Cached expanded type and the relevant session
+    var cachedExpandedType: Pair<*, ConeClassLikeType>? = null
+
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (javaClass != other?.javaClass) return false

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

@@ -63,9 +63,27 @@ fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): FirClassLikeSym
     return firSymbolProvider.getSymbolByLookupTag(this)
 }
 
-tailrec fun ConeClassLikeType.fullyExpandedType(
+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 computedExpandedType = fullyExpandedTypeNoCache(useSiteSession, expandedConeType)
+        cachedExpandedType = Pair(useSiteSession, computedExpandedType)
+        return computedExpandedType
+    }
+
+    return fullyExpandedTypeNoCache(useSiteSession, expandedConeType)
+}
+
+private fun ConeClassLikeType.fullyExpandedTypeNoCache(
+    useSiteSession: FirSession,
+    expandedConeType: (FirTypeAlias) -> ConeClassLikeType?
 ): ConeClassLikeType {
     val directExpansionType = directExpansionType(useSiteSession, expandedConeType) ?: return this
     return directExpansionType.fullyExpandedType(useSiteSession, expandedConeType)