[NI] Improve completing callable references with type variable as expected type

#KT-32462 Fixed

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
 import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
+import org.jetbrains.kotlin.builtins.isBuiltinFunctionalTypeOrSubtype
 import org.jetbrains.kotlin.resolve.calls.components.KotlinResolutionStatelessCallbacks
 import org.jetbrains.kotlin.resolve.calls.components.transformToResolvedLambda
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
@@ -121,51 +122,83 @@ class KotlinConstraintSystemCompleter(
     ): Boolean {
         val variable = variableForFixation.variable as TypeConstructor
         val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
+        val hasProperAtom = postponedArguments.any {
+            when (it) {
+                is LambdaWithTypeVariableAsExpectedTypeAtom,
+                is PostponedCallableReferenceAtom -> it.expectedType?.constructor == variable
+                else -> false
+            }
+        }
+
         if (
-            !postponedArguments.any { (it as? LambdaWithTypeVariableAsExpectedTypeAtom)?.expectedType?.constructor == variable } &&
+            !hasProperAtom &&
             variableForFixation.hasProperConstraint &&
             !variableForFixation.hasOnlyTrivialProperConstraint
         ) return false
 
         val postponedAtom = postponedArguments.firstOrNull() ?: return false
-        when (postponedAtom) {
-            is PostponedCallableReferenceAtom -> {
-                analyze(postponedAtom)
-            }
-            is LambdaWithTypeVariableAsExpectedTypeAtom -> {
-                var atomToAnalyze = postponedAtom
-                if (postponedAtom.atom.parametersTypes?.all { it != null } != true) {
-                    val functionalType = resultTypeResolver.findResultType(
-                        c,
-                        c.notFixedTypeVariables.getValue(variable),
-                        TypeVariableDirectionCalculator.ResolveDirection.TO_SUPERTYPE
-                    ) as KotlinType
-                    if (functionalType.isBuiltinFunctionalType) {
-                        val csBuilder = c as ConstraintSystemBuilder
-                        val builtIns = (variable as TypeVariableTypeConstructor).builtIns
-                        val returnVariable = TypeVariableForLambdaReturnType(postponedAtom.atom, builtIns, "_R")
-                        csBuilder.registerVariable(returnVariable)
-                        val expectedType = KotlinTypeFactory.simpleType(
-                            functionalType.annotations,
-                            functionalType.constructor,
-                            functionalType.arguments.dropLast(1) + returnVariable.defaultType.asTypeProjection(),
-                            functionalType.isMarkedNullable
-                        )
-                        csBuilder.addSubtypeConstraint(
-                            expectedType,
-                            variable.typeForTypeVariable(),
-                            ArgumentConstraintPosition(postponedAtom.atom)
-                        )
-                        atomToAnalyze = postponedAtom.transformToResolvedLambda(csBuilder, expectedType, returnVariable)
+
+        val builtIns = (variable as TypeVariableTypeConstructor).builtIns
+        val csBuilder = (c as NewConstraintSystemImpl).getBuilder()
+        val atomToAnalyze = when (postponedAtom) {
+            is PostponedCallableReferenceAtom -> postponedAtom.preparePostponedAtomWithTypeVariableAsExpectedType(
+                c, csBuilder, variable,
+                condition = { true },
+                isSuitable = KotlinType::isBuiltinFunctionalTypeOrSubtype,
+                typeVariableCreator = { TypeVariableForCallableReferenceReturnType(builtIns, "_Q") },
+                newAtomCreator = { returnVariable, expectedType ->
+                    CallableReferenceWithTypeVariableAsExpectedTypeAtom(postponedAtom.atom, expectedType, returnVariable).also {
+                        postponedAtom.setAnalyzedResults(null, listOf(it))
                     }
                 }
-                analyze(atomToAnalyze)
-            }
+            )
+            is LambdaWithTypeVariableAsExpectedTypeAtom -> postponedAtom.preparePostponedAtomWithTypeVariableAsExpectedType(
+                c, csBuilder, variable,
+                condition = { it.atom.parametersTypes?.all { type -> type != null } != true },
+                isSuitable = KotlinType::isBuiltinFunctionalType,
+                typeVariableCreator = { TypeVariableForLambdaReturnType(postponedAtom.atom, builtIns, "_R") },
+                newAtomCreator = { returnVariable, expectedType ->
+                    postponedAtom.transformToResolvedLambda(csBuilder, expectedType, returnVariable)
+                }
+            )
             else -> return false
         }
+        analyze(atomToAnalyze)
         return true
     }
 
+    private inline fun <T : PostponedResolvedAtom, V : NewTypeVariable> T.preparePostponedAtomWithTypeVariableAsExpectedType(
+        c: Context,
+        csBuilder: ConstraintSystemBuilder,
+        variable: TypeConstructor,
+        condition: (T) -> Boolean,
+        isSuitable: KotlinType.() -> Boolean,
+        typeVariableCreator: () -> V,
+        newAtomCreator: (V, SimpleType) -> PostponedResolvedAtom
+    ): PostponedResolvedAtom {
+        if (!condition(this)) return this
+        val functionalType = resultTypeResolver.findResultType(
+            c,
+            c.notFixedTypeVariables.getValue(variable),
+            TypeVariableDirectionCalculator.ResolveDirection.TO_SUPERTYPE
+        ) as KotlinType
+        if (!functionalType.isSuitable()) return this
+        val returnVariable = typeVariableCreator()
+        csBuilder.registerVariable(returnVariable)
+        val expectedType = KotlinTypeFactory.simpleType(
+            functionalType.annotations,
+            functionalType.constructor,
+            functionalType.arguments.dropLast(1) + returnVariable.defaultType.asTypeProjection(),
+            functionalType.isMarkedNullable
+        )
+        csBuilder.addSubtypeConstraint(
+            expectedType,
+            variable.typeForTypeVariable(),
+            ArgumentConstraintPosition(atom as KotlinCallArgument)
+        )
+        return newAtomCreator(returnVariable, expectedType)
+    }
+
     // true if we do analyze
     private fun analyzePostponeArgumentIfPossible(
         c: Context,
@@ -208,6 +241,7 @@ class KotlinConstraintSystemCompleter(
         fun ResolvedAtom.process(to: LinkedHashSet<TypeConstructor>) {
             val typeVariables = when (this) {
                 is ResolvedCallAtom -> freshVariablesSubstitutor.freshVariables
+                is CallableReferenceWithTypeVariableAsExpectedTypeAtom -> listOfNotNull(typeVariableForReturnType)
                 is ResolvedCallableReferenceAtom -> candidate?.freshSubstitutor?.freshVariables.orEmpty()
                 is ResolvedLambdaAtom -> listOfNotNull(typeVariableForLambdaReturnType)
                 else -> emptyList()

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/model/TypeVariable.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.resolve.calls.model.CallableReferenceKotlinCallArgument
 import org.jetbrains.kotlin.resolve.calls.model.LambdaKotlinCallArgument
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasOnlyInputTypesAnnotation
@@ -82,3 +83,10 @@ class TypeVariableForLambdaReturnType(
 ) : NewTypeVariable(builtIns, name) {
     override fun hasOnlyInputTypesAnnotation(): Boolean = false
 }
+
+class TypeVariableForCallableReferenceReturnType(
+    builtIns: KotlinBuiltIns,
+    name: String
+) : NewTypeVariable(builtIns, name) {
+    override fun hasOnlyInputTypesAnnotation(): Boolean = false
+}

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolutionAtoms.kt

@@ -14,7 +14,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.components.FreshVariableNewT
 import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintError
-import org.jetbrains.kotlin.resolve.calls.inference.model.NewTypeVariable
+import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableForCallableReferenceReturnType
 import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableForLambdaReturnType
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.types.KotlinType
@@ -91,11 +91,12 @@ interface PostponedResolvedAtomMarker {
 sealed class PostponedResolvedAtom : ResolvedAtom(), PostponedResolvedAtomMarker {
     abstract override val inputTypes: Collection<UnwrappedType>
     abstract override val outputType: UnwrappedType?
+    abstract val expectedType: UnwrappedType?
 }
 
 class LambdaWithTypeVariableAsExpectedTypeAtom(
     override val atom: LambdaKotlinCallArgument,
-    val expectedType: UnwrappedType
+    override val expectedType: UnwrappedType
 ) : PostponedResolvedAtom() {
     override val inputTypes: Collection<UnwrappedType> get() = listOf(expectedType)
     override val outputType: UnwrappedType? get() = null
@@ -112,7 +113,7 @@ class ResolvedLambdaAtom(
     val parameters: List<UnwrappedType>,
     val returnType: UnwrappedType,
     val typeVariableForLambdaReturnType: TypeVariableForLambdaReturnType?,
-    val expectedType: UnwrappedType?
+    override val expectedType: UnwrappedType?
 ) : PostponedResolvedAtom() {
     lateinit var resultArguments: List<KotlinCallArgument>
         private set
@@ -131,7 +132,7 @@ class ResolvedLambdaAtom(
 
 abstract class ResolvedCallableReferenceAtom(
     override val atom: CallableReferenceKotlinCallArgument,
-    val expectedType: UnwrappedType?
+    override val expectedType: UnwrappedType?
 ) : PostponedResolvedAtom() {
     var candidate: CallableReferenceCandidate? = null
         private set
@@ -157,10 +158,10 @@ class EagerCallableReferenceAtom(
     fun transformToPostponed(): PostponedCallableReferenceAtom = PostponedCallableReferenceAtom(this)
 }
 
-class PostponedCallableReferenceAtom(
-    eagerCallableReferenceAtom: EagerCallableReferenceAtom
-) : ResolvedCallableReferenceAtom(eagerCallableReferenceAtom.atom, eagerCallableReferenceAtom.expectedType) {
-
+sealed class AbstractPostponedCallableReferenceAtom(
+    atom: CallableReferenceKotlinCallArgument,
+    expectedType: UnwrappedType?
+) : ResolvedCallableReferenceAtom(atom, expectedType) {
     override val inputTypes: Collection<UnwrappedType>
         get() = extractInputOutputTypesFromCallableReferenceExpectedType(expectedType)?.inputTypes ?: listOfNotNull(expectedType)
 
@@ -168,6 +169,16 @@ class PostponedCallableReferenceAtom(
         get() = extractInputOutputTypesFromCallableReferenceExpectedType(expectedType)?.outputType
 }
 
+class CallableReferenceWithTypeVariableAsExpectedTypeAtom(
+    atom: CallableReferenceKotlinCallArgument,
+    expectedType: UnwrappedType?,
+    val typeVariableForReturnType: TypeVariableForCallableReferenceReturnType?
+) : AbstractPostponedCallableReferenceAtom(atom, expectedType)
+
+class PostponedCallableReferenceAtom(
+    eagerCallableReferenceAtom: EagerCallableReferenceAtom
+) : AbstractPostponedCallableReferenceAtom(eagerCallableReferenceAtom.atom, eagerCallableReferenceAtom.expectedType)
+
 class ResolvedCollectionLiteralAtom(
     override val atom: CollectionLiteralKotlinCallArgument,
     val expectedType: UnwrappedType?