Support type inference for self type materialization calls

2021-06-30 16:21:58 +03:00
parent 44cf4be1e5
commit 51c5a54e31
21 changed files with 183 additions and 89 deletions
@@ -270,6 +270,11 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
        return this.withArguments(newArguments.cast<List<ConeTypeProjection>>().toTypedArray(), this@ConeInferenceContext)
    }

+    override fun SimpleTypeMarker.replaceArguments(replacement: (TypeArgumentMarker) -> TypeArgumentMarker): SimpleTypeMarker {
+        require(this is ConeKotlinType)
+        return this.withArguments({ replacement(it).cast() }, this@ConeInferenceContext)
+    }
+
    override fun KotlinTypeMarker.hasExactAnnotation(): Boolean {
        require(this is ConeKotlinType)
        return attributes.exact != null
@@ -239,6 +239,15 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
        }
    }

+    override fun TypeConstructorMarker.getParameters(): List<TypeParameterMarker> {
+        return when (val symbol = toClassLikeSymbol()) {
+            is FirAnonymousObjectSymbol -> symbol.fir.typeParameters.map { it.symbol.toLookupTag() }
+            is FirRegularClassSymbol -> symbol.fir.typeParameters.map { it.symbol.toLookupTag() }
+            is FirTypeAliasSymbol -> symbol.fir.typeParameters.map { it.symbol.toLookupTag() }
+            else -> error("Unexpected FirClassLikeSymbol $symbol for ${this::class}, with classId ${(this as? ConeClassLikeLookupTag)?.classId}")
+        }
+    }
+
    private fun TypeConstructorMarker.toClassLikeSymbol(): FirClassLikeSymbol<*>? = (this as? ConeClassLikeLookupTag)?.toSymbol(session)

    override fun TypeConstructorMarker.supertypes(): Collection<ConeKotlinType> {
@@ -297,11 +306,11 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
        return this
    }

-    override fun TypeParameterMarker.hasRecursiveBounds(selfConstructor: TypeConstructorMarker): Boolean {
+    override fun TypeParameterMarker.hasRecursiveBounds(selfConstructor: TypeConstructorMarker?): Boolean {
        require(this is ConeTypeParameterLookupTag)
        return this.typeParameterSymbol.fir.bounds.any { typeRef ->
            typeRef.coneType.contains { it.typeConstructor() == this.getTypeConstructor() }
-                    && typeRef.coneType.typeConstructor() == selfConstructor
+                    && (selfConstructor == null || typeRef.coneType.typeConstructor() == selfConstructor)
        }
    }

@@ -89,6 +89,9 @@ fun <T : ConeKotlinType> T.withArguments(arguments: Array<out ConeTypeProjection
    }
 }

+fun <T : ConeKotlinType> T.withArguments(replacement: (ConeTypeProjection) -> ConeTypeProjection, typeSystemContext: ConeTypeContext) =
+    withArguments(typeArguments.map(replacement).toTypedArray(), typeSystemContext)
+
 fun <T : ConeKotlinType> T.withAttributes(attributes: ConeAttributes, typeSystemContext: ConeTypeContext): T {
    if (this.attributes == attributes) {
        return this
@@ -139,6 +139,8 @@ interface IrTypeSystemContext : TypeSystemContext, TypeSystemCommonSuperTypesCon

    override fun TypeConstructorMarker.getParameter(index: Int) = getTypeParameters(this)[index].symbol

+    override fun TypeConstructorMarker.getParameters() = getTypeParameters(this).map { it.symbol }
+
    override fun TypeConstructorMarker.supertypes(): Collection<KotlinTypeMarker> {
        return when (this) {
            is IrCapturedType.Constructor -> superTypes
@@ -179,10 +181,10 @@ interface IrTypeSystemContext : TypeSystemContext, TypeSystemCommonSuperTypesCon
        return false
    }

-    override fun TypeParameterMarker.hasRecursiveBounds(selfConstructor: TypeConstructorMarker): Boolean {
+    override fun TypeParameterMarker.hasRecursiveBounds(selfConstructor: TypeConstructorMarker?): Boolean {
        for (i in 0 until this.upperBoundCount()) {
            val upperBound = this.getUpperBound(i)
-            if (upperBound.containsTypeConstructor(selfConstructor) && upperBound.typeConstructor() == selfConstructor) {
+            if (upperBound.containsTypeConstructor(this.getTypeConstructor()) && (selfConstructor == null || upperBound.typeConstructor() == selfConstructor)) {
                return true
            }
        }
@@ -5,7 +5,9 @@

 package org.jetbrains.kotlin.resolve.calls.inference

+import org.jetbrains.kotlin.resolve.calls.inference.model.Constraint
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
+import org.jetbrains.kotlin.resolve.calls.inference.model.DeclaredUpperBoundConstraintPosition
 import org.jetbrains.kotlin.types.model.*

 fun ConstraintStorage.buildCurrentSubstitutor(
@@ -46,3 +48,11 @@ fun ConstraintStorage.buildNotFixedVariablesToNonSubtypableTypesSubstitutor(
        notFixedTypeVariables.mapValues { context.createStubTypeForTypeVariablesInSubtyping(it.value.typeVariable) }
    )
 }
+
+fun TypeSystemInferenceExtensionContext.hasDeclaredUpperBoundSelfTypes(constraint: Constraint): Boolean {
+    val typeConstructor = constraint.type.typeConstructor()
+
+    return constraint.position.from is DeclaredUpperBoundConstraintPosition<*>
+            && (typeConstructor.getParameters().any { it.hasRecursiveBounds(typeConstructor) }
+            || typeConstructor.getTypeParameterClassifier()?.hasRecursiveBounds() == true)
+}
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.resolve.calls.inference.components

 import org.jetbrains.kotlin.resolve.calls.NewCommonSuperTypeCalculator
 import org.jetbrains.kotlin.resolve.calls.inference.components.TypeVariableDirectionCalculator.ResolveDirection
+import org.jetbrains.kotlin.resolve.calls.inference.hasDeclaredUpperBoundSelfTypes
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
 import org.jetbrains.kotlin.types.AbstractTypeChecker
@@ -210,8 +211,9 @@ class ResultTypeResolver(
    }

    private fun Context.findSuperType(variableWithConstraints: VariableWithConstraints): KotlinTypeMarker? {
-        val upperConstraints =
-            variableWithConstraints.constraints.filter { it.kind == ConstraintKind.UPPER && this@findSuperType.isProperTypeForFixation(it.type) }
+        val upperConstraints = variableWithConstraints.constraints.filter {
+            it.kind == ConstraintKind.UPPER && (hasDeclaredUpperBoundSelfTypes(it) || isProperTypeForFixation(it.type))
+        }
        if (upperConstraints.isNotEmpty()) {
            val intersectionUpperType = intersectTypes(upperConstraints.map { it.type })
            val resultIsActuallyIntersection = intersectionUpperType.typeConstructor().isIntersection()
@@ -238,7 +240,7 @@ class ResultTypeResolver(

            return typeApproximator.approximateToSubType(
                upperType,
-                TypeApproximatorConfiguration.InternalTypesApproximation
+                TypeApproximatorConfiguration.InternalAndSelfTypesApproximation
            ) ?: upperType
        }
        return null
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.resolve.calls.inference.components
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode.PARTIAL
+import org.jetbrains.kotlin.resolve.calls.inference.hasDeclaredUpperBoundSelfTypes
 import org.jetbrains.kotlin.resolve.calls.inference.model.Constraint
 import org.jetbrains.kotlin.resolve.calls.inference.model.DeclaredUpperBoundConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
@@ -48,6 +49,7 @@ class VariableFixationFinder(
        FROM_INCORPORATION_OF_DECLARED_UPPER_BOUND,
        READY_FOR_FIXATION_UPPER,
        READY_FOR_FIXATION_LOWER,
+        READY_FOR_FIXATION_DECLARED_UPPER_BOUND_WITH_SELF_TYPES,
        READY_FOR_FIXATION,
        READY_FOR_FIXATION_REIFIED,
    }
@@ -61,6 +63,7 @@ class VariableFixationFinder(
    ): TypeVariableFixationReadiness = when {
        !notFixedTypeVariables.contains(variable) ||
                dependencyProvider.isVariableRelatedToTopLevelType(variable) -> TypeVariableFixationReadiness.FORBIDDEN
+        hasDeclaredUpperBoundSelfTypes(variable) -> TypeVariableFixationReadiness.READY_FOR_FIXATION_DECLARED_UPPER_BOUND_WITH_SELF_TYPES
        !variableHasProperArgumentConstraints(variable) -> TypeVariableFixationReadiness.WITHOUT_PROPER_ARGUMENT_CONSTRAINT
        hasDependencyToOtherTypeVariables(variable) -> TypeVariableFixationReadiness.WITH_COMPLEX_DEPENDENCY
        variableHasTrivialOrNonProperConstraints(variable) -> TypeVariableFixationReadiness.WITH_TRIVIAL_OR_NON_PROPER_CONSTRAINTS
@@ -166,6 +169,11 @@ class VariableFixationFinder(
            it.kind.isLower() && isProperArgumentConstraint(it) && !it.type.typeConstructor().isNothingConstructor()
        }
    }
+
+    private fun Context.hasDeclaredUpperBoundSelfTypes(variable: TypeConstructorMarker): Boolean {
+        val constraints = notFixedTypeVariables[variable]?.constraints ?: return false
+        return constraints.isNotEmpty() && constraints.all(::hasDeclaredUpperBoundSelfTypes)
+    }
 }

 inline fun TypeSystemInferenceExtensionContext.isProperTypeForFixation(type: KotlinTypeMarker, isProper: (KotlinTypeMarker) -> Boolean) =
@@ -404,6 +404,15 @@ abstract class AbstractTypeApproximator(

            if (argument.isStarProjection()) continue

+            val argumentTypeConstructor = argument.getType().typeConstructor()
+
+            if (argumentTypeConstructor is TypeVariableTypeConstructorMarker && conf.selfTypesWithTypeVariablesToCapturedStarProjection) {
+                // If we have Self<TypeVariable(T)> where T is bounded by Self<T>, we approximate it to CapturedType(*) to satisfy constraints
+                if (argumentTypeConstructor.typeParameter?.hasRecursiveBounds(type.typeConstructor()) == true) {
+                    return createCapturedStarProjectionForSelfType(argumentTypeConstructor, type)
+                }
+            }
+
            val effectiveVariance = AbstractTypeChecker.effectiveVariance(parameter.getVariance(), argument.getVariance())

            val argumentType = newArguments[index]?.getType() ?: argument.getType()
@@ -24,6 +24,7 @@ open class TypeApproximatorConfiguration {
    open val intersection: IntersectionStrategy = IntersectionStrategy.TO_COMMON_SUPERTYPE
    open val intersectionTypesInContravariantPositions = false
    open val localTypes = false
+    open val selfTypesWithTypeVariablesToCapturedStarProjection = false

    open val typeVariable: (TypeVariableTypeConstructorMarker) -> Boolean = { false }
    open fun capturedType(ctx: TypeSystemInferenceExtensionContext, type: CapturedTypeMarker): Boolean =
@@ -73,6 +74,12 @@ open class TypeApproximatorConfiguration {
        override val intersectionTypesInContravariantPositions: Boolean get() = true
    }

+    object InternalAndSelfTypesApproximation : AbstractCapturedTypesApproximation(CaptureStatus.FROM_EXPRESSION) {
+        override val integerLiteralType: Boolean get() = true
+        override val intersectionTypesInContravariantPositions: Boolean get() = true
+        override val selfTypesWithTypeVariablesToCapturedStarProjection = true
+    }
+
    object FinalApproximationAfterResolutionAndInference :
        AbstractCapturedTypesApproximation(CaptureStatus.FROM_EXPRESSION) {
        override val integerLiteralType: Boolean get() = true
@@ -41,7 +41,7 @@ class ClassicTypeSystemContextForCS(override val builtIns: KotlinBuiltIns) : Typ
        captureStatus: CaptureStatus
    ): CapturedTypeMarker {
        require(lowerType is UnwrappedType?, lowerType::errorMessage)
-        require(constructorProjection is TypeProjectionImpl, constructorProjection::errorMessage)
+        require(constructorProjection is TypeProjectionBase, constructorProjection::errorMessage)

        @Suppress("UNCHECKED_CAST")
        val newCapturedTypeConstructor = NewCapturedTypeConstructor(
@@ -49,7 +49,7 @@ class ClassicTypeSystemContextForCS(override val builtIns: KotlinBuiltIns) : Typ
            constructorSupertypes as List<UnwrappedType>
        )
        return NewCapturedType(
-            CaptureStatus.FOR_INCORPORATION,
+            captureStatus,
            newCapturedTypeConstructor,
            lowerType = lowerType
        )
@@ -299,11 +299,6 @@ class KotlinConstraintSystemCompleter(
                    return true
                }

-                hasDeclaredUpperBoundSelfTypes(variableWithConstraints) -> {
-                    fixVariable(this, variableWithConstraints, topLevelAtoms)
-                    return true
-                }
-
                else -> processVariableWhenNotEnoughInformation(this, variableWithConstraints, topLevelAtoms, diagnosticsHolder)
            }
        }
@@ -311,33 +306,6 @@ class KotlinConstraintSystemCompleter(
        return false
    }

-    private fun ConstraintSystemCompletionContext.hasDeclaredUpperBoundSelfTypes(variable: VariableWithConstraints): Boolean {
-        val constraints = variable.constraints
-        return constraints.isNotEmpty() && constraints.all {
-            it.position.from is DeclaredUpperBoundConstraintPosition<*> && isSelfType(it.type.typeConstructor())
-        }
-    }
-
-    private fun ConstraintSystemCompletionContext.isSelfType(typeConstructor: TypeConstructorMarker): Boolean {
-        if (typeConstructor.isTypeParameterTypeConstructor()) {
-            val supertype = typeConstructor.supertypes().firstOrNull() ?: return false
-            return isSelfType(supertype.typeConstructor())
-        }
-
-        val parametersCount = typeConstructor.parametersCount()
-        if (parametersCount == 0) return false
-
-        for (parameterId in 0 until parametersCount) {
-            val parameter = typeConstructor.getParameter(parameterId)
-            if (parameter.upperBoundCount() != 1) return false
-
-            val upperBound = parameter.getUpperBound(0)
-            if (upperBound.typeConstructor() == typeConstructor) return true
-        }
-
-        return false
-    }
-
    private fun processVariableWhenNotEnoughInformation(
        c: ConstraintSystemCompletionContext,
        variableWithConstraints: VariableWithConstraints,
@@ -18,12 +18,4 @@ fun <K : Builder<K>> testTypeParam(builder: Builder<K>) {
    builder
        .test()
        .foo()
-}
-
-interface BodySpec<B, S : BodySpec<B, S>> {
-    fun <T : S> isEqualTo(expected: B): T
-}
-
-fun test(b: BodySpec<String, *>) {
-    b.isEqualTo("")
 }
@@ -1,15 +1,24 @@
 fun test() {
-    WriterAppender.newBuilder()
+    <!DEBUG_INFO_EXPRESSION_TYPE("WriterAppender.Builder1<*>")!>WriterAppender.newBuilder()<!>
+    <!DEBUG_INFO_EXPRESSION_TYPE("WriterAppender.Builder1<out WriterAppender.Builder1<*>>")!>WriterAppender.Builder1()<!>
+
+    <!DEBUG_INFO_EXPRESSION_TYPE("{Builder1<*> & Builder2<*>}")!>WriterAppender.intersectTwoSelfTypes()<!>
 }

 object WriterAppender {
-    class Builder<B : Builder<B>> {
+    interface Builder2<K : Builder2<K>>
+
+    class Builder1<B : Builder1<B>> {
        fun asBuilder(): B {
-            return this as B
+            return this <!UNCHECKED_CAST!>as B<!>
        }
    }

-    fun <B : Builder<B>> newBuilder(): B {
-        return Builder<B>().asBuilder()
+    fun <B : Builder1<B>> newBuilder(): B {
+        return Builder1<B>().asBuilder()
+    }
+
+    fun <B> intersectTwoSelfTypes(): B where B : Builder1<B>, B: Builder2<B> {
+        return Builder1<B>().asBuilder()
    }
 }
@@ -1,15 +1,24 @@
 fun test() {
-    WriterAppender.<!TYPE_MISMATCH!>newBuilder<!>()
+    <!DEBUG_INFO_EXPRESSION_TYPE("WriterAppender.Builder1<*>")!>WriterAppender.newBuilder()<!>
+    <!DEBUG_INFO_EXPRESSION_TYPE("WriterAppender.Builder1<out WriterAppender.Builder1<*>>")!>WriterAppender.Builder1()<!>
+
+    <!DEBUG_INFO_EXPRESSION_TYPE("{Builder1<*> & Builder2<*>}")!>WriterAppender.intersectTwoSelfTypes()<!>
 }

 object WriterAppender {
-    class Builder<B : Builder<B>> {
+    interface Builder2<K : Builder2<K>>
+
+    class Builder1<B : Builder1<B>> {
        fun asBuilder(): B {
            return this <!UNCHECKED_CAST!>as B<!>
        }
    }

-    fun <B : Builder<B>> newBuilder(): B {
-        return Builder<B>().asBuilder()
+    fun <B : Builder1<B>> newBuilder(): B {
+        return Builder1<B>().asBuilder()
+    }
+
+    fun <B> intersectTwoSelfTypes(): B where B : Builder1<B>, B: Builder2<B> {
+        return Builder1<B>().asBuilder()
    }
 }
@@ -6,14 +6,21 @@ public object WriterAppender {
    private constructor WriterAppender()
    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
-    public final fun </*0*/ B : WriterAppender.Builder<B>> newBuilder(): B
+    public final fun </*0*/ B : WriterAppender.Builder1<B>> intersectTwoSelfTypes(): B where B : WriterAppender.Builder2<B>
+    public final fun </*0*/ B : WriterAppender.Builder1<B>> newBuilder(): B
    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String

-    public final class Builder</*0*/ B : WriterAppender.Builder<B>> {
-        public constructor Builder</*0*/ B : WriterAppender.Builder<B>>()
+    public final class Builder1</*0*/ B : WriterAppender.Builder1<B>> {
+        public constructor Builder1</*0*/ B : WriterAppender.Builder1<B>>()
        public final fun asBuilder(): B
        public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
        public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
        public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
    }
+
+    public interface Builder2</*0*/ K : WriterAppender.Builder2<K>> {
+        public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+        public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+        public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    }
 }
@@ -528,7 +528,9 @@ object AbstractTypeChecker {
         * so if CapturedType(out Bar) is the same as a type of Foo's argument and Foo is a self type, then subtyping should return true.
         * If we don't handle this case separately, subtyping may not converge due to the nature of the capturing.
         */
-        if (subType is CapturedTypeMarker) {
+        val subTypeConstructor = subType.typeConstructor()
+        if (subType is CapturedTypeMarker
+            || (subTypeConstructor.isIntersection() && subTypeConstructor.supertypes().all { it is CapturedTypeMarker })) {
            val typeParameter =
                context.typeSystemContext.getTypeParameterForArgumentInBaseIfItEqualToTarget(baseType = superType, targetType = subType)
            if (typeParameter != null && typeParameter.hasRecursiveBounds(superType.typeConstructor())) {
@@ -30,7 +30,6 @@ interface TypeSystemCommonBackendContext : TypeSystemContext {
     */
    fun KotlinTypeMarker.getAnnotationFirstArgumentValue(fqName: FqName): Any?

-    fun TypeConstructorMarker.getTypeParameterClassifier(): TypeParameterMarker?
    fun TypeConstructorMarker.isInlineClass(): Boolean
    fun TypeConstructorMarker.isInnerClass(): Boolean
    fun TypeParameterMarker.getRepresentativeUpperBound(): KotlinTypeMarker
@@ -175,6 +175,7 @@ interface TypeSystemInferenceExtensionContext : TypeSystemContext, TypeSystemBui
    fun KotlinTypeMarker.removeExactAnnotation(): KotlinTypeMarker

    fun SimpleTypeMarker.replaceArguments(newArguments: List<TypeArgumentMarker>): SimpleTypeMarker
+    fun SimpleTypeMarker.replaceArguments(replacement: (TypeArgumentMarker) -> TypeArgumentMarker): SimpleTypeMarker

    fun KotlinTypeMarker.hasExactAnnotation(): Boolean
    fun KotlinTypeMarker.hasNoInferAnnotation(): Boolean
@@ -248,6 +249,22 @@ interface TypeSystemInferenceExtensionContext : TypeSystemContext, TypeSystemBui
     * In future once we have only FIR (or FE 1.0 behavior is fixed) this method should be inlined to the use-site
     */
    fun SimpleTypeMarker.createConstraintPartForLowerBoundAndFlexibleTypeVariable(): KotlinTypeMarker
+
+    fun createCapturedStarProjectionForSelfType(
+        typeVariable: TypeVariableTypeConstructorMarker,
+        selfType: SimpleTypeMarker,
+    ): SimpleTypeMarker? {
+        val typeParameter = typeVariable.typeParameter ?: return null
+        val starProjection = createStarProjection(typeParameter)
+        val superType = selfType.replaceArguments {
+            val constructor = it.getType().typeConstructor()
+            if (constructor is TypeVariableTypeConstructorMarker && constructor == typeVariable) {
+                starProjection
+            } else it
+        }
+
+        return createCapturedType(starProjection, listOf(superType), lowerType = null, CaptureStatus.FROM_EXPRESSION)
+    }
 }


@@ -313,12 +330,14 @@ interface TypeSystemContext : TypeSystemOptimizationContext {

    fun TypeConstructorMarker.parametersCount(): Int
    fun TypeConstructorMarker.getParameter(index: Int): TypeParameterMarker
+    fun TypeConstructorMarker.getParameters(): List<TypeParameterMarker>
    fun TypeConstructorMarker.supertypes(): Collection<KotlinTypeMarker>
    fun TypeConstructorMarker.isIntersection(): Boolean
    fun TypeConstructorMarker.isClassTypeConstructor(): Boolean
    fun TypeConstructorMarker.isInterface(): Boolean
    fun TypeConstructorMarker.isIntegerLiteralTypeConstructor(): Boolean
    fun TypeConstructorMarker.isLocalType(): Boolean
+    fun TypeConstructorMarker.getTypeParameterClassifier(): TypeParameterMarker?

    val TypeVariableTypeConstructorMarker.typeParameter: TypeParameterMarker?

@@ -327,7 +346,7 @@ interface TypeSystemContext : TypeSystemOptimizationContext {
    fun TypeParameterMarker.getUpperBound(index: Int): KotlinTypeMarker
    fun TypeParameterMarker.getUpperBounds(): List<KotlinTypeMarker>
    fun TypeParameterMarker.getTypeConstructor(): TypeConstructorMarker
-    fun TypeParameterMarker.hasRecursiveBounds(selfConstructor: TypeConstructorMarker): Boolean
+    fun TypeParameterMarker.hasRecursiveBounds(selfConstructor: TypeConstructorMarker? = null): Boolean

    fun areEqualTypeConstructors(c1: TypeConstructorMarker, c2: TypeConstructorMarker): Boolean

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.types

 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.ClassifierDescriptorWithTypeParameters
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.types.checker.KotlinTypeRefiner
@@ -45,18 +46,38 @@ class StarProjectionImpl(
    }
 }

-fun TypeParameterDescriptor.starProjectionType(): KotlinType {
-    val classDescriptor = this.containingDeclaration as ClassifierDescriptorWithTypeParameters
-    val typeParameters = classDescriptor.typeConstructor.parameters.map { it.typeConstructor }
-    return TypeSubstitutor.create(
-        object : TypeConstructorSubstitution() {
-            override fun get(key: TypeConstructor) =
-                if (key in typeParameters)
-                    TypeUtils.makeStarProjection(key.declarationDescriptor as TypeParameterDescriptor)
-                else null
+private fun buildStarProjectionTypeByTypeParameters(
+    typeParameters: List<TypeConstructor>,
+    upperBounds: List<KotlinType>,
+    builtIns: KotlinBuiltIns
+) = TypeSubstitutor.create(
+    object : TypeConstructorSubstitution() {
+        override fun get(key: TypeConstructor) =
+            if (key in typeParameters)
+                TypeUtils.makeStarProjection(key.declarationDescriptor as TypeParameterDescriptor)
+            else null

+    }
+).substitute(upperBounds.first(), Variance.OUT_VARIANCE) ?: builtIns.defaultBound
+
+fun TypeParameterDescriptor.starProjectionType(): KotlinType {
+    return when (val descriptor = this.containingDeclaration) {
+        is ClassifierDescriptorWithTypeParameters -> {
+            buildStarProjectionTypeByTypeParameters(
+                typeParameters = descriptor.typeConstructor.parameters.map { it.typeConstructor },
+                upperBounds,
+                builtIns
+            )
        }
-    ).substitute(this.upperBounds.first(), Variance.OUT_VARIANCE) ?: builtIns.defaultBound
+        is FunctionDescriptor -> {
+            buildStarProjectionTypeByTypeParameters(
+                typeParameters = descriptor.typeParameters.map { it.typeConstructor },
+                upperBounds,
+                builtIns
+            )
+        }
+        else -> throw IllegalArgumentException("Unsupported descriptor type to build star projection type based on type parameters of it")
+    }
 }

 // It should only be used in rare cases when type parameter for the relevant argument is not available
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.isCaptured
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.checker.*
+import org.jetbrains.kotlin.types.model.TypeArgumentMarker
 import org.jetbrains.kotlin.types.model.TypeVariableTypeConstructorMarker
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import java.util.*
@@ -182,8 +183,8 @@ fun KotlinType.getImmediateSuperclassNotAny(): KotlinType? {
 fun KotlinType.asTypeProjection(): TypeProjection = TypeProjectionImpl(this)
 fun KotlinType.contains(predicate: (UnwrappedType) -> Boolean) = TypeUtils.contains(this, predicate)

-fun KotlinType.replaceArgumentsWithStarProjections() = replaceArgumentsWith(::StarProjectionImpl)
-fun KotlinType.replaceArgumentsWithNothing() = replaceArgumentsWith { it.builtIns.nothingType.asTypeProjection() }
+fun KotlinType.replaceArgumentsWithStarProjections() = replaceArgumentsByParametersWith(::StarProjectionImpl)
+fun KotlinType.replaceArgumentsWithNothing() = replaceArgumentsByParametersWith { it.builtIns.nothingType.asTypeProjection() }

 fun KotlinType.extractTypeParametersFromUpperBounds(visitedTypeParameters: Set<TypeParameterDescriptor>?): Set<TypeParameterDescriptor> =
    mutableSetOf<TypeParameterDescriptor>().also { extractTypeParametersFromUpperBounds(this, it, visitedTypeParameters) }
@@ -247,7 +248,7 @@ fun KotlinType.replaceArgumentsWithStarProjectionOrMapped(
    variance: Variance,
    visitedTypeParameters: Set<TypeParameterDescriptor>?
 ) =
-    replaceArgumentsWith { typeParameterDescriptor ->
+    replaceArgumentsByParametersWith { typeParameterDescriptor ->
        val argument = arguments.getOrNull(typeParameterDescriptor.index)
        val isTypeParameterVisited = visitedTypeParameters != null && typeParameterDescriptor in visitedTypeParameters
        if (!isTypeParameterVisited && argument != null && argument.type.constructor in substitutionMap) {
@@ -256,18 +257,18 @@ fun KotlinType.replaceArgumentsWithStarProjectionOrMapped(
    }.let { substitutor.safeSubstitute(it, variance) }


-inline fun KotlinType.replaceArgumentsWith(replacement: (TypeParameterDescriptor) -> TypeProjection): KotlinType {
+inline fun KotlinType.replaceArgumentsByParametersWith(replacement: (TypeParameterDescriptor) -> TypeProjection): KotlinType {
    val unwrapped = unwrap()
    return when (unwrapped) {
        is FlexibleType -> KotlinTypeFactory.flexibleType(
-            unwrapped.lowerBound.replaceArgumentsWith(replacement),
-            unwrapped.upperBound.replaceArgumentsWith(replacement)
+            unwrapped.lowerBound.replaceArgumentsByParametersWith(replacement),
+            unwrapped.upperBound.replaceArgumentsByParametersWith(replacement)
        )
-        is SimpleType -> unwrapped.replaceArgumentsWith(replacement)
+        is SimpleType -> unwrapped.replaceArgumentsByParametersWith(replacement)
    }.inheritEnhancement(unwrapped)
 }

-inline fun SimpleType.replaceArgumentsWith(replacement: (TypeParameterDescriptor) -> TypeProjection): SimpleType {
+inline fun SimpleType.replaceArgumentsByParametersWith(replacement: (TypeParameterDescriptor) -> TypeProjection): SimpleType {
    if (constructor.parameters.isEmpty() || constructor.declarationDescriptor == null) return this

    val newArguments = constructor.parameters.map(replacement)
@@ -275,6 +276,11 @@ inline fun SimpleType.replaceArgumentsWith(replacement: (TypeParameterDescriptor
    return replace(newArguments)
 }

+inline fun SimpleType.replaceArgumentsByExistingArgumentsWith(replacement: (TypeArgumentMarker) -> TypeArgumentMarker): SimpleType {
+    if (arguments.isEmpty()) return this
+    return replace(newArguments = arguments.map { replacement(it) as TypeProjection })
+}
+
 fun KotlinType.containsTypeAliasParameters(): Boolean =
    contains {
        it.constructor.declarationDescriptor?.isTypeAliasParameter() ?: false
@@ -17,7 +17,6 @@ import org.jetbrains.kotlin.name.FqNameUnsafe
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.inference.CapturedType
-import org.jetbrains.kotlin.resolve.calls.inference.CapturedTypeConstructor
 import org.jetbrains.kotlin.resolve.constants.IntegerLiteralTypeConstructor
 import org.jetbrains.kotlin.resolve.descriptorUtil.*
 import org.jetbrains.kotlin.resolve.isInlineClass
@@ -25,10 +24,7 @@ import org.jetbrains.kotlin.resolve.substitutedUnderlyingType
 import org.jetbrains.kotlin.resolve.unsubstitutedUnderlyingType
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.model.*
-import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
-import org.jetbrains.kotlin.types.typeUtil.contains
-import org.jetbrains.kotlin.types.typeUtil.hasTypeParameterRecursiveBounds
-import org.jetbrains.kotlin.types.typeUtil.representativeUpperBound
+import org.jetbrains.kotlin.types.typeUtil.*
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.types.typeUtil.isSignedOrUnsignedNumberType as classicIsSignedOrUnsignedNumberType

@@ -210,6 +206,11 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
        return this.parameters[index]
    }

+    override fun TypeConstructorMarker.getParameters(): List<TypeParameterMarker> {
+        require(this is TypeConstructor, this::errorMessage)
+        return this.parameters
+    }
+
    override fun TypeConstructorMarker.supertypes(): Collection<KotlinTypeMarker> {
        require(this is TypeConstructor, this::errorMessage)
        return this.supertypes
@@ -240,9 +241,9 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
        return this.typeConstructor
    }

-    override fun TypeParameterMarker.hasRecursiveBounds(selfConstructor: TypeConstructorMarker): Boolean {
+    override fun TypeParameterMarker.hasRecursiveBounds(selfConstructor: TypeConstructorMarker?): Boolean {
        require(this is TypeParameterDescriptor, this::errorMessage)
-        require(selfConstructor is TypeConstructor, this::errorMessage)
+        require(selfConstructor is TypeConstructor?, this::errorMessage)

        return hasTypeParameterRecursiveBounds(this, selfConstructor)
    }
@@ -528,6 +529,12 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
        return this.replace(newArguments as List<TypeProjection>)
    }

+    override fun SimpleTypeMarker.replaceArguments(replacement: (TypeArgumentMarker) -> TypeArgumentMarker): SimpleTypeMarker {
+        require(this is SimpleType, this::errorMessage)
+        @Suppress("UNCHECKED_CAST")
+        return this.replaceArgumentsByExistingArgumentsWith(replacement)
+    }
+
    override fun DefinitelyNotNullTypeMarker.original(): SimpleTypeMarker {
        require(this is DefinitelyNotNullType, this::errorMessage)
        return this.original