[NI] Don't process lambda until expected type will be fixed

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

@@ -45,52 +45,80 @@ fun resolveKtPrimitive(
 private fun preprocessLambdaArgument(
         csBuilder: ConstraintSystemBuilder,
         argument: LambdaKotlinCallArgument,
-        expectedType: UnwrappedType?
+        expectedType: UnwrappedType?,
+        forceResolution: Boolean = false
 ): ResolvedAtom {
+    if (expectedType != null && !forceResolution && csBuilder.isTypeVariable(expectedType)) {
+        return LambdaWithTypeVariableAsExpectedTypeAtom(argument, expectedType)
+    }
+
+    val resolvedArgument = extractLambdaInfoFromFunctionalType(expectedType, argument) ?: extraLambdaInfo(expectedType, argument, csBuilder)
+
+    if (expectedType != null) {
+        val lambdaType = createFunctionType(csBuilder.builtIns, Annotations.EMPTY, resolvedArgument.receiver,
+                                            resolvedArgument.parameters, null, resolvedArgument.returnType, resolvedArgument.isSuspend)
+        csBuilder.addSubtypeConstraint(lambdaType, expectedType, ArgumentConstraintPosition(argument))
+    }
+
+    return resolvedArgument
+}
+
+private fun extraLambdaInfo(
+        expectedType: UnwrappedType?,
+        argument: LambdaKotlinCallArgument,
+        csBuilder: ConstraintSystemBuilder
+): ResolvedLambdaAtom {
     val builtIns = csBuilder.builtIns
     val isSuspend = expectedType?.isSuspendFunctionType ?: false
 
-    val receiverType: UnwrappedType? // null means that there is no receiver
-    val parameters: List<UnwrappedType>
-    val returnType: UnwrappedType
+    val isFunctionSupertype = expectedType != null && KotlinBuiltIns.isNotNullOrNullableFunctionSupertype(expectedType)
+    val argumentAsFunctionExpression = argument.safeAs<FunctionExpression>()
+
     val typeVariable = TypeVariableForLambdaReturnType(argument, builtIns, "_L")
 
-    if (expectedType?.isBuiltinFunctionalType == true) {
-        receiverType = if (argument is FunctionExpression) argument.receiverType else expectedType.getReceiverTypeFromFunctionType()?.unwrap()
-
-        val expectedParameters = expectedType.getValueParameterTypesFromFunctionType()
-        parameters = if (argument.parametersTypes != null) {
-            argument.parametersTypes!!.mapIndexed {
-                index, type ->
-                type ?: expectedParameters.getOrNull(index)?.type?.unwrap() ?: builtIns.nullableAnyType
-            }
-        }
-        else {
-            // lambda without explicit parameters: { }
-            expectedParameters.map { it.type.unwrap() }
-        }
-        returnType = argument.safeAs<FunctionExpression>()?.returnType ?: expectedType.getReturnTypeFromFunctionType().unwrap()
-    }
-    else {
-        val isFunctionSupertype = expectedType != null && KotlinBuiltIns.isNotNullOrNullableFunctionSupertype(expectedType)
-        receiverType = argument.safeAs<FunctionExpression>()?.receiverType
-        parameters = argument.parametersTypes?.map { it ?: builtIns.nothingType } ?: emptyList()
-        returnType = argument.safeAs<FunctionExpression>()?.returnType ?:
+    val receiverType = argumentAsFunctionExpression?.receiverType
+    val returnType = argumentAsFunctionExpression?.returnType ?:
                      expectedType?.arguments?.singleOrNull()?.type?.unwrap()?.takeIf { isFunctionSupertype } ?:
                      typeVariable.defaultType
 
-        // what about case where expected type is type variable? In old TY such cases was not supported. => do nothing for now. todo design
-    }
+    val parameters = argument.parametersTypes?.map { it ?: builtIns.nothingType } ?: emptyList()
 
     val newTypeVariableUsed = returnType == typeVariable.defaultType
     if (newTypeVariableUsed) csBuilder.registerVariable(typeVariable)
 
-    if (expectedType != null) {
-        val lambdaType = createFunctionType(returnType.builtIns, Annotations.EMPTY, receiverType, parameters, null, returnType, isSuspend)
-        csBuilder.addSubtypeConstraint(lambdaType, expectedType, ArgumentConstraintPosition(argument))
+    return ResolvedLambdaAtom(argument, isSuspend, receiverType, parameters, returnType, typeVariable.takeIf { newTypeVariableUsed })
+}
+
+private fun extractLambdaInfoFromFunctionalType(expectedType: UnwrappedType?, argument: LambdaKotlinCallArgument): ResolvedLambdaAtom? {
+    if (expectedType == null || !expectedType.isBuiltinFunctionalType) return null
+    val parameters = extractLambdaParameters(expectedType, argument)
+
+    val argumentAsFunctionExpression = argument.safeAs<FunctionExpression>()
+    val receiverType = argumentAsFunctionExpression?.receiverType ?: expectedType.getReceiverTypeFromFunctionType()?.unwrap()
+    val returnType = argumentAsFunctionExpression?.returnType ?: expectedType.getReturnTypeFromFunctionType().unwrap()
+
+    return ResolvedLambdaAtom(argument, expectedType.isSuspendFunctionType, receiverType, parameters, returnType, typeVariableForLambdaReturnType = null)
+}
+
+private fun extractLambdaParameters(expectedType: UnwrappedType, argument: LambdaKotlinCallArgument): List<UnwrappedType> {
+    val parametersTypes = argument.parametersTypes
+    val expectedParameters = expectedType.getValueParameterTypesFromFunctionType()
+    if (parametersTypes == null) {
+        return expectedParameters.map { it.type.unwrap() }
     }
 
-    return ResolvedLambdaAtom(argument, isSuspend, receiverType, parameters, returnType, typeVariable.takeIf { newTypeVariableUsed })
+    return parametersTypes.mapIndexed { index, type ->
+        type ?: expectedParameters.getOrNull(index)?.type?.unwrap() ?: expectedType.builtIns.nullableAnyType
+    }
+}
+
+fun LambdaWithTypeVariableAsExpectedTypeAtom.transformToResolvedLambda(csBuilder: ConstraintSystemBuilder): ResolvedLambdaAtom {
+    val fixedExpectedType = csBuilder.buildCurrentSubstitutor().safeSubstitute(expectedType)
+    val resolvedLambdaAtom = preprocessLambdaArgument(csBuilder, atom, fixedExpectedType, forceResolution = true) as ResolvedLambdaAtom
+
+    setAnalyzed(resolvedLambdaAtom)
+
+    return resolvedLambdaAtom
 }
 
 private fun preprocessCallableReference(

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

@@ -42,6 +42,7 @@ class PostponedArgumentsAnalyzer(
     fun analyze(c: Context, resolutionCallbacks: KotlinResolutionCallbacks, argument: ResolvedAtom) {
         when (argument) {
             is ResolvedLambdaAtom -> analyzeLambda(c, resolutionCallbacks, argument)
+            is LambdaWithTypeVariableAsExpectedTypeAtom -> analyzeLambda(c, resolutionCallbacks, argument.transformToResolvedLambda(c.getBuilder()))
             is ResolvedCallableReferenceAtom -> callableReferenceResolver.processCallableReferenceArgument(c.getBuilder(), argument)
             is ResolvedCollectionLiteralAtom -> TODO("Not supported")
             else -> error("Unexpected resolved primitive: ${argument.javaClass.canonicalName}")

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

@@ -37,6 +37,7 @@ interface ConstraintSystemOperation {
     fun addEqualityConstraint(a: UnwrappedType, b: UnwrappedType, position: ConstraintPosition)
 
     fun isProperType(type: UnwrappedType): Boolean
+    fun isTypeVariable(type: UnwrappedType): Boolean
 
     fun getProperSuperTypeConstructors(type: UnwrappedType): List<TypeConstructor>
 }

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

@@ -61,8 +61,9 @@ class KotlinConstraintSystemCompleter(
             val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
                     c, allTypeVariables, postponedKtPrimitives, completionMode, topLevelType)
 
-            if (shouldWeForceCallableReferenceResolution(completionMode, variableForFixation)) {
-                if (forceCallableReferenceResolution(topLevelPrimitive, analyze)) continue
+            if (shouldForceCallableReferenceOrLambdaResolution(completionMode, variableForFixation)) {
+                if (forcePostponedAtomResolution<ResolvedCallableReferenceAtom>(topLevelPrimitive, analyze)) continue
+                if (forcePostponedAtomResolution<LambdaWithTypeVariableAsExpectedTypeAtom>(topLevelPrimitive, analyze)) continue
             }
 
             if (variableForFixation != null) {
@@ -86,7 +87,7 @@ class KotlinConstraintSystemCompleter(
         }
     }
 
-    private fun shouldWeForceCallableReferenceResolution(
+    private fun shouldForceCallableReferenceOrLambdaResolution(
             completionMode: ConstraintSystemCompletionMode,
             variableForFixation: VariableFixationFinder.VariableForFixation?
     ): Boolean {
@@ -108,11 +109,12 @@ class KotlinConstraintSystemCompleter(
     }
 
     // true if we find some callable reference and run resolution for it. Note that such resolution can be unsuccessful
-    private fun forceCallableReferenceResolution(topLevelPrimitive: ResolvedAtom, analyze: (PostponedResolvedAtom) -> Unit): Boolean {
-        val callableReferenceArgument = getOrderedNotAnalyzedPostponedArguments(topLevelPrimitive).
-                firstIsInstanceOrNull<ResolvedCallableReferenceAtom>() ?: return false
-
-        analyze(callableReferenceArgument)
+    private inline fun <reified T : PostponedResolvedAtom> forcePostponedAtomResolution(
+            topLevelPrimitive: ResolvedAtom,
+            analyze: (PostponedResolvedAtom) -> Unit
+    ): Boolean {
+        val postponedArgument = getOrderedNotAnalyzedPostponedArguments(topLevelPrimitive).firstIsInstanceOrNull<T>() ?: return false
+        analyze(postponedArgument)
         return true
     }
 

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

@@ -174,6 +174,11 @@ class NewConstraintSystemImpl(
         }
     }
 
+    override fun isTypeVariable(type: UnwrappedType): Boolean {
+        checkState(State.BUILDING, State.COMPLETION, State.TRANSACTION)
+        return notFixedTypeVariables.containsKey(type.constructor)
+    }
+
     // ConstraintInjector.Context
     override val allTypeVariables: Map<TypeConstructor, NewTypeVariable> get() {
         checkState(State.BUILDING, State.COMPLETION, State.TRANSACTION)

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

@@ -82,6 +82,18 @@ sealed class PostponedResolvedAtom : ResolvedAtom() {
     abstract val outputType: UnwrappedType?
 }
 
+class LambdaWithTypeVariableAsExpectedTypeAtom(
+        override val atom: LambdaKotlinCallArgument,
+        val expectedType: UnwrappedType
+) : PostponedResolvedAtom() {
+    override val inputTypes: Collection<UnwrappedType> get() = listOf(expectedType)
+    override val outputType: UnwrappedType? get() = null
+
+    fun setAnalyzed(resolvedLambdaAtom: ResolvedLambdaAtom) {
+        setAnalyzedResults(listOf(resolvedLambdaAtom), listOf())
+    }
+}
+
 class ResolvedLambdaAtom(
         override val atom: LambdaKotlinCallArgument,
         val isSuspend: Boolean,

compiler/testData/codegen/box/regressions/lambdaPostponeConstruction.kt

@@ -0,0 +1,14 @@
+class MyList<T>
+
+operator fun <T> MyList<T>.plusAssign(element: T) {}
+
+val listOfFunctions = MyList<(Int) -> Int>()
+
+fun foo() {
+    listOfFunctions.plusAssign({ it -> it })
+    listOfFunctions += { it -> it }
+}
+
+fun box(): String {
+    return "OK"
+}

compiler/testData/codegen/box/regressions/lambdaWrongReturnType.kt

@@ -0,0 +1,7 @@
+fun test() = foo({ line: String -> line })
+
+fun <T> foo(x: T): T = TODO()
+
+fun box(): String {
+    return "OK"
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -17018,6 +17018,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("lambdaPostponeConstruction.kt")
+        public void testLambdaPostponeConstruction() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/lambdaPostponeConstruction.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lambdaWrongReturnType.kt")
+        public void testLambdaWrongReturnType() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/lambdaWrongReturnType.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("nestedIntersection.kt")
         public void testNestedIntersection() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/nestedIntersection.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -17018,6 +17018,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("lambdaPostponeConstruction.kt")
+        public void testLambdaPostponeConstruction() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/lambdaPostponeConstruction.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lambdaWrongReturnType.kt")
+        public void testLambdaWrongReturnType() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/lambdaWrongReturnType.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("nestedIntersection.kt")
         public void testNestedIntersection() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/nestedIntersection.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -17018,6 +17018,18 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("lambdaPostponeConstruction.kt")
+        public void testLambdaPostponeConstruction() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/lambdaPostponeConstruction.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lambdaWrongReturnType.kt")
+        public void testLambdaWrongReturnType() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/lambdaWrongReturnType.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("nestedIntersection.kt")
         public void testNestedIntersection() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/nestedIntersection.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -20798,6 +20798,18 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("lambdaPostponeConstruction.kt")
+        public void testLambdaPostponeConstruction() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/lambdaPostponeConstruction.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lambdaWrongReturnType.kt")
+        public void testLambdaWrongReturnType() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/lambdaWrongReturnType.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("nestedIntersection.kt")
         public void testNestedIntersection() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/regressions/nestedIntersection.kt");