Extract type preparation during type checking to a separate component

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

@@ -393,22 +393,6 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
         return ConeTypeIntersector.intersectTypes(this as ConeInferenceContext, types as List<ConeKotlinType>)
     }
 
-    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
-        return when (type) {
-            is ConeClassLikeType -> type.fullyExpandedType(session)
-            is ConeFlexibleType -> {
-                val lowerBound = prepareType(type.lowerBound)
-                if (lowerBound === type.lowerBound) return type
-
-                ConeFlexibleType(
-                    lowerBound as ConeKotlinType,
-                    prepareType(type.upperBound) as ConeKotlinType
-                )
-            }
-            else -> type
-        }
-    }
-
     override fun KotlinTypeMarker.isNullableType(): Boolean {
         require(this is ConeKotlinType)
         if (this.isMarkedNullable)
@@ -551,7 +535,8 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
 class ConeTypeCheckerContext(
     override val isErrorTypeEqualsToAnything: Boolean,
     override val isStubTypeEqualsToAnything: Boolean,
-    override val typeSystemContext: ConeInferenceContext
+    override val typeSystemContext: ConeInferenceContext,
+    val kotlinTypePreparator: ConeTypePreparator = ConeTypePreparator.getDefault(typeSystemContext.session),
 ) : AbstractTypeCheckerContext() {
 
     val session: FirSession = typeSystemContext.session
@@ -585,7 +570,11 @@ class ConeTypeCheckerContext(
     }
 
     override fun refineType(type: KotlinTypeMarker): KotlinTypeMarker {
-        return typeSystemContext.prepareType(type)
+        return kotlinTypePreparator.prepareType(type)
+    }
+
+    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
+        return kotlinTypePreparator.prepareType(type)
     }
 
     override val KotlinTypeMarker.isAllowedTypeVariable: Boolean

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

@@ -0,0 +1,33 @@
+/*
+ * 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.types
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
+import org.jetbrains.kotlin.types.AbstractTypePreparator
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+
+abstract class ConeTypePreparator(val session: FirSession) : AbstractTypePreparator {
+    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
+        return when (type) {
+            is ConeClassLikeType -> type.fullyExpandedType(session)
+            is ConeFlexibleType -> {
+                val lowerBound = prepareType(type.lowerBound)
+                if (lowerBound === type.lowerBound) return type
+
+                ConeFlexibleType(
+                    lowerBound as ConeKotlinType,
+                    prepareType(type.upperBound) as ConeKotlinType
+                )
+            }
+            else -> type
+        }
+    }
+
+    companion object {
+        fun getDefault(session: FirSession) = object : ConeTypePreparator(session) {}
+    }
+}

compiler/frontend/src/org/jetbrains/kotlin/frontend/di/injection.kt

@@ -41,6 +41,7 @@ import org.jetbrains.kotlin.resolve.checkers.ExperimentalUsageChecker
 import org.jetbrains.kotlin.resolve.lazy.*
 import org.jetbrains.kotlin.resolve.lazy.declarations.DeclarationProviderFactory
 import org.jetbrains.kotlin.types.KotlinTypeRefinerImpl
+import org.jetbrains.kotlin.types.checker.KotlinTypePreparator
 import org.jetbrains.kotlin.types.checker.KotlinTypeRefiner
 import org.jetbrains.kotlin.types.checker.NewKotlinTypeCheckerImpl
 import org.jetbrains.kotlin.types.expressions.DeclarationScopeProviderForLocalClassifierAnalyzer
@@ -89,6 +90,8 @@ fun StorageComponentContainer.configureModule(
         useInstance(KotlinTypeRefiner.Default)
     }
 
+    useInstance(KotlinTypePreparator.Default)
+
     configurePlatformIndependentComponents()
 }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/types/IrTypeSystemContext.kt

@@ -519,10 +519,6 @@ interface IrTypeSystemContext : TypeSystemContext, TypeSystemCommonSuperTypesCon
     override fun SimpleTypeMarker.makeSimpleTypeDefinitelyNotNullOrNotNull(): SimpleTypeMarker {
         error("makeSimpleTypeDefinitelyNotNullOrNotNull is not yet supported in IR")
     }
-
-    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
-        return type
-    }
 }
 
 fun extractTypeParameters(parent: IrDeclarationParent): List<IrTypeParameter> {

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

@@ -257,6 +257,8 @@ class ConstraintInjector(
             }
         }
 
+        override fun prepareType(type: KotlinTypeMarker) = baseContext.prepareType(type)
+
         fun runIsSubtypeOf(
             lowerType: KotlinTypeMarker,
             upperType: KotlinTypeMarker,

compiler/resolution.common/src/org/jetbrains/kotlin/resolve/calls/results/FlatSignature.kt

@@ -74,8 +74,8 @@ fun <T> SimpleConstraintSystem.isSignatureNotLessSpecific(
              * Here, when we try solve this CS(Y is variables) then Array<out X> <: Array<out Y> and this system impossible to solve,
              * so we capture types from receiver and value parameters.
              */
-            val specificCapturedType =
+            val specificCapturedType = AbstractTypeChecker.prepareType(context, specificType)
-                context.prepareType(specificType).let { if (captureFromArgument) context.captureFromExpression(it) ?: it else it }
+                .let { if (captureFromArgument) context.captureFromExpression(it) ?: it else it }
             addSubtypeConstraint(specificCapturedType, substitutedGeneralType)
         }
     }

compiler/resolution.common/src/org/jetbrains/kotlin/types/AbstractTypeApproximator.kt

@@ -86,7 +86,7 @@ abstract class AbstractTypeApproximator(
 
         return cachedValue(type, conf, toSuper = true) {
             approximateTo(
-                prepareType(type), conf, { upperBound() },
+                AbstractTypeChecker.prepareType(ctx, type), conf, { upperBound() },
                 referenceApproximateToSuperType, depth
             )
         }
@@ -97,7 +97,7 @@ abstract class AbstractTypeApproximator(
 
         return cachedValue(type, conf, toSuper = false) {
             approximateTo(
-                prepareType(type), conf, { lowerBound() },
+                AbstractTypeChecker.prepareType(ctx, type), conf, { lowerBound() },
                 referenceApproximateToSubType, depth
             )
         }

core/compiler.common/src/org/jetbrains/kotlin/types/AbstractTypeChecker.kt

@@ -31,6 +31,10 @@ abstract class AbstractTypeCheckerContext() {
         return type
     }
 
+    open fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
+        return type
+    }
+
     open fun customIsSubtypeOf(subType: KotlinTypeMarker, superType: KotlinTypeMarker): Boolean = true
 
     abstract val isErrorTypeEqualsToAnything: Boolean
@@ -160,13 +164,19 @@ object AbstractTypeChecker {
     @JvmField
     var RUN_SLOW_ASSERTIONS = false
 
+    fun prepareType(
+        context: TypeCheckerProviderContext,
+        type: KotlinTypeMarker,
+        stubTypesEqualToAnything: Boolean = true
+    ) = context.newBaseTypeCheckerContext(true, stubTypesEqualToAnything).prepareType(type)
+
     fun isSubtypeOf(
         context: TypeCheckerProviderContext,
         subType: KotlinTypeMarker,
         superType: KotlinTypeMarker,
         stubTypesEqualToAnything: Boolean = true
     ): Boolean {
-        return AbstractTypeChecker.isSubtypeOf(context.newBaseTypeCheckerContext(true, stubTypesEqualToAnything), subType, superType)
+        return isSubtypeOf(context.newBaseTypeCheckerContext(true, stubTypesEqualToAnything), subType, superType)
     }
 
     fun isSubtypeOfClass(
@@ -204,12 +214,7 @@ object AbstractTypeChecker {
 
         if (!context.customIsSubtypeOf(subType, superType)) return false
 
-        return completeIsSubTypeOf(
+        return completeIsSubTypeOf(context, subType, superType, isFromNullabilityConstraint)
-            context,
-            context.typeSystemContext.prepareType(context.refineType(subType)),
-            context.typeSystemContext.prepareType(context.refineType(superType)),
-            isFromNullabilityConstraint
-        )
     }
 
     fun equalTypes(context: AbstractTypeCheckerContext, a: KotlinTypeMarker, b: KotlinTypeMarker): Boolean =
@@ -238,15 +243,18 @@ object AbstractTypeChecker {
         superType: KotlinTypeMarker,
         isFromNullabilityConstraint: Boolean
     ): Boolean = with(context.typeSystemContext) {
-        checkSubtypeForSpecialCases(context, subType.lowerBoundIfFlexible(), superType.upperBoundIfFlexible())?.let {
+        val preparedSubType = context.prepareType(context.refineType(subType))
-            context.addSubtypeConstraint(subType, superType, isFromNullabilityConstraint)
+        val preparedSuperType = context.prepareType(context.refineType(superType))
+
+        checkSubtypeForSpecialCases(context, preparedSubType.lowerBoundIfFlexible(), preparedSuperType.upperBoundIfFlexible())?.let {
+            context.addSubtypeConstraint(preparedSubType, preparedSuperType, isFromNullabilityConstraint)
             return it
         }
 
         // we should add constraints with flexible types, otherwise we never get flexible type as answer in constraint system
-        context.addSubtypeConstraint(subType, superType, isFromNullabilityConstraint)?.let { return it }
+        context.addSubtypeConstraint(preparedSubType, preparedSuperType, isFromNullabilityConstraint)?.let { return it }
 
-        return isSubtypeOfForSingleClassifierType(context, subType.lowerBoundIfFlexible(), superType.upperBoundIfFlexible())
+        return isSubtypeOfForSingleClassifierType(context, preparedSubType.lowerBoundIfFlexible(), preparedSuperType.upperBoundIfFlexible())
     }
 
     private fun checkSubtypeForIntegerLiteralType(
@@ -332,6 +340,7 @@ object AbstractTypeChecker {
         if (superType.typeConstructor().isAnyConstructor()) return true
 
         val supertypesWithSameConstructor = findCorrespondingSupertypes(context, subType, superConstructor)
+            .map { context.prepareType(it).asSimpleType() ?: it }
         when (supertypesWithSameConstructor.size) {
             0 -> return hasNothingSupertype(context, subType) // todo Nothing & Array<Number> <: Array<String>
             1 -> return context.isSubtypeForSameConstructor(supertypesWithSameConstructor.first().asArgumentList(), superType)

core/compiler.common/src/org/jetbrains/kotlin/types/AbstractTypePreparator.kt

@@ -0,0 +1,12 @@
+/*
+ * 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.types
+
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+
+interface AbstractTypePreparator {
+    fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker
+}

core/compiler.common/src/org/jetbrains/kotlin/types/model/TypeSystemContext.kt

@@ -412,8 +412,6 @@ interface TypeSystemContext : TypeSystemOptimizationContext {
 
     fun KotlinTypeMarker.isSimpleType(): Boolean = asSimpleType() != null
 
-    fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker
-
     fun SimpleTypeMarker.isPrimitiveType(): Boolean
 
     fun KotlinTypeMarker.getAnnotations(): List<AnnotationMarker>

core/descriptors/src/org/jetbrains/kotlin/resolve/OverridingUtil.java

@@ -392,7 +392,7 @@ public class OverridingUtil {
         assert firstParameters.size() == secondParameters.size() :
                 "Should be the same number of type parameters: " + firstParameters + " vs " + secondParameters;
 
-        NewKotlinTypeCheckerImpl typeChecker = new NewKotlinTypeCheckerImpl(kotlinTypeRefiner);
+        NewKotlinTypeCheckerImpl typeChecker = new NewKotlinTypeCheckerImpl(kotlinTypeRefiner, KotlinTypePreparator.Default.INSTANCE);
         ClassicTypeCheckerContext context = createTypeCheckerContext(firstParameters, secondParameters);
 
         return new Pair<NewKotlinTypeCheckerImpl, ClassicTypeCheckerContext>(typeChecker, context);

core/descriptors/src/org/jetbrains/kotlin/resolve/OverridingUtilTypeSystemContext.kt

@@ -34,7 +34,7 @@ class OverridingUtilTypeSystemContext(
             stubTypesEqualToAnything,
             allowedTypeVariable = true,
             kotlinTypeRefiner,
-            this
+            typeSystemContext = this
         )
     }
 

core/descriptors/src/org/jetbrains/kotlin/types/checker/ClassicTypeCheckerContext.kt

@@ -26,6 +26,7 @@ open class ClassicTypeCheckerContext(
     val stubTypeEqualsToAnything: Boolean = true,
     val allowedTypeVariable: Boolean = true,
     val kotlinTypeRefiner: KotlinTypeRefiner = KotlinTypeRefiner.Default,
+    val kotlinTypePreparator: KotlinTypePreparator = KotlinTypePreparator.Default,
     override val typeSystemContext: ClassicTypeSystemContext = SimpleClassicTypeSystemContext
 ) : AbstractTypeCheckerContext() {
 
@@ -35,6 +36,11 @@ open class ClassicTypeCheckerContext(
         return kotlinTypeRefiner.refineType(type)
     }
 
+    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
+        require(type is KotlinType, type::errorMessage)
+        return kotlinTypePreparator.prepareType(type.unwrap())
+    }
+
     override val isErrorTypeEqualsToAnything: Boolean
         get() = errorTypeEqualsToAnything
 

core/descriptors/src/org/jetbrains/kotlin/types/checker/ClassicTypeSystemContext.kt

@@ -507,11 +507,6 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
         return this.replace(newArguments as List<TypeProjection>)
     }
 
-    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
-        require(type is KotlinType, type::errorMessage)
-        return NewKotlinTypeChecker.Default.transformToNewType(type.unwrap())
-    }
-
     override fun DefinitelyNotNullTypeMarker.original(): SimpleTypeMarker {
         require(this is DefinitelyNotNullType, this::errorMessage)
         return this.original

core/descriptors/src/org/jetbrains/kotlin/types/checker/KotlinTypePreparator.kt

@@ -0,0 +1,75 @@
+/*
+ * 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.types.checker
+
+import org.jetbrains.kotlin.container.DefaultImplementation
+import org.jetbrains.kotlin.resolve.calls.inference.CapturedTypeConstructorImpl
+import org.jetbrains.kotlin.resolve.constants.IntegerValueTypeConstructor
+import org.jetbrains.kotlin.types.*
+import org.jetbrains.kotlin.types.model.CaptureStatus
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.typeUtil.makeNullable
+
+@DefaultImplementation(impl = KotlinTypePreparator.Default::class)
+abstract class KotlinTypePreparator : AbstractTypePreparator {
+    private fun transformToNewType(type: SimpleType): SimpleType {
+        when (val constructor = type.constructor) {
+            // Type itself can be just SimpleTypeImpl, not CapturedType. see KT-16147
+            is CapturedTypeConstructorImpl -> {
+                val lowerType = constructor.projection.takeIf { it.projectionKind == Variance.IN_VARIANCE }?.type?.unwrap()
+
+                // it is incorrect calculate this type directly because of recursive star projections
+                if (constructor.newTypeConstructor == null) {
+                    constructor.newTypeConstructor =
+                        NewCapturedTypeConstructor(constructor.projection, constructor.supertypes.map { it.unwrap() })
+                }
+                return NewCapturedType(
+                    CaptureStatus.FOR_SUBTYPING, constructor.newTypeConstructor!!,
+                    lowerType, type.annotations, type.isMarkedNullable
+                )
+            }
+
+            is IntegerValueTypeConstructor -> {
+                val newConstructor =
+                    IntersectionTypeConstructor(constructor.supertypes.map { TypeUtils.makeNullableAsSpecified(it, type.isMarkedNullable) })
+                return KotlinTypeFactory.simpleTypeWithNonTrivialMemberScope(
+                    type.annotations,
+                    newConstructor,
+                    listOf(),
+                    false,
+                    type.memberScope
+                )
+            }
+
+            is IntersectionTypeConstructor -> if (type.isMarkedNullable) {
+                val newConstructor = constructor.transformComponents(transform = { it.makeNullable() }) ?: constructor
+                return newConstructor.createType()
+
+            }
+        }
+
+        return type
+    }
+
+    override fun prepareType(type: KotlinTypeMarker): UnwrappedType {
+        require(type is KotlinType)
+        val unwrappedType = type.unwrap()
+        return when (unwrappedType) {
+            is SimpleType -> transformToNewType(unwrappedType)
+            is FlexibleType -> {
+                val newLower = transformToNewType(unwrappedType.lowerBound)
+                val newUpper = transformToNewType(unwrappedType.upperBound)
+                if (newLower !== unwrappedType.lowerBound || newUpper !== unwrappedType.upperBound) {
+                    KotlinTypeFactory.flexibleType(newLower, newUpper)
+                } else {
+                    unwrappedType
+                }
+            }
+        }.inheritEnhancement(unwrappedType)
+    }
+
+    object Default : KotlinTypePreparator()
+}

core/descriptors/src/org/jetbrains/kotlin/types/checker/NewCapturedType.kt

@@ -173,11 +173,11 @@ private fun captureArguments(type: UnwrappedType, status: CaptureStatus): List<T
 
         if (oldProjection.projectionKind == Variance.INVARIANT) continue
         val capturedTypeSupertypes = type.constructor.parameters[index].upperBounds.mapTo(mutableListOf()) {
-            NewKotlinTypeChecker.Default.transformToNewType(substitutor.safeSubstitute(it, Variance.INVARIANT).unwrap())
+            KotlinTypePreparator.Default.prepareType(substitutor.safeSubstitute(it, Variance.INVARIANT).unwrap())
         }
 
         if (!oldProjection.isStarProjection && oldProjection.projectionKind == Variance.OUT_VARIANCE) {
-            capturedTypeSupertypes += NewKotlinTypeChecker.Default.transformToNewType(oldProjection.type.unwrap())
+            capturedTypeSupertypes += KotlinTypePreparator.Default.prepareType(oldProjection.type.unwrap())
         }
 
         val capturedType = newProjection.type as NewCapturedType

core/descriptors/src/org/jetbrains/kotlin/types/checker/NewKotlinTypeChecker.kt

@@ -61,17 +61,19 @@ object ErrorTypesAreEqualToAnything : KotlinTypeChecker {
 
 interface NewKotlinTypeChecker : KotlinTypeChecker {
     val kotlinTypeRefiner: KotlinTypeRefiner
+    val kotlinTypePreparator: KotlinTypePreparator
     val overridingUtil: OverridingUtil
 
-    fun transformToNewType(type: UnwrappedType): UnwrappedType
-
     companion object {
         val Default = NewKotlinTypeCheckerImpl(KotlinTypeRefiner.Default)
     }
 }
 
 
-class NewKotlinTypeCheckerImpl(override val kotlinTypeRefiner: KotlinTypeRefiner) : NewKotlinTypeChecker {
+class NewKotlinTypeCheckerImpl(
+    override val kotlinTypeRefiner: KotlinTypeRefiner,
+    override val kotlinTypePreparator: KotlinTypePreparator = KotlinTypePreparator.Default
+) : NewKotlinTypeChecker {
     override val overridingUtil: OverridingUtil = OverridingUtil.createWithTypeRefiner(kotlinTypeRefiner)
 
     override fun isSubtypeOf(subtype: KotlinType, supertype: KotlinType): Boolean =
@@ -88,60 +90,6 @@ class NewKotlinTypeCheckerImpl(override val kotlinTypeRefiner: KotlinTypeRefiner
     fun ClassicTypeCheckerContext.isSubtypeOf(subType: UnwrappedType, superType: UnwrappedType): Boolean {
         return AbstractTypeChecker.isSubtypeOf(this as AbstractTypeCheckerContext, subType, superType)
     }
-
-    fun transformToNewType(type: SimpleType): SimpleType {
-        val constructor = type.constructor
-        when (constructor) {
-            // Type itself can be just SimpleTypeImpl, not CapturedType. see KT-16147
-            is CapturedTypeConstructorImpl -> {
-                val lowerType = constructor.projection.takeIf { it.projectionKind == Variance.IN_VARIANCE }?.type?.unwrap()
-
-                // it is incorrect calculate this type directly because of recursive star projections
-                if (constructor.newTypeConstructor == null) {
-                    constructor.newTypeConstructor =
-                        NewCapturedTypeConstructor(constructor.projection, constructor.supertypes.map { it.unwrap() })
-                }
-                return NewCapturedType(
-                    CaptureStatus.FOR_SUBTYPING, constructor.newTypeConstructor!!,
-                    lowerType, type.annotations, type.isMarkedNullable
-                )
-            }
-
-            is IntegerValueTypeConstructor -> {
-                val newConstructor =
-                    IntersectionTypeConstructor(constructor.supertypes.map { TypeUtils.makeNullableAsSpecified(it, type.isMarkedNullable) })
-                return KotlinTypeFactory.simpleTypeWithNonTrivialMemberScope(
-                    type.annotations,
-                    newConstructor,
-                    listOf(),
-                    false,
-                    type.memberScope
-                )
-            }
-
-            is IntersectionTypeConstructor -> if (type.isMarkedNullable) {
-                val newConstructor = constructor.transformComponents(transform = { it.makeNullable() }) ?: constructor
-                return newConstructor.createType()
-
-            }
-        }
-
-        return type
-    }
-
-    override fun transformToNewType(type: UnwrappedType): UnwrappedType =
-        when (type) {
-            is SimpleType -> transformToNewType(type)
-            is FlexibleType -> {
-                val newLower = transformToNewType(type.lowerBound)
-                val newUpper = transformToNewType(type.upperBound)
-                if (newLower !== type.lowerBound || newUpper !== type.upperBound) {
-                    KotlinTypeFactory.flexibleType(newLower, newUpper)
-                } else {
-                    type
-                }
-            }
-        }.inheritEnhancement(type)
 }
 
 object NullabilityChecker {