[NI] Preserve constraint position for OnlyInputType during incorporation

Variable with complex dependency on other variable loses its
original constraint position when variable it depends on is fixed.
To prevent that, original input position is saved in incorporation
position.

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

@@ -5,9 +5,7 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
-import org.jetbrains.kotlin.resolve.calls.inference.model.Constraint
-import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind
-import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
+import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.model.*
 import org.jetbrains.kotlin.utils.SmartSet
@@ -174,7 +172,10 @@ class ConstraintIncorporator(
 
         val kind = if (isSubtype) ConstraintKind.LOWER else ConstraintKind.UPPER
 
-        addNewIncorporatedConstraint(targetVariable, newConstraint, ConstraintContext(kind, derivedFrom))
+        val inputTypePosition = if (baseConstraint.position.from is OnlyInputTypeConstraintPosition)
+            baseConstraint.position else null
+
+        addNewIncorporatedConstraint(targetVariable, newConstraint, ConstraintContext(kind, derivedFrom, inputTypePosition))
     }
 
     fun Context.containsConstrainingTypeWithoutProjection(

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

@@ -48,21 +48,19 @@ class ConstraintInjector(
 
     fun addInitialSubtypeConstraint(c: Context, lowerType: KotlinTypeMarker, upperType: KotlinTypeMarker, position: ConstraintPosition) {
         val initialConstraint = InitialConstraint(lowerType, upperType, UPPER, position)
-        val incorporationPosition = IncorporationConstraintPosition(position, initialConstraint)
         c.addInitialConstraint(initialConstraint)
         updateAllowedTypeDepth(c, lowerType)
         updateAllowedTypeDepth(c, upperType)
-        addSubTypeConstraintAndIncorporateIt(c, lowerType, upperType, incorporationPosition)
+        addSubTypeConstraintAndIncorporateIt(c, lowerType, upperType, IncorporationConstraintPosition(position, initialConstraint))
     }
 
     fun addInitialEqualityConstraint(c: Context, a: KotlinTypeMarker, b: KotlinTypeMarker, position: ConstraintPosition) {
         val initialConstraint = InitialConstraint(a, b, ConstraintKind.EQUALITY, position)
-        val incorporationPosition = IncorporationConstraintPosition(position, initialConstraint)
         c.addInitialConstraint(initialConstraint)
         updateAllowedTypeDepth(c, a)
         updateAllowedTypeDepth(c, b)
-        addSubTypeConstraintAndIncorporateIt(c, a, b, incorporationPosition)
-        addSubTypeConstraintAndIncorporateIt(c, b, a, incorporationPosition)
+        addSubTypeConstraintAndIncorporateIt(c, a, b, IncorporationConstraintPosition(position, initialConstraint))
+        addSubTypeConstraintAndIncorporateIt(c, b, a, IncorporationConstraintPosition(position, initialConstraint))
     }
 
     private fun addSubTypeConstraintAndIncorporateIt(
@@ -72,7 +70,7 @@ class ConstraintInjector(
         incorporatePosition: IncorporationConstraintPosition
     ) {
         val possibleNewConstraints = Stack<Pair<TypeVariableMarker, Constraint>>()
-        val typeCheckerContext = TypeCheckerContext(c, incorporatePosition, lowerType, upperType, possibleNewConstraints)
+        val typeCheckerContext = TypeCheckerContext(c, incorporatePosition, lowerType, upperType, possibleNewConstraints, mutableSetOf())
         typeCheckerContext.runIsSubtypeOf(lowerType, upperType)
 
         while (possibleNewConstraints.isNotEmpty()) {
@@ -87,6 +85,8 @@ class ConstraintInjector(
                 constraintIncorporator.incorporate(typeCheckerContext, typeVariable, it)
             }
         }
+
+        incorporatePosition.inputTypePositions = typeCheckerContext.preservedInputTypePositions
     }
 
     private fun updateAllowedTypeDepth(c: Context, initialType: KotlinTypeMarker) = with(c) {
@@ -122,7 +122,8 @@ class ConstraintInjector(
         val position: IncorporationConstraintPosition,
         val baseLowerType: KotlinTypeMarker,
         val baseUpperType: KotlinTypeMarker,
-        val possibleNewConstraints: MutableList<Pair<TypeVariableMarker, Constraint>>
+        val possibleNewConstraints: MutableList<Pair<TypeVariableMarker, Constraint>>,
+        val preservedInputTypePositions: MutableSet<ConstraintPosition>
     ) : AbstractTypeCheckerContextForConstraintSystem(), ConstraintIncorporator.Context, TypeSystemInferenceExtensionContext by c {
 
         val baseContext: AbstractTypeCheckerContext = newBaseTypeCheckerContext(isErrorTypeEqualsToAnything, isStubTypeEqualsToAnything)
@@ -193,7 +194,7 @@ class ConstraintInjector(
             type: KotlinTypeMarker,
             constraintContext: ConstraintContext
         ) {
-            val (kind, derivedFrom) = constraintContext
+            val (kind, derivedFrom, inputTypePosition) = constraintContext
 
             var targetType = type
             if (targetType.isUninferredParameter()) {
@@ -231,6 +232,7 @@ class ConstraintInjector(
             }
 
             possibleNewConstraints.add(typeVariable to Constraint(kind, targetType, position, derivedFrom = derivedFrom))
+            inputTypePosition?.let { preservedInputTypePositions.add(it) }
         }
 
         override val allTypeVariablesWithConstraints: Collection<VariableWithConstraints>
@@ -259,4 +261,8 @@ class ConstraintInjector(
     }
 }
 
-data class ConstraintContext(val kind: ConstraintKind, val derivedFrom: Set<TypeVariableMarker>)
+data class ConstraintContext(
+    val kind: ConstraintKind,
+    val derivedFrom: Set<TypeVariableMarker>,
+    val inputTypePosition: ConstraintPosition? = null
+)

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

@@ -29,11 +29,11 @@ import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 sealed class ConstraintPosition
 
-class ExplicitTypeParameterConstraintPosition(val typeArgument: SimpleTypeArgument) : ConstraintPosition() {
+class ExplicitTypeParameterConstraintPosition(val typeArgument: SimpleTypeArgument) : ConstraintPosition(), OnlyInputTypeConstraintPosition {
     override fun toString() = "TypeParameter $typeArgument"
 }
 
-class ExpectedTypeConstraintPosition(val topLevelCall: KotlinCall) : ConstraintPosition() {
+class ExpectedTypeConstraintPosition(val topLevelCall: KotlinCall) : ConstraintPosition(), OnlyInputTypeConstraintPosition {
     override fun toString() = "ExpectedType for call $topLevelCall"
 }
 
@@ -45,11 +45,13 @@ class DeclaredUpperBoundConstraintPositionImpl(val typeParameterDescriptor: Type
 
 class FirDeclaredUpperBoundConstraintPosition : DeclaredUpperBoundConstraintPosition()
 
-class ArgumentConstraintPosition(val argument: KotlinCallArgument) : ConstraintPosition() {
+interface OnlyInputTypeConstraintPosition
+
+class ArgumentConstraintPosition(val argument: KotlinCallArgument) : ConstraintPosition(), OnlyInputTypeConstraintPosition {
     override fun toString() = "Argument $argument"
 }
 
-class ReceiverConstraintPosition(val argument: KotlinCallArgument) : ConstraintPosition() {
+class ReceiverConstraintPosition(val argument: KotlinCallArgument) : ConstraintPosition(), OnlyInputTypeConstraintPosition {
     override fun toString() = "Receiver $argument"
 }
 
@@ -77,8 +79,14 @@ class DelegatedPropertyConstraintPosition(val topLevelCall: KotlinCall) : Constr
     override fun toString() = "Constraint from call $topLevelCall for delegated property"
 }
 
-class IncorporationConstraintPosition(val from: ConstraintPosition, val initialConstraint: InitialConstraint) : ConstraintPosition() {
-    override fun toString() = "Incorporate $initialConstraint from position $from"
+class IncorporationConstraintPosition(
+    val from: ConstraintPosition,
+    val initialConstraint: InitialConstraint
+) : ConstraintPosition() {
+    lateinit var inputTypePositions: Set<ConstraintPosition>
+
+    override fun toString() =
+        "Incorporate $initialConstraint from position $from"
 }
 
 class CoroutinePosition() : ConstraintPosition() {

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

@@ -12,9 +12,7 @@ import org.jetbrains.kotlin.types.KotlinType
 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.TypeConstructor
 import org.jetbrains.kotlin.types.UnwrappedType
-import org.jetbrains.kotlin.types.checker.NewCapturedType
 import org.jetbrains.kotlin.types.typeUtil.unCapture
 import kotlin.collections.ArrayList
 import kotlin.collections.LinkedHashMap
@@ -34,14 +32,16 @@ class MutableVariableWithConstraints(
 
     // see @OnlyInputTypes annotation
     val projectedInputCallTypes: Collection<UnwrappedType>
-        get() =
-            mutableConstraints.filter {
-                val position = it.position.from
-                position is ArgumentConstraintPosition || position is ReceiverConstraintPosition ||
-                        position is ExpectedTypeConstraintPosition || position is ExplicitTypeParameterConstraintPosition
-            }.map {
-                (it.type as KotlinType).unCapture().unwrap()
-            }
+        get() = mutableConstraints
+            .filter { it.position.isInputTypePosition }
+            .map { (it.type as KotlinType).unCapture().unwrap() }
+
+    private val ConstraintPosition?.isInputTypePosition: Boolean
+        get() = when (this) {
+            is OnlyInputTypeConstraintPosition -> true
+            !is IncorporationConstraintPosition -> false
+            else -> from.isInputTypePosition || inputTypePositions.any { it.isInputTypePosition }
+        }
 
     private val mutableConstraints = ArrayList(constraints)
 

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

@@ -298,10 +298,10 @@ class NewConstraintSystemImpl(
         variableWithConstraints: MutableVariableWithConstraints?,
         resultType: KotlinTypeMarker
     ) {
-        if (resultType !is KotlinType) return
-        if (variableWithConstraints == null || variableWithConstraints.typeVariable.safeAs<NewTypeVariable>()?.hasOnlyInputTypesAnnotation() != true ) return
-        val projectedInputCallTypes = variableWithConstraints.projectedInputCallTypes
-        val resultTypeIsInputType = projectedInputCallTypes.any { inputType ->
+        if (resultType !is KotlinType || variableWithConstraints == null) return
+        if (variableWithConstraints.typeVariable.safeAs<NewTypeVariable>()?.hasOnlyInputTypesAnnotation() != true) return
+
+        val resultTypeIsInputType = variableWithConstraints.projectedInputCallTypes.any { inputType ->
             NewKotlinTypeChecker.Default.equalTypes(resultType, inputType) ||
                     inputType.constructor is IntersectionTypeConstructor
                     && inputType.constructor.supertypes.any { NewKotlinTypeChecker.Default.equalTypes(resultType, it) }