Reuse revised variables during lambda analysis against type variables

#KT-41400 Fixed

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

@@ -8,10 +8,7 @@ package org.jetbrains.kotlin.resolve.calls.inference
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
-import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.model.*
-import org.jetbrains.kotlin.types.model.TypeConstructorMarker
-import org.jetbrains.kotlin.types.model.TypeSubstitutorMarker
-import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 interface ConstraintSystemOperation {
     val hasContradiction: Boolean
@@ -20,6 +17,12 @@ interface ConstraintSystemOperation {
     fun unmarkPostponedVariable(variable: TypeVariableMarker)
     fun removePostponedVariables()
 
+    fun getRevisedVariableForParameter(expectedType: TypeVariableMarker, index: Int): TypeVariableMarker?
+    fun getRevisedVariableForReturnType(expectedType: TypeVariableMarker): TypeVariableMarker?
+
+    fun putRevisedVariableForParameter(expectedType: TypeVariableMarker, index: Int, newVariable: TypeVariableMarker)
+    fun putRevisedVariableForReturnType(expectedType: TypeVariableMarker, newVariable: TypeVariableMarker)
+
     fun addSubtypeConstraint(lowerType: KotlinTypeMarker, upperType: KotlinTypeMarker, position: ConstraintPosition)
     fun addEqualityConstraint(a: KotlinTypeMarker, b: KotlinTypeMarker, position: ConstraintPosition)
 

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

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.resolve.calls.inference.model
 
 import org.jetbrains.kotlin.types.AbstractTypeChecker
-import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.model.*
-import org.jetbrains.kotlin.types.model.TypeCheckerProviderContext
-import org.jetbrains.kotlin.types.model.TypeConstructorMarker
-import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 /**
  * Every type variable can be in the following states:
@@ -42,6 +39,8 @@ interface ConstraintStorage {
     val hasContradiction: Boolean
     val fixedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker>
     val postponedTypeVariables: List<TypeVariableMarker>
+    val revisedVariablesForParameters: Map<Pair<TypeVariableMarker, Int>, TypeVariableMarker>
+    val revisedReturnTypes: Map<TypeVariableMarker, TypeVariableMarker>
 
     object Empty : ConstraintStorage {
         override val allTypeVariables: Map<TypeConstructorMarker, TypeVariableMarker> get() = emptyMap()
@@ -52,6 +51,8 @@ interface ConstraintStorage {
         override val hasContradiction: Boolean get() = false
         override val fixedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker> get() = emptyMap()
         override val postponedTypeVariables: List<TypeVariableMarker> get() = emptyList()
+        override val revisedVariablesForParameters: Map<Pair<TypeVariableMarker, Int>, TypeVariableMarker> = emptyMap()
+        override val revisedReturnTypes: Map<TypeVariableMarker, TypeVariableMarker> = emptyMap()
     }
 }
 

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

@@ -202,4 +202,6 @@ internal class MutableConstraintStorage : ConstraintStorage {
     override val hasContradiction: Boolean get() = errors.any { !it.applicability.isSuccess }
     override val fixedTypeVariables: MutableMap<TypeConstructorMarker, KotlinTypeMarker> = LinkedHashMap()
     override val postponedTypeVariables: MutableList<TypeVariableMarker> = SmartList()
+    override val revisedVariablesForParameters: MutableMap<Pair<TypeVariableMarker, Int>, TypeVariableMarker> = LinkedHashMap()
+    override val revisedReturnTypes: MutableMap<TypeVariableMarker, TypeVariableMarker> = LinkedHashMap()
 }

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

@@ -107,6 +107,22 @@ class NewConstraintSystemImpl(
         storage.postponedTypeVariables.clear()
     }
 
+    override fun getRevisedVariableForParameter(expectedType: TypeVariableMarker, index: Int): TypeVariableMarker? {
+        return storage.revisedVariablesForParameters[expectedType to index]
+    }
+
+    override fun getRevisedVariableForReturnType(expectedType: TypeVariableMarker): TypeVariableMarker? {
+        return storage.revisedReturnTypes[expectedType]
+    }
+
+    override fun putRevisedVariableForParameter(expectedType: TypeVariableMarker, index: Int, newVariable: TypeVariableMarker) {
+        storage.revisedVariablesForParameters[expectedType to index] = newVariable
+    }
+
+    override fun putRevisedVariableForReturnType(expectedType: TypeVariableMarker, newVariable: TypeVariableMarker) {
+        storage.revisedReturnTypes[expectedType] = newVariable
+    }
+
     override fun addSubtypeConstraint(lowerType: KotlinTypeMarker, upperType: KotlinTypeMarker, position: ConstraintPosition) =
         constraintInjector.addInitialSubtypeConstraint(
             apply { checkState(State.BUILDING, State.COMPLETION, State.TRANSACTION) },

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

@@ -155,6 +155,12 @@ class PostponedArgumentInputTypesResolver(
         val expectedType = argument.expectedType
             ?: throw IllegalStateException("Postponed argument's expected type must not be null")
 
+        val variable = getBuilder().currentStorage().allTypeVariables[expectedType.constructor]
+        if (variable != null) {
+            val revisedVariableForReturnType = getBuilder().getRevisedVariableForReturnType(variable)
+            if (revisedVariableForReturnType != null) return revisedVariableForReturnType as NewTypeVariable
+        }
+
         return when (argument) {
             is LambdaWithTypeVariableAsExpectedTypeAtom -> TypeVariableForLambdaReturnType(
                 expectedType.builtIns,
@@ -165,7 +171,11 @@ class PostponedArgumentInputTypesResolver(
                 TYPE_VARIABLE_NAME_FOR_CR_RETURN_TYPE
             )
             else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
-        }.also { getBuilder().registerVariable(it) }
+        }.also {
+            if (variable != null) getBuilder().putRevisedVariableForReturnType(variable, it)
+
+            getBuilder().registerVariable(it)
+        }
     }
 
     private fun Context.createTypeVariableForParameterType(
@@ -175,6 +185,12 @@ class PostponedArgumentInputTypesResolver(
         val expectedType = argument.expectedType
             ?: throw IllegalStateException("Postponed argument's expected type must not be null")
 
+        val variable = getBuilder().currentStorage().allTypeVariables[expectedType.constructor]
+        if (variable != null) {
+            val revisedVariableForParameter = getBuilder().getRevisedVariableForParameter(variable, index)
+            if (revisedVariableForParameter != null) return revisedVariableForParameter as NewTypeVariable
+        }
+
         return when (argument) {
             is LambdaWithTypeVariableAsExpectedTypeAtom -> TypeVariableForLambdaParameterType(
                 argument.atom,
@@ -187,7 +203,11 @@ class PostponedArgumentInputTypesResolver(
                 TYPE_VARIABLE_NAME_PREFIX_FOR_CR_PARAMETER_TYPE + (index + 1)
             )
             else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
-        }.also { getBuilder().registerVariable(it) }
+        }.also {
+            if (variable != null) getBuilder().putRevisedVariableForParameter(variable, index, it)
+
+            getBuilder().registerVariable(it)
+        }
     }
 
     private fun Context.createTypeVariablesForParameters(

compiler/testData/codegen/box/inference/lambdasWithExtensionFunctionType.kt

@@ -18,7 +18,7 @@ interface Canvas {
     fun circle(circle: Circle, stroke: Stroke, fill: Fill?)
 }
 
-fun test() {
+fun test1() {
     val rect = Rectangle(100, 100)
     val circle = Circle(100)
 
@@ -34,9 +34,34 @@ fun test() {
     }
 }
 
+fun test2() {
+    val rect = Rectangle(100, 100)
+    val circle = Circle(100)
+
+    val l: List<Canvas.(Stroke, Fill) -> Unit> = listOf(
+        { _, fill -> rect(rect, fill) },
+        { _, fill -> rect(rect, 10.0, fill) },
+        { stroke, fill -> rect(rect, stroke, fill) },
+        { stroke, fill -> rect(rect, 10.0, stroke, fill) },
+        { _, fill -> circle(circle, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+        { stroke, fill -> circle(circle, stroke, fill) },
+    )
+}
+
 fun check(block: Canvas.(Stroke, Fill) -> Unit) {}
 
 fun box(): String {
-    test()
+    test1()
+    test2()
+
     return "OK"
 }