[NI] Commonize detection of @OnlyInputTypes in NewConstraintSystemImpl

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/ConeConstraintSystemUtilContext.kt

@@ -5,7 +5,9 @@
 
 package org.jetbrains.kotlin.fir.resolve.inference
 
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemUtilContext
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
@@ -13,4 +15,15 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
         // TODO
         return false
     }
+
+    override fun TypeVariableMarker.hasOnlyInputTypesAttribute(): Boolean {
+        // TODO
+        return false
+    }
+
+    override fun KotlinTypeMarker.unCapture(): KotlinTypeMarker {
+        require(this is ConeKotlinType)
+        // TODO, see TypeUtils.kt
+        return this
+    }
 }

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

@@ -6,11 +6,25 @@
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
 import org.jetbrains.kotlin.resolve.calls.components.CreateFreshVariablesSubstitutor.shouldBeFlexible
+import org.jetbrains.kotlin.resolve.calls.inference.model.NewTypeVariable
 import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableFromCallableDescriptor
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
+import org.jetbrains.kotlin.types.typeUtil.unCapture as unCaptureKotlinType
 
 object ClassicConstraintSystemUtilContext : ConstraintSystemUtilContext {
     override fun TypeVariableMarker.shouldBeFlexible(): Boolean {
         return this is TypeVariableFromCallableDescriptor && this.originalTypeParameter.shouldBeFlexible()
     }
+
+    override fun TypeVariableMarker.hasOnlyInputTypesAttribute(): Boolean {
+        require(this is NewTypeVariable)
+        return hasOnlyInputTypesAnnotation()
+    }
+
+    override fun KotlinTypeMarker.unCapture(): KotlinTypeMarker {
+        require(this is KotlinType)
+        return unCaptureKotlinType().unwrap()
+    }
 }

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 /*
@@ -14,4 +15,6 @@ import org.jetbrains.kotlin.types.model.TypeVariableMarker
  */
 interface ConstraintSystemUtilContext {
     fun TypeVariableMarker.shouldBeFlexible(): Boolean
+    fun TypeVariableMarker.hasOnlyInputTypesAttribute(): Boolean
+    fun KotlinTypeMarker.unCapture(): KotlinTypeMarker
 }

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

@@ -5,16 +5,15 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.model
 
+import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemUtilContext
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 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.unCapture
 import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.addToStdlib.trimToSize
 
-
 class MutableVariableWithConstraints private constructor(
     override val typeVariable: TypeVariableMarker,
     constraints: List<Constraint>? // assume simplified and deduplicated
@@ -33,13 +32,16 @@ class MutableVariableWithConstraints private constructor(
         }
 
     // see @OnlyInputTypes annotation
-    val projectedInputCallTypes: Collection<UnwrappedType>
-        get() = mutableConstraints
-            .mapNotNullTo(SmartList()) {
-                if (it.position.from is OnlyInputTypeConstraintPosition || it.inputTypePositionBeforeIncorporation != null)
-                    (it.type as KotlinType).unCapture().unwrap()
-                else null
-            }
+    fun getProjectedInputCallTypes(utilContext: ConstraintSystemUtilContext): Collection<KotlinTypeMarker> {
+        return with(utilContext) {
+            mutableConstraints
+                .mapNotNullTo(SmartList()) {
+                    if (it.position.from is OnlyInputTypeConstraintPosition || it.inputTypePositionBeforeIncorporation != null)
+                        it.type.unCapture()
+                    else null
+                }
+        }
+    }
 
     private val mutableConstraints = if (constraints == null) SmartList() else SmartList(constraints)
 

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

@@ -8,26 +8,26 @@ package org.jetbrains.kotlin.resolve.calls.inference.model
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzerContext
 import org.jetbrains.kotlin.resolve.calls.inference.*
 import org.jetbrains.kotlin.resolve.calls.inference.components.*
-import org.jetbrains.kotlin.types.*
-import org.jetbrains.kotlin.types.checker.NewKotlinTypeChecker
+import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.model.*
 import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.SmartSet
-import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.utils.addToStdlib.trimToSize
 import kotlin.math.max
 
 class NewConstraintSystemImpl(
     private val constraintInjector: ConstraintInjector,
     val typeSystemContext: TypeSystemInferenceExtensionContext
-) :
-    TypeSystemInferenceExtensionContext by typeSystemContext,
+) : TypeSystemInferenceExtensionContext by typeSystemContext,
     NewConstraintSystem,
     ConstraintSystemBuilder,
     ConstraintInjector.Context,
     ResultTypeResolver.Context,
     ConstraintSystemCompletionContext,
-    PostponedArgumentsAnalyzerContext {
+    PostponedArgumentsAnalyzerContext
+{
+    private val utilContext = constraintInjector.constraintIncorporator.utilContext
+
     private val storage = MutableConstraintStorage()
     private var state = State.BUILDING
     private val typeVariablesTransaction: MutableList<TypeVariableMarker> = SmartList()
@@ -316,16 +316,17 @@ class NewConstraintSystemImpl(
         variableWithConstraints: MutableVariableWithConstraints?,
         resultType: KotlinTypeMarker
     ) {
-        if (resultType !is KotlinType || variableWithConstraints == null) return
-        if (variableWithConstraints.typeVariable.safeAs<NewTypeVariable>()?.hasOnlyInputTypesAnnotation() != true) return
+        if (variableWithConstraints == null) return
+        val variableHasOnlyInputTypes = with(utilContext) { variableWithConstraints.typeVariable.hasOnlyInputTypesAttribute() }
+        if (!variableHasOnlyInputTypes) 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) }
+        val resultTypeIsInputType = variableWithConstraints.getProjectedInputCallTypes(utilContext).any { inputType ->
+            if (AbstractTypeChecker.equalTypes(this, resultType, inputType)) return@any true
+            val constructor = inputType.typeConstructor()
+            constructor.isIntersection() && constructor.supertypes().any { AbstractTypeChecker.equalTypes(this, resultType, it) }
         }
         if (!resultTypeIsInputType) {
-            addError(OnlyInputTypesDiagnostic(variableWithConstraints.typeVariable as NewTypeVariable))
+            addError(OnlyInputTypesDiagnostic(variableWithConstraints.typeVariable))
         }
     }