[NI] Improve completing calls with multiple postponed arguments

#KT-27999 Fixed
#KT-30244 Fixed
#KT-31102 Fixed

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

@@ -57,13 +57,15 @@ private fun preprocessLambdaArgument(
     csBuilder: ConstraintSystemBuilder,
     argument: LambdaKotlinCallArgument,
     expectedType: UnwrappedType?,
-    forceResolution: Boolean = false
+    forceResolution: Boolean = false,
+    returnTypeVariable: TypeVariableForLambdaReturnType? = null
 ): ResolvedAtom {
     if (expectedType != null && !forceResolution && csBuilder.isTypeVariable(expectedType)) {
         return LambdaWithTypeVariableAsExpectedTypeAtom(argument, expectedType)
     }
 
-    val resolvedArgument = extractLambdaInfoFromFunctionalType(expectedType, argument) ?: extraLambdaInfo(expectedType, argument, csBuilder)
+    val resolvedArgument = extractLambdaInfoFromFunctionalType(expectedType, argument, returnTypeVariable)
+        ?: extraLambdaInfo(expectedType, argument, csBuilder)
 
     if (expectedType != null) {
         val lambdaType = createFunctionType(
@@ -110,7 +112,11 @@ private fun extraLambdaInfo(
     )
 }
 
-private fun extractLambdaInfoFromFunctionalType(expectedType: UnwrappedType?, argument: LambdaKotlinCallArgument): ResolvedLambdaAtom? {
+private fun extractLambdaInfoFromFunctionalType(
+    expectedType: UnwrappedType?,
+    argument: LambdaKotlinCallArgument,
+    returnTypeVariable: TypeVariableForLambdaReturnType? = null
+): ResolvedLambdaAtom? {
     if (expectedType == null || !expectedType.isBuiltinFunctionalType) return null
     val parameters = extractLambdaParameters(expectedType, argument)
 
@@ -124,7 +130,7 @@ private fun extractLambdaInfoFromFunctionalType(expectedType: UnwrappedType?, ar
         receiverType,
         parameters,
         returnType,
-        typeVariableForLambdaReturnType = null,
+        typeVariableForLambdaReturnType = returnTypeVariable,
         expectedType = expectedType
     )
 }
@@ -141,9 +147,20 @@ private fun extractLambdaParameters(expectedType: UnwrappedType, argument: Lambd
     }
 }
 
-fun LambdaWithTypeVariableAsExpectedTypeAtom.transformToResolvedLambda(csBuilder: ConstraintSystemBuilder): ResolvedLambdaAtom {
-    val fixedExpectedType = (csBuilder.buildCurrentSubstitutor() as NewTypeSubstitutor).safeSubstitute(expectedType)
-    val resolvedLambdaAtom = preprocessLambdaArgument(csBuilder, atom, fixedExpectedType, forceResolution = true) as ResolvedLambdaAtom
+fun LambdaWithTypeVariableAsExpectedTypeAtom.transformToResolvedLambda(
+    csBuilder: ConstraintSystemBuilder,
+    expectedType: UnwrappedType? = null,
+    returnTypeVariable: TypeVariableForLambdaReturnType? = null
+): ResolvedLambdaAtom {
+    val fixedExpectedType = (csBuilder.buildCurrentSubstitutor() as NewTypeSubstitutor)
+        .safeSubstitute(expectedType ?: this.expectedType)
+    val resolvedLambdaAtom = preprocessLambdaArgument(
+        csBuilder,
+        atom,
+        fixedExpectedType,
+        forceResolution = true,
+        returnTypeVariable
+    ) as ResolvedLambdaAtom
 
     setAnalyzed(resolvedLambdaAtom)
 

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

@@ -5,19 +5,18 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
+import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
 import org.jetbrains.kotlin.resolve.calls.components.KotlinResolutionStatelessCallbacks
-import org.jetbrains.kotlin.resolve.calls.inference.model.NotEnoughInformationForTypeParameter
-import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableFromCallableDescriptor
-import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
+import org.jetbrains.kotlin.resolve.calls.components.transformToResolvedLambda
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
+import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.*
-import org.jetbrains.kotlin.types.ErrorUtils
-import org.jetbrains.kotlin.types.TypeConstructor
-import org.jetbrains.kotlin.types.UnwrappedType
+import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
+import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.utils.addIfNotNull
-import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class KotlinConstraintSystemCompleter(
@@ -80,9 +79,11 @@ class KotlinConstraintSystemCompleter(
                     c, allTypeVariables, postponedKtPrimitives, completionMode, topLevelType
                 ) ?: break
 
-            if (shouldForceCallableReferenceOrLambdaResolution(completionMode, variableForFixation)) {
-                if (forcePostponedAtomResolution<PostponedCallableReferenceAtom>(topLevelAtoms, analyze)) continue
-                if (forcePostponedAtomResolution<LambdaWithTypeVariableAsExpectedTypeAtom>(topLevelAtoms, analyze)) continue
+            if (
+                completionMode == ConstraintSystemCompletionMode.FULL &&
+                resolveLambdaOrCallableReferenceWithTypeVariableAsExpectedType(c, variableForFixation, topLevelAtoms, analyze)
+            ) {
+                continue
             }
 
             if (variableForFixation.hasProperConstraint || completionMode == ConstraintSystemCompletionMode.FULL) {
@@ -109,12 +110,60 @@ class KotlinConstraintSystemCompleter(
         }
     }
 
-    private fun shouldForceCallableReferenceOrLambdaResolution(
-        completionMode: ConstraintSystemCompletionMode,
-        variableForFixation: VariableFixationFinder.VariableForFixation
+    /*
+     * returns true -> analyzed
+     */
+    private fun resolveLambdaOrCallableReferenceWithTypeVariableAsExpectedType(
+        c: Context,
+        variableForFixation: VariableFixationFinder.VariableForFixation,
+        topLevelAtoms: List<ResolvedAtom>,
+        analyze: (PostponedResolvedAtom) -> Unit
     ): Boolean {
-        if (completionMode == ConstraintSystemCompletionMode.PARTIAL) return false
-        return !variableForFixation.hasProperConstraint || variableForFixation.hasOnlyTrivialProperConstraint
+        val variable = variableForFixation.variable as TypeConstructor
+        val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
+        if (
+            !postponedArguments.any { (it as? LambdaWithTypeVariableAsExpectedTypeAtom)?.expectedType?.constructor == variable } &&
+            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)
+                    }
+                }
+                analyze(atomToAnalyze)
+            }
+            else -> return false
+        }
+        return true
     }
 
     // true if we do analyze
@@ -132,16 +181,6 @@ class KotlinConstraintSystemCompleter(
         return false
     }
 
-    // true if we find some callable reference and run resolution for it. Note that such resolution can be unsuccessful
-    private inline fun <reified T : PostponedResolvedAtom> forcePostponedAtomResolution(
-        topLevelAtoms: List<ResolvedAtom>,
-        analyze: (PostponedResolvedAtom) -> Unit
-    ): Boolean {
-        val postponedArgument = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms).firstIsInstanceOrNull<T>() ?: return false
-        analyze(postponedArgument)
-        return true
-    }
-
     private fun getOrderedNotAnalyzedPostponedArguments(topLevelAtoms: List<ResolvedAtom>): List<PostponedResolvedAtom> {
         fun ResolvedAtom.process(to: MutableList<PostponedResolvedAtom>) {
             to.addIfNotNull(this.safeAs<PostponedResolvedAtom>()?.takeUnless { it.analyzed })
@@ -190,7 +229,7 @@ class KotlinConstraintSystemCompleter(
         }
 
         assert(result.size == c.notFixedTypeVariables.size) {
-            val notFoundTypeVariables = c.notFixedTypeVariables.keys.toMutableSet().removeAll(result)
+            val notFoundTypeVariables = c.notFixedTypeVariables.keys.toMutableSet().apply { removeAll(result) }
             "Not all type variables found: $notFoundTypeVariables"
         }
 

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

@@ -55,16 +55,20 @@ sealed class NewTypeVariable(builtIns: KotlinBuiltIns, name: String) : TypeVaria
 
     // member scope is used if we have receiver with type TypeVariable(T)
     // todo add to member scope methods from supertypes for type variable
-    val defaultType: SimpleType = KotlinTypeFactory.simpleTypeWithNonTrivialMemberScope(
-        Annotations.EMPTY, freshTypeConstructor, arguments = emptyList(),
-        nullable = false, memberScope = builtIns.any.unsubstitutedMemberScope
-    )
-
+    val defaultType: SimpleType = freshTypeConstructor.typeForTypeVariable()
     abstract fun hasOnlyInputTypesAnnotation(): Boolean
 
     override fun toString() = freshTypeConstructor.toString()
 }
 
+fun TypeConstructor.typeForTypeVariable(): SimpleType {
+    require(this is TypeVariableTypeConstructor)
+    return KotlinTypeFactory.simpleTypeWithNonTrivialMemberScope(
+        Annotations.EMPTY, this, arguments = emptyList(),
+        nullable = false, memberScope = builtIns.any.unsubstitutedMemberScope
+    )
+}
+
 class TypeVariableFromCallableDescriptor(
     val originalTypeParameter: TypeParameterDescriptor
 ) : NewTypeVariable(originalTypeParameter.builtIns, originalTypeParameter.name.identifier) {

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

@@ -14,6 +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.TypeVariableForLambdaReturnType
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.types.KotlinType