FIR: Support captureFromArguments for ConeTypes

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

@@ -114,10 +114,33 @@ abstract class ConeFunctionType : ConeClassLikeType() {
     abstract val returnType: ConeKotlinType
 }
 
-class ConeFlexibleType(val lowerBound: ConeLookupTagBasedType, val upperBound: ConeLookupTagBasedType) : ConeKotlinType(), FlexibleTypeMarker {
+class ConeFlexibleType(val lowerBound: ConeLookupTagBasedType, val upperBound: ConeLookupTagBasedType) : ConeKotlinType(),
+    FlexibleTypeMarker {
     override val typeArguments: Array<out ConeKotlinTypeProjection>
         get() = emptyArray()
 
     override val nullability: ConeNullability
         get() = lowerBound.nullability.takeIf { it == upperBound.nullability } ?: ConeNullability.UNKNOWN
 }
+
+class ConeCapturedTypeConstructor(val projection: ConeKotlinTypeProjection) : TypeConstructorMarker {
+    var supertypes: List<ConeKotlinType>? = null
+}
+
+class ConeCapturedType(
+    val captureStatus: CaptureStatus,
+    val lowerType: ConeKotlinType?,
+    override val nullability: ConeNullability = ConeNullability.NOT_NULL,
+    val constructor: ConeCapturedTypeConstructor
+) : ConeKotlinType(), SimpleTypeMarker, CapturedTypeMarker {
+    constructor(captureStatus: CaptureStatus, lowerType: ConeKotlinType?, projection: ConeKotlinTypeProjection) : this(
+        captureStatus,
+        lowerType,
+        constructor = ConeCapturedTypeConstructor(
+            projection
+        )
+    )
+
+    override val typeArguments: Array<out ConeKotlinTypeProjection>
+        get() = emptyArray()
+}

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

@@ -7,11 +7,12 @@ package org.jetbrains.kotlin.fir.types
 
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.descriptors.TypeAliasDescriptor
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.expandedConeType
 import org.jetbrains.kotlin.fir.declarations.superConeTypes
+import org.jetbrains.kotlin.fir.resolve.directExpansionType
 import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.resolve.withArguments
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeSymbol
@@ -23,13 +24,8 @@ import org.jetbrains.kotlin.fir.types.impl.ConeAbbreviatedTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeClassTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeFunctionTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
-import org.jetbrains.kotlin.resolve.calls.inference.CapturedType
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
-import org.jetbrains.kotlin.types.DefinitelyNotNullType
-import org.jetbrains.kotlin.types.SimpleType
-import org.jetbrains.kotlin.types.checker.NewCapturedType
 import org.jetbrains.kotlin.types.checker.convertVariance
-import org.jetbrains.kotlin.types.isError
 import org.jetbrains.kotlin.types.model.*
 
 
@@ -40,7 +36,12 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
 
     override fun KotlinTypeMarker.asSimpleType(): SimpleTypeMarker? {
         assert(this is ConeKotlinType)
-        return this as? ConeLookupTagBasedType
+        return when (this) {
+            is ConeAbbreviatedType -> directExpansionType(session)
+            is ConeCapturedType -> this
+            is ConeLookupTagBasedType -> this
+            else -> null
+        }
     }
 
     override fun KotlinTypeMarker.asFlexibleType(): FlexibleTypeMarker? {
@@ -74,25 +75,25 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
     override fun SimpleTypeMarker.asCapturedType(): CapturedTypeMarker? {
-        require(this is ConeLookupTagBasedType)
+        //require(this is ConeLookupTagBasedType)
-        return null // TODO
+        return this as? ConeCapturedType
     }
 
     override fun SimpleTypeMarker.asDefinitelyNotNullType(): DefinitelyNotNullTypeMarker? {
-        require(this is ConeLookupTagBasedType)
+        require(this is ConeKotlinType)
         return null // TODO
     }
 
     override fun SimpleTypeMarker.isMarkedNullable(): Boolean {
-        require(this is ConeLookupTagBasedType)
+        require(this is ConeKotlinType)
         return this.nullability.isNullable
     }
 
     override fun SimpleTypeMarker.withNullability(nullable: Boolean): SimpleTypeMarker {
-        require(this is ConeLookupTagBasedType)
+        require(this is ConeKotlinType)
         if (nullability.isNullable == nullable) return this
         return when (this) {
-
+            is ConeCapturedType -> ConeCapturedType(captureStatus, lowerType, ConeNullability.create(nullable), constructor)
             is ConeTypeParameterType -> ConeTypeParameterTypeImpl(lookupTag, nullable)
             is ConeClassErrorType -> this
             is ConeClassType -> ConeClassTypeImpl(lookupTag, typeArguments, nullable)
@@ -108,23 +109,27 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
                 lookupTag,
                 nullable
             )
+            else -> error("!")
         }
     }
 
     override fun SimpleTypeMarker.typeConstructor(): TypeConstructorMarker {
-        require(this is ConeLookupTagBasedType)
+        return when (this) {
-        if (this is ConeClassErrorType) return ErrorTypeConstructor("No constructor: $reason")
+            is ConeCapturedType -> constructor
-        return this.lookupTag.toSymbol(session) ?: ErrorTypeConstructor("Unresolved: ${this.lookupTag}")
+            is ConeLookupTagBasedType -> this.lookupTag.toSymbol(session) ?: ErrorTypeConstructor("Unresolved: ${this.lookupTag}")
+            else -> error("?: ${this}")
+        }
+
     }
 
     override fun SimpleTypeMarker.argumentsCount(): Int {
-        require(this is ConeLookupTagBasedType)
+        require(this is ConeKotlinType)
 
         return this.typeArguments.size
     }
 
     override fun SimpleTypeMarker.getArgument(index: Int): TypeArgumentMarker {
-        require(this is ConeLookupTagBasedType)
+        require(this is ConeKotlinType)
 
         return this.typeArguments[index]
     }
@@ -136,8 +141,8 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
     override fun CapturedTypeMarker.lowerType(): KotlinTypeMarker? {
-        require(this is ConeKotlinType)
+        require(this is ConeCapturedType)
-        return null // TODO
+        return this.lowerType
     }
 
     override fun TypeArgumentMarker.isStarProjection(): Boolean {
@@ -163,9 +168,9 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
     override fun TypeConstructorMarker.parametersCount(): Int {
-        require(this is ConeSymbol)
+        //require(this is ConeSymbol)
         return when (this) {
-            is ConeTypeParameterSymbol -> 0
+            is ConeTypeParameterSymbol, is ConeCapturedTypeConstructor, is ErrorTypeConstructor -> 0
             is FirClassSymbol -> fir.typeParameters.size
             is FirTypeAliasSymbol -> fir.typeParameters.size
             else -> error("?!:10")
@@ -173,7 +178,7 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
     override fun TypeConstructorMarker.getParameter(index: Int): TypeParameterMarker {
-        require(this is ConeSymbol)
+        //require(this is ConeSymbol)
         return when (this) {
             is ConeTypeParameterSymbol -> error("?!:11")
             is FirClassSymbol -> fir.typeParameters[index].symbol
@@ -184,11 +189,12 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
 
     override fun TypeConstructorMarker.supertypes(): Collection<KotlinTypeMarker> {
         if (this is ErrorTypeConstructor) return emptyList()
-        require(this is ConeSymbol)
+        //require(this is ConeSymbol)
         return when (this) {
             is ConeTypeParameterSymbol -> emptyList()
             is FirClassSymbol -> fir.superConeTypes
             is FirTypeAliasSymbol -> listOfNotNull(fir.expandedConeType)
+            is ConeCapturedTypeConstructor -> supertypes!!
             else -> error("?!:13")
         }
     }
@@ -198,7 +204,7 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
     override fun TypeConstructorMarker.isClassTypeConstructor(): Boolean {
-        assert(this is ConeSymbol)
+        //assert(this is ConeSymbol)
         return this is FirClassSymbol
     }
 
@@ -224,8 +230,9 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
 
     override fun isEqualTypeConstructors(c1: TypeConstructorMarker, c2: TypeConstructorMarker): Boolean {
         if (c1 is ErrorTypeConstructor || c2 is ErrorTypeConstructor) return false
-        require(c1 is ConeSymbol)
+
-        require(c2 is ConeSymbol)
+        //assert(c1 is ConeSymbol)
+        //assert(c2 is ConeSymbol)
         return c1 == c2
     }
 
@@ -234,7 +241,6 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
     override fun TypeConstructorMarker.isCommonFinalClassConstructor(): Boolean {
-        require(this is ConeSymbol)
         val classSymbol = this as? ConeClassSymbol ?: return false
         val fir = (classSymbol as FirClassSymbol).fir
         return fir.modality == Modality.FINAL &&
@@ -243,7 +249,43 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
     override fun captureFromArguments(type: SimpleTypeMarker, status: CaptureStatus): SimpleTypeMarker? {
-        return type //TODO
+        require(type is ConeLookupTagBasedType) { ":( $type" }
+        val typeConstructor = type.typeConstructor()
+        if (type.argumentsCount() != typeConstructor.parametersCount()) return null
+        if (type.asArgumentList().all(this) { !it.isStarProjection() && it.getVariance() == TypeVariance.INV }) return null
+        val newArguments = Array(type.argumentsCount()) { index ->
+            val argument = type.getArgument(index)
+            if (!argument.isStarProjection() && argument.getVariance() == TypeVariance.INV) return@Array argument as ConeKotlinTypeProjection
+
+            val lowerType = if (!argument.isStarProjection() && argument.getVariance() == TypeVariance.IN) {
+                argument.getType() as ConeKotlinType
+            } else {
+                null
+            }
+
+            ConeCapturedType(status, lowerType, argument as ConeKotlinTypeProjection)
+        }
+
+        for (index in 0 until type.argumentsCount()) {
+            val oldArgument = type.getArgument(index)
+            val newArgument = newArguments[index]
+
+            if (!oldArgument.isStarProjection() && oldArgument.getVariance() == TypeVariance.INV) continue
+
+            val parameter = typeConstructor.getParameter(index)
+            val upperBounds = (0 until parameter.upperBoundCount()).mapTo(mutableListOf()) { paramIndex ->
+                parameter.getUpperBound(paramIndex) // TODO: substitution
+            }
+
+            if (!oldArgument.isStarProjection() && oldArgument.getVariance() == TypeVariance.OUT) {
+                upperBounds += oldArgument.getType()
+            }
+
+            require(newArgument is ConeCapturedType)
+            newArgument.constructor.supertypes = upperBounds as List<ConeKotlinType>
+        }
+
+        return type.withArguments(newArguments)
     }
 
     override fun SimpleTypeMarker.asArgumentList(): TypeArgumentListMarker {
@@ -252,8 +294,8 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
     override fun identicalArguments(a: SimpleTypeMarker, b: SimpleTypeMarker): Boolean {
-        require(a is ConeLookupTagBasedType)
+        require(a is ConeKotlinType)
-        require(b is ConeLookupTagBasedType)
+        require(b is ConeKotlinType)
         return a.typeArguments === b.typeArguments
     }
 
@@ -273,15 +315,16 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext {
     }
 
 
-
     override fun SimpleTypeMarker.isSingleClassifierType(): Boolean {
+        if (isError()) return false
+        if (this is ConeCapturedType) return true
         require(this is ConeLookupTagBasedType)
         val symbol = this.lookupTag.toSymbol(session)
-        return !isError() &&
+        return symbol is FirClassSymbol ||
-                (symbol is FirClassSymbol ||
+                symbol is FirTypeParameterSymbol
-                        symbol is FirTypeParameterSymbol)
-        // || TODO: (this is CapturedType || this is NewCapturedType || this is DefinitelyNotNullType)
     }
+
+
 }
 
 class ConeTypeCheckerContext(override val isErrorTypeEqualsToAnything: Boolean, override val session: FirSession) :
@@ -300,7 +343,7 @@ class ConeTypeCheckerContext(override val isErrorTypeEqualsToAnything: Boolean,
     }
 
     override fun intersectTypes(types: List<KotlinTypeMarker>): KotlinTypeMarker {
-        TODO("not implemented")
+        return types.first() // TODO: proper implementation
     }
 
     override val KotlinTypeMarker.isAllowedTypeVariable: Boolean