[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.*
-import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind
-import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
 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, a, b, IncorporationConstraintPosition(position, initialConstraint))
-        addSubTypeConstraintAndIncorporateIt(c, b, a, incorporationPosition)
+        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() {
+class IncorporationConstraintPosition(
-    override fun toString() = "Incorporate $initialConstraint from position $from"
+    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() =
+        get() = mutableConstraints
-            mutableConstraints.filter {
+            .filter { it.position.isInputTypePosition }
-                val position = it.position.from
+            .map { (it.type as KotlinType).unCapture().unwrap() }
-                position is ArgumentConstraintPosition || position is ReceiverConstraintPosition ||
+
-                        position is ExpectedTypeConstraintPosition || position is ExplicitTypeParameterConstraintPosition
+    private val ConstraintPosition?.isInputTypePosition: Boolean
-            }.map {
+        get() = when (this) {
-                (it.type as KotlinType).unCapture().unwrap()
+            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 (resultType !is KotlinType || variableWithConstraints == null) return
-        if (variableWithConstraints == null || variableWithConstraints.typeVariable.safeAs<NewTypeVariable>()?.hasOnlyInputTypesAnnotation() != true ) return
+        if (variableWithConstraints.typeVariable.safeAs<NewTypeVariable>()?.hasOnlyInputTypesAnnotation() != true) return
-        val projectedInputCallTypes = variableWithConstraints.projectedInputCallTypes
+
-        val resultTypeIsInputType = projectedInputCallTypes.any { inputType ->
+        val resultTypeIsInputType = variableWithConstraints.projectedInputCallTypes.any { inputType ->
             NewKotlinTypeChecker.Default.equalTypes(resultType, inputType) ||
                     inputType.constructor is IntersectionTypeConstructor
                     && inputType.constructor.supertypes.any { NewKotlinTypeChecker.Default.equalTypes(resultType, it) }

compiler/testData/diagnostics/testsWithStdLib/inference/annotationsForResolve/onlyInputTypesCommonConstraintSystem.kt

@@ -0,0 +1,54 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE
+// !WITH_NEW_INFERENCE
+
+@file:Suppress("INVISIBLE_MEMBER", "INVISIBLE_REFERENCE")
+
+import kotlin.internal.OnlyInputTypes
+
+interface Bound
+class First : Bound
+class Second : Bound
+class Inv<I >(val v: I)
+class InvB<I : Bound>(val v: I)
+class In<in C>(v: C)
+class InB<in C : Bound>(v: C)
+class Out<out O>(val v: O)
+class OutB<out O : Bound>(val v: O)
+
+fun <@OnlyInputTypes M> strictId(arg: M): M = arg
+fun <@OnlyInputTypes S> strictSelect(arg1: S, arg2: S): S = arg1
+
+fun testOK(first: First, bound: Bound, second: Second) {
+    strictId(Inv(15))
+    strictId(Inv("foo"))
+    strictId(Inv(first))
+    strictId(InvB(first))
+    strictId(In(first))
+    strictId(InB(first))
+    strictId(Out(first))
+    strictId(OutB(first))
+    strictId(Inv(Inv(Inv(first))))
+
+    strictSelect(Inv(first), Inv(first))
+    strictSelect(InvB(first), InvB(first))
+
+    strictSelect(Out(first), Out(bound))
+    strictSelect(OutB(first), OutB(bound))
+    strictSelect(In(first), In(bound))
+    strictSelect(InB(first), InB(bound))
+    <!OI;TYPE_INFERENCE_ONLY_INPUT_TYPES!>strictSelect<!>(InB(first), InB(second)) // different behaviour in contravariant position
+
+    val out: Out<Bound> = strictSelect(Out(first), Out(second))
+    val outb: OutB<Bound> = strictSelect(OutB(first), OutB(second))
+    strictSelect<Out<Bound>>(Out(first), Out(second))
+    strictSelect<OutB<Bound>>(OutB(first), OutB(second))
+}
+
+fun testFail(first: First, bound: Bound, second: Second) {
+    <!TYPE_INFERENCE_ONLY_INPUT_TYPES!>strictSelect<!>(InvB(first), InvB(bound))
+    <!TYPE_INFERENCE_ONLY_INPUT_TYPES!>strictSelect<!>(Inv(first), Inv(bound))
+    <!TYPE_INFERENCE_ONLY_INPUT_TYPES!>strictSelect<!>(Out(first), Out(second))
+    <!TYPE_INFERENCE_ONLY_INPUT_TYPES!>strictSelect<!>(In(first), In(second))
+    <!TYPE_INFERENCE_ONLY_INPUT_TYPES!>strictSelect<!>(Out(Inv(first)), Out(Inv(second)))
+    <!TYPE_INFERENCE_ONLY_INPUT_TYPES!>strictSelect<!>(In(Inv(first)), In(Inv(second)))
+}

compiler/testData/diagnostics/testsWithStdLib/inference/annotationsForResolve/onlyInputTypesCommonConstraintSystem.txt

@@ -0,0 +1,72 @@
+package
+
+public fun </*0*/ @kotlin.internal.OnlyInputTypes M> strictId(/*0*/ arg: M): M
+public fun </*0*/ @kotlin.internal.OnlyInputTypes S> strictSelect(/*0*/ arg1: S, /*1*/ arg2: S): S
+public fun testFail(/*0*/ first: First, /*1*/ bound: Bound, /*2*/ second: Second): kotlin.Unit
+public fun testOK(/*0*/ first: First, /*1*/ bound: Bound, /*2*/ second: Second): kotlin.Unit
+
+public interface Bound {
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class First : Bound {
+    public constructor First()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class In</*0*/ in C> {
+    public constructor In</*0*/ in C>(/*0*/ v: C)
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class InB</*0*/ in C : Bound> {
+    public constructor InB</*0*/ in C : Bound>(/*0*/ v: C)
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Inv</*0*/ I> {
+    public constructor Inv</*0*/ I>(/*0*/ v: I)
+    public final val v: I
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class InvB</*0*/ I : Bound> {
+    public constructor InvB</*0*/ I : Bound>(/*0*/ v: I)
+    public final val v: I
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Out</*0*/ out O> {
+    public constructor Out</*0*/ out O>(/*0*/ v: O)
+    public final val v: O
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class OutB</*0*/ out O : Bound> {
+    public constructor OutB</*0*/ out O : Bound>(/*0*/ v: O)
+    public final val v: O
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Second : Bound {
+    public constructor Second()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestWithStdLibGenerated.java

@@ -2843,6 +2843,11 @@ public class DiagnosticsTestWithStdLibGenerated extends AbstractDiagnosticsTestW
                 runTest("compiler/testData/diagnostics/testsWithStdLib/inference/annotationsForResolve/onlyInputTypesCaptured.kt");
             }
 
+            @TestMetadata("onlyInputTypesCommonConstraintSystem.kt")
+            public void testOnlyInputTypesCommonConstraintSystem() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/inference/annotationsForResolve/onlyInputTypesCommonConstraintSystem.kt");
+            }
+
             @TestMetadata("onlyInputTypesWithVarargs.kt")
             public void testOnlyInputTypesWithVarargs() throws Exception {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/inference/annotationsForResolve/onlyInputTypesWithVarargs.kt");

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsTestWithStdLibUsingJavacGenerated.java

@@ -2843,6 +2843,11 @@ public class DiagnosticsTestWithStdLibUsingJavacGenerated extends AbstractDiagno
                 runTest("compiler/testData/diagnostics/testsWithStdLib/inference/annotationsForResolve/onlyInputTypesCaptured.kt");
             }
 
+            @TestMetadata("onlyInputTypesCommonConstraintSystem.kt")
+            public void testOnlyInputTypesCommonConstraintSystem() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/inference/annotationsForResolve/onlyInputTypesCommonConstraintSystem.kt");
+            }
+
             @TestMetadata("onlyInputTypesWithVarargs.kt")
             public void testOnlyInputTypesWithVarargs() throws Exception {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/inference/annotationsForResolve/onlyInputTypesWithVarargs.kt");