FIR: Implement fast path for classes subtyping

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve
 
-import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.render
@@ -33,6 +32,8 @@ inline fun <K, V, VA : V> MutableMap<K, V>.getOrPut(key: K, defaultValue: (K) ->
 }
 
 val FirSession.firSymbolProvider: FirSymbolProvider get() = _firSymbolProvider as FirSymbolProvider? ?: service()
+val FirSession.correspondingSupertypesCache: FirCorrespondingSupertypesCache
+    get() = _correspondingSupertypesCache as FirCorrespondingSupertypesCache? ?: service()
 
 fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): ConeClassifierSymbol? {
     val firSymbolProvider = useSiteSession.firSymbolProvider

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeTypeContext.kt

@@ -11,14 +11,11 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.expandedConeType
 import org.jetbrains.kotlin.fir.declarations.superConeTypes
+import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.ConeTypeVariableTypeConstructor
-import org.jetbrains.kotlin.fir.resolve.constructType
-import org.jetbrains.kotlin.fir.resolve.directExpansionType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
-import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.resolve.transformers.firUnsafe
-import org.jetbrains.kotlin.fir.resolve.withArguments
 import org.jetbrains.kotlin.fir.service
 import org.jetbrains.kotlin.fir.symbols.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
@@ -46,6 +43,11 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
         TODO("not implemented")
     }
 
+    override fun SimpleTypeMarker.fastCorrespondingSupertypes(constructor: TypeConstructorMarker): List<SimpleTypeMarker>? {
+        require(this is ConeKotlinType)
+        return session.correspondingSupertypesCache.getCorrespondingSupertypes(this, constructor)
+    }
+
     override fun SimpleTypeMarker.isIntegerLiteralType(): Boolean {
         return false
     }
@@ -391,7 +393,8 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
 
 class ConeTypeCheckerContext(override val isErrorTypeEqualsToAnything: Boolean, override val session: FirSession) :
     AbstractTypeCheckerContext(), ConeTypeContext {
-    override fun substitutionSupertypePolicy(type: SimpleTypeMarker): SupertypesPolicy.DoCustomTransform {
+    override fun substitutionSupertypePolicy(type: SimpleTypeMarker): SupertypesPolicy {
+        if (type.argumentsCount() == 0) return SupertypesPolicy.LowerIfFlexible
         require(type is ConeKotlinType)
         val declaration = when (type) {
             is ConeClassType -> type.lookupTag.toSymbol(session)?.firUnsafe<FirRegularClass>()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/FirCorrespondingSupertypesCache.kt

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2010-2019 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.types
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.toFirClassLike
+import org.jetbrains.kotlin.fir.resolve.constructClassType
+import org.jetbrains.kotlin.fir.resolve.constructType
+import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.symbols.ConeClassLikeSymbol
+import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
+import org.jetbrains.kotlin.types.model.CaptureStatus
+import org.jetbrains.kotlin.types.model.SimpleTypeMarker
+import org.jetbrains.kotlin.types.model.TypeConstructorMarker
+
+class FirCorrespondingSupertypesCache(private val session: FirSession) {
+    private val context = ConeTypeCheckerContext(false, session)
+    private val cache = HashMap<ConeClassLikeSymbol, Map<ConeClassLikeSymbol, List<ConeClassLikeType>>?>(1000, 0.5f)
+
+    fun getCorrespondingSupertypes(
+        type: ConeKotlinType,
+        supertypeConstructor: TypeConstructorMarker
+    ): List<ConeClassLikeType>? {
+        if (type !is ConeClassLikeType || supertypeConstructor !is ConeClassLikeSymbol) return null
+
+        val symbol = type.lookupTag.toSymbol(session) as? ConeClassLikeSymbol ?: return null
+        if (symbol == supertypeConstructor) return listOf(captureType(type))
+
+        if (symbol !in cache) {
+            cache[symbol] = computeSupertypesMap(type, symbol)
+        }
+
+        val resultTypes = cache[symbol]?.getOrDefault(supertypeConstructor, emptyList()) ?: return null
+        if (type.typeArguments.isEmpty()) return resultTypes
+
+        val capturedType = captureType(type)
+        val substitutionSupertypePolicy = context.substitutionSupertypePolicy(capturedType)
+        return resultTypes.map {
+            substitutionSupertypePolicy.transformType(context, it) as ConeClassLikeType
+        }
+    }
+
+    private fun captureType(type: ConeClassLikeType): ConeClassLikeType =
+        (context.captureFromArguments(type, CaptureStatus.FOR_SUBTYPING) ?: type) as ConeClassLikeType
+
+    private fun computeSupertypesMap(
+        subtype: ConeLookupTagBasedType,
+        subtypeSymbol: ConeClassLikeSymbol
+    ): Map<ConeClassLikeSymbol, List<ConeClassLikeType>>? {
+        val resultingMap = HashMap<ConeClassLikeSymbol, List<ConeClassLikeType>>()
+
+        val subtypeClassSymbol = with(context) {
+            subtype.typeConstructor() as? ConeClassLikeSymbol ?: return null
+        }
+        val subtypeFirClass = subtypeClassSymbol.toFirClassLike() ?: return null
+
+        val defaultType = subtypeClassSymbol.toLookupTag().constructClassType(
+            subtypeFirClass.typeParameters.map {
+                it.symbol.toLookupTag().constructType(emptyArray(), isNullable = false)
+            }.toTypedArray(),
+            isNullable = false
+        )
+
+        if (context.anySupertype(
+                defaultType,
+                { it !is ConeClassLikeType || it.lookupTag.toSymbol(session) !is ConeClassLikeSymbol }
+            ) { supertype -> computeSupertypePolicyAndPutInMap(supertype, subtypeSymbol, resultingMap) }
+        ) {
+            return null
+        }
+
+        return resultingMap
+    }
+
+    private fun computeSupertypePolicyAndPutInMap(
+        supertype: SimpleTypeMarker,
+        subtypeSymbol: ConeClassLikeSymbol,
+        resultingMap: MutableMap<ConeClassLikeSymbol, List<ConeClassLikeType>>
+    ): AbstractTypeCheckerContext.SupertypesPolicy {
+        val supertypeSymbol = (supertype as ConeClassLikeType).lookupTag.toSymbol(session) as ConeClassLikeSymbol
+        val captured = context.captureFromArguments(supertype, CaptureStatus.FOR_SUBTYPING) as ConeClassLikeType? ?: supertype
+
+        if (supertypeSymbol != subtypeSymbol) {
+            resultingMap[supertypeSymbol] = listOf(captured)
+        }
+
+        return when {
+            with(context) { captured.argumentsCount() } == 0 -> {
+                AbstractTypeCheckerContext.SupertypesPolicy.LowerIfFlexible
+            }
+            else -> {
+                context.substitutionSupertypePolicy(captured)
+            }
+        }
+    }
+}