FIR: Implement fast path for classes subtyping

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.resolve.impl.FirCompositeSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.impl.FirDependenciesSymbolProviderImpl
 import org.jetbrains.kotlin.fir.resolve.impl.FirLibrarySymbolProviderImpl
 import org.jetbrains.kotlin.fir.scopes.impl.FirClassDeclaredMemberScopeProvider
+import org.jetbrains.kotlin.fir.types.FirCorrespondingSupertypesCache
 import org.jetbrains.kotlin.load.java.JavaClassFinder
 import org.jetbrains.kotlin.load.java.JavaClassFinderImpl
 import org.jetbrains.kotlin.load.kotlin.KotlinClassFinder
@@ -32,6 +33,7 @@ class FirJavaModuleBasedSession(
 
     init {
         sessionProvider.sessionCache[moduleInfo] = this
+
         _firSymbolProvider = FirCompositeSymbolProvider(
             listOf(
                 service<FirProvider>(),
@@ -44,6 +46,12 @@ class FirJavaModuleBasedSession(
             FirSymbolProvider::class,
             _firSymbolProvider as FirSymbolProvider
         )
+
+        _correspondingSupertypesCache = FirCorrespondingSupertypesCache(this)
+        registerComponent(
+            FirCorrespondingSupertypesCache::class,
+            _correspondingSupertypesCache as FirCorrespondingSupertypesCache
+        )
     }
 }
 
@@ -81,6 +89,12 @@ class FirLibrarySession private constructor(
             _firSymbolProvider as FirSymbolProvider
         )
         registerComponent(FirClassDeclaredMemberScopeProvider::class, FirClassDeclaredMemberScopeProvider())
+
+        _correspondingSupertypesCache = FirCorrespondingSupertypesCache(this)
+        registerComponent(
+            FirCorrespondingSupertypesCache::class,
+            _correspondingSupertypesCache as FirCorrespondingSupertypesCache
+        )
     }
 
     companion object {

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)
+            }
+        }
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt

@@ -19,6 +19,7 @@ abstract class FirSession(val sessionProvider: FirSessionProvider?) {
 
     var _firSymbolProvider: Any? = null
 
+    var _correspondingSupertypesCache: Any? = null
 
     fun <T : Any> getService(kclass: KClass<T>): T =
         components[kclass] as T
@@ -34,4 +35,4 @@ interface FirSessionProvider {
 }
 
 inline fun <reified T : Any> FirSession.service(): T =
-    getService(T::class)
+    getService(T::class)

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/ConstraintInjector.kt

@@ -21,8 +21,10 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind.LOWER
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind.UPPER
 import org.jetbrains.kotlin.resolve.calls.model.KotlinCallDiagnostic
-import org.jetbrains.kotlin.types.*
-import org.jetbrains.kotlin.types.checker.NewCapturedType
+import org.jetbrains.kotlin.types.AbstractTypeApproximator
+import org.jetbrains.kotlin.types.AbstractTypeChecker
+import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
+import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.model.*
 import java.util.*
 import kotlin.math.max
@@ -123,7 +125,7 @@ class ConstraintInjector(val constraintIncorporator: ConstraintIncorporator, val
 
         val baseContext: AbstractTypeCheckerContext = newBaseTypeCheckerContext(isErrorTypeEqualsToAnything)
 
-        override fun substitutionSupertypePolicy(type: SimpleTypeMarker): SupertypesPolicy.DoCustomTransform {
+        override fun substitutionSupertypePolicy(type: SimpleTypeMarker): SupertypesPolicy {
             return baseContext.substitutionSupertypePolicy(type)
         }
 

core/type-system/src/org/jetbrains/kotlin/types/AbstractTypeChecker.kt

@@ -17,7 +17,7 @@ import java.util.*
 abstract class AbstractTypeCheckerContext : TypeSystemContext {
 
 
-    abstract fun substitutionSupertypePolicy(type: SimpleTypeMarker): SupertypesPolicy.DoCustomTransform
+    abstract fun substitutionSupertypePolicy(type: SimpleTypeMarker): SupertypesPolicy
 
     abstract fun areEqualTypeConstructors(a: TypeConstructorMarker, b: TypeConstructorMarker): Boolean
 
@@ -379,6 +379,10 @@ object AbstractTypeChecker {
         baseType: SimpleTypeMarker,
         constructor: TypeConstructorMarker
     ): List<SimpleTypeMarker> {
+        baseType.fastCorrespondingSupertypes(constructor)?.let {
+            return it
+        }
+
         if (constructor.isCommonFinalClassConstructor()) {
             return if (areEqualTypeConstructors(baseType.typeConstructor(), constructor))
                 listOf(captureFromArguments(baseType, CaptureStatus.FOR_SUBTYPING) ?: baseType)

core/type-system/src/org/jetbrains/kotlin/types/model/TypeSystemContext.kt

@@ -219,6 +219,8 @@ interface TypeSystemContext : TypeSystemOptimizationContext {
 
     fun SimpleTypeMarker.isClassType(): Boolean = typeConstructor().isClassTypeConstructor()
 
+    fun SimpleTypeMarker.fastCorrespondingSupertypes(constructor: TypeConstructorMarker): List<SimpleTypeMarker>? = null
+
     fun SimpleTypeMarker.isIntegerLiteralType(): Boolean = typeConstructor().isIntegerLiteralTypeConstructor()
 
     fun SimpleTypeMarker.possibleIntegerTypes(): Collection<KotlinTypeMarker>
@@ -290,4 +292,4 @@ inline fun TypeArgumentListMarker.all(
         if (!predicate(get(index))) return false
     }
     return true
-}
+}