ConeFlexibleType: statically require to have simple bounds

compiler/fir/checkers/checkers.jvm/src/org/jetbrains/kotlin/fir/analysis/jvm/checkers/expression/FirJavaGenericVarianceViolationTypeChecker.kt

@@ -132,6 +132,12 @@ object FirJavaGenericVarianceViolationTypeChecker : FirFunctionCallChecker() {
                 lowerBound.removeOutProjection(isCovariant),
                 upperBound.removeOutProjection(isCovariant)
             )
+            is ConeSimpleKotlinType -> removeOutProjection(isCovariant)
+        }
+    }
+
+    private fun ConeSimpleKotlinType.removeOutProjection(isCovariant: Boolean): ConeSimpleKotlinType {
+        return when (this) {
             is ConeCapturedType -> ConeCapturedType(
                 captureStatus,
                 lowerType?.removeOutProjection(isCovariant),

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeTypes.kt

@@ -128,16 +128,10 @@ abstract class ConeClassLikeType : ConeLookupTagBasedType() {
 }
 
 open class ConeFlexibleType(
-    val lowerBound: ConeKotlinType,
-    val upperBound: ConeKotlinType
+    val lowerBound: ConeSimpleKotlinType,
+    val upperBound: ConeSimpleKotlinType
 ) : ConeKotlinType(), FlexibleTypeMarker {
 
-    init {
-        val message = { "Bounds violation: $lowerBound, $upperBound" }
-        require(lowerBound is SimpleTypeMarker, message)
-        require(upperBound is SimpleTypeMarker, message)
-    }
-
     override val typeArguments: Array<out ConeTypeProjection>
         get() = lowerBound.typeArguments
 
@@ -167,8 +161,19 @@ open class ConeFlexibleType(
 
 }
 
-fun ConeKotlinType.upperBoundIfFlexible() = (this as? ConeFlexibleType)?.upperBound ?: this
-fun ConeKotlinType.lowerBoundIfFlexible() = (this as? ConeFlexibleType)?.lowerBound ?: this
+fun ConeKotlinType.upperBoundIfFlexible(): ConeSimpleKotlinType {
+    return when (this) {
+        is ConeSimpleKotlinType -> this
+        is ConeFlexibleType -> upperBound
+    }
+}
+
+fun ConeKotlinType.lowerBoundIfFlexible(): ConeSimpleKotlinType {
+    return when (this) {
+        is ConeSimpleKotlinType -> this
+        is ConeFlexibleType -> lowerBound
+    }
+}
 
 class ConeCapturedTypeConstructor(
     val projection: ConeTypeProjection,
@@ -264,7 +269,10 @@ data class ConeDefinitelyNotNullType(val original: ConeKotlinType) : ConeSimpleK
     companion object
 }
 
-class ConeRawType(lowerBound: ConeKotlinType, upperBound: ConeKotlinType) : ConeFlexibleType(lowerBound, upperBound), RawTypeMarker
+class ConeRawType(
+    lowerBound: ConeSimpleKotlinType,
+    upperBound: ConeSimpleKotlinType
+) : ConeFlexibleType(lowerBound, upperBound), RawTypeMarker
 
 /*
  * Contract of the intersection type: it is flat. It means that

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/enhancement/javaTypeUtils.kt

@@ -86,14 +86,13 @@ private fun ClassId.mutableToReadOnly(): ClassId? {
     return JavaToKotlinClassMap.mutableToReadOnly(this)
 }
 
-private fun ConeKotlinType.enhanceInflexibleType(
+private fun ConeSimpleKotlinType.enhanceInflexibleType(
     session: FirSession,
     position: TypeComponentPosition,
     qualifiers: IndexedJavaTypeQualifiers,
     index: Int,
     subtreeSizes: List<Int>,
-): ConeKotlinType? {
-    require(this !is ConeFlexibleType) { "$this should not be flexible" }
+): ConeSimpleKotlinType? {
     val shouldEnhance = position.shouldEnhance()
     if ((!shouldEnhance && typeArguments.isEmpty()) || this !is ConeLookupTagBasedType) {
         return null

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

@@ -49,6 +49,13 @@ fun ConeKotlinType.fullyExpandedType(
     else -> this
 }
 
+fun ConeSimpleKotlinType.fullyExpandedType(
+    useSiteSession: FirSession
+): ConeSimpleKotlinType = when (this) {
+    is ConeClassLikeType -> fullyExpandedType(useSiteSession)
+    else -> this
+}
+
 private fun ConeClassLikeType.fullyExpandedTypeNoCache(
     useSiteSession: FirSession,
     expandedConeType: (FirTypeAlias) -> ConeClassLikeType?,

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

@@ -104,12 +104,12 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
 
     override fun FlexibleTypeMarker.upperBound(): SimpleTypeMarker {
         require(this is ConeFlexibleType)
-        return this.upperBound as SimpleTypeMarker
+        return this.upperBound
     }
 
     override fun FlexibleTypeMarker.lowerBound(): SimpleTypeMarker {
         require(this is ConeFlexibleType)
-        return this.lowerBound as SimpleTypeMarker
+        return this.lowerBound
     }
 
     override fun SimpleTypeMarker.asCapturedType(): CapturedTypeMarker? {

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

@@ -11,19 +11,25 @@ import org.jetbrains.kotlin.types.AbstractTypePreparator
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 
 class ConeTypePreparator(val session: FirSession) : AbstractTypePreparator() {
-    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
+    private fun prepareType(type: ConeSimpleKotlinType): ConeSimpleKotlinType {
         return when (type) {
             is ConeClassLikeType -> type.fullyExpandedType(session)
+            else -> type
+        }
+    }
+
+    override fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker {
+        if (type !is ConeKotlinType) {
+            throw AssertionError("Unexpected type in ConeTypePreparator: ${this::class.java}")
+        }
+        return when (type) {
             is ConeFlexibleType -> {
                 val lowerBound = prepareType(type.lowerBound)
                 if (lowerBound === type.lowerBound) return type
 
-                ConeFlexibleType(
-                    lowerBound as ConeKotlinType,
-                    prepareType(type.upperBound) as ConeKotlinType
-                )
+                ConeFlexibleType(lowerBound, prepareType(type.upperBound))
             }
-            else -> type
+            is ConeSimpleKotlinType -> prepareType(type)
         }
     }
 }

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/TypeUtils.kt

@@ -28,6 +28,8 @@ import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.resolve.calls.NewCommonSuperTypeCalculator
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.model.*
+import org.jetbrains.kotlin.fir.types.lowerBoundIfFlexible as coneLowerBoundIfFlexible
+import org.jetbrains.kotlin.fir.types.upperBoundIfFlexible as coneUpperBoundIfFlexible
 
 fun ConeInferenceContext.commonSuperTypeOrNull(types: List<ConeKotlinType>): ConeKotlinType? {
     return when (types.size) {
@@ -179,20 +181,20 @@ fun coneFlexibleOrSimpleType(
     lowerBound: ConeKotlinType,
     upperBound: ConeKotlinType,
 ): ConeKotlinType {
-    if (lowerBound is ConeFlexibleType) {
-        return coneFlexibleOrSimpleType(typeContext, lowerBound.lowerBound, upperBound)
-    }
-    if (upperBound is ConeFlexibleType) {
-        return coneFlexibleOrSimpleType(typeContext, lowerBound, upperBound.upperBound)
-    }
-    return when {
-        AbstractStrictEqualityTypeChecker.strictEqualTypes(typeContext, lowerBound, upperBound) -> lowerBound
-        else -> ConeFlexibleType(lowerBound, upperBound)
+    return when (lowerBound) {
+        is ConeFlexibleType -> coneFlexibleOrSimpleType(typeContext, lowerBound.lowerBound, upperBound)
+        is ConeSimpleKotlinType -> when (upperBound) {
+            is ConeFlexibleType -> coneFlexibleOrSimpleType(typeContext, lowerBound, upperBound.upperBound)
+            is ConeSimpleKotlinType -> when {
+                AbstractStrictEqualityTypeChecker.strictEqualTypes(typeContext, lowerBound, upperBound) -> lowerBound
+                else -> ConeFlexibleType(lowerBound, upperBound)
+            }
+        }
     }
 }
 
 fun ConeKotlinType.isExtensionFunctionType(session: FirSession): Boolean {
-    val type = this.lowerBoundIfFlexible().fullyExpandedType(session)
+    val type = this.coneLowerBoundIfFlexible().fullyExpandedType(session)
     return type.attributes.extensionFunctionType != null
 }
 
@@ -201,7 +203,7 @@ fun FirTypeRef.isExtensionFunctionType(session: FirSession): Boolean {
 }
 
 fun ConeKotlinType.isUnsafeVarianceType(session: FirSession): Boolean {
-    val type = this.lowerBoundIfFlexible().fullyExpandedType(session)
+    val type = this.coneLowerBoundIfFlexible().fullyExpandedType(session)
     return type.attributes.unsafeVarianceType != null
 }
 
@@ -471,7 +473,7 @@ fun ConeTypeContext.captureFromExpressionInternal(type: ConeKotlinType): ConeKot
     fun findCorrespondingCapturedArgumentsForType(type: ConeKotlinType) =
         capturedArgumentsByComponents.find { typeToCapture -> typeToCapture.isSuitableForType(type, this) }?.capturedArguments
 
-    fun replaceArgumentsWithCapturedArgumentsByIntersectionComponents(typeToReplace: ConeKotlinType): List<ConeKotlinType> {
+    fun replaceArgumentsWithCapturedArgumentsByIntersectionComponents(typeToReplace: ConeSimpleKotlinType): List<ConeKotlinType> {
         return if (typeToReplace is ConeIntersectionType) {
             typeToReplace.intersectedTypes.map { componentType ->
                 val capturedArguments = findCorrespondingCapturedArgumentsForType(componentType)
@@ -485,15 +487,18 @@ fun ConeTypeContext.captureFromExpressionInternal(type: ConeKotlinType): ConeKot
         }
     }
 
-    return if (type is ConeFlexibleType) {
-        val lowerIntersectedType = intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type.lowerBound))
-            .withNullability(type.lowerBound.isMarkedNullable) as ConeKotlinType
-        val upperIntersectedType = intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type.upperBound))
-            .withNullability(type.upperBound.isMarkedNullable) as ConeKotlinType
+    return when (type) {
+        is ConeFlexibleType -> {
+            val lowerIntersectedType = intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type.lowerBound))
+                .withNullability(ConeNullability.create(type.lowerBound.isMarkedNullable), this)
+            val upperIntersectedType = intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type.upperBound))
+                .withNullability(ConeNullability.create(type.upperBound.isMarkedNullable), this)
 
-        ConeFlexibleType(lowerIntersectedType, upperIntersectedType)
-    } else {
-        intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type)).withNullability(type.isMarkedNullable) as ConeKotlinType
+            ConeFlexibleType(lowerIntersectedType.coneLowerBoundIfFlexible(), upperIntersectedType.coneUpperBoundIfFlexible())
+        }
+        is ConeSimpleKotlinType -> {
+            intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type)).withNullability(type.isMarkedNullable) as ConeKotlinType
+        }
     }
 }
 

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

@@ -58,8 +58,8 @@ class FirSamResolverImpl(
         return when (type) {
             is ConeClassLikeType -> getFunctionTypeForPossibleSamType(type.fullyExpandedType(session))
             is ConeFlexibleType -> ConeFlexibleType(
-                getFunctionTypeForPossibleSamType(type.lowerBound) ?: return null,
-                getFunctionTypeForPossibleSamType(type.upperBound) ?: return null,
+                getFunctionTypeForPossibleSamType(type.lowerBound)?.lowerBoundIfFlexible() ?: return null,
+                getFunctionTypeForPossibleSamType(type.upperBound)?.upperBoundIfFlexible() ?: return null,
             )
             is ConeClassErrorType, is ConeStubType -> null
             // TODO: support those types as well

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CreateFreshTypeVariableSubstitutorStage.kt

@@ -92,7 +92,7 @@ internal object CreateFreshTypeVariableSubstitutorStage : ResolutionStage() {
         session: FirSession,
     ): ConeKotlinType {
         return if (typeParameter.shouldBeFlexible(session.typeContext)) {
-            val notNullType = type.withNullability(ConeNullability.NOT_NULL, session.typeContext)
+            val notNullType = type.withNullability(ConeNullability.NOT_NULL, session.typeContext) as ConeSimpleKotlinType
             ConeFlexibleType(notNullType, notNullType.withNullability(ConeNullability.NULLABLE, session.typeContext))
         } else {
             type