Support type inference for self type materialization calls

core/descriptors/src/org/jetbrains/kotlin/types/StarProjectionImpl.kt

@@ -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

core/descriptors/src/org/jetbrains/kotlin/types/TypeUtils.kt

@@ -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

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

@@ -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