NI: Support forking inference with heuristics

Mostly, it only affects FIR

It partially allows to consider several variance of constraints like
A<Int> & A<T> <: A<X_var> that are mostly brought by smart casts

^KT-49542 Fixed
^KT-50489 Relates

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

@@ -7,8 +7,10 @@ package org.jetbrains.kotlin.resolve.calls.inference
 
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzerContext
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionContext
+import org.jetbrains.kotlin.resolve.calls.inference.model.Constraint
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintSystemError
+import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 interface NewConstraintSystem {
     val hasContradiction: Boolean
@@ -21,4 +23,8 @@ interface NewConstraintSystem {
 
     fun asConstraintSystemCompleterContext(): ConstraintSystemCompletionContext
     fun asPostponedArgumentsAnalyzerContext(): PostponedArgumentsAnalyzerContext
+    fun processForkConstraints()
 }
+
+typealias ForkPointData = List<ConstraintsFromSingleFork>
+typealias ConstraintsFromSingleFork = Set<Pair<TypeVariableMarker, Constraint>>

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

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.resolve.calls.inference.components
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemOperation
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintsFromSingleFork
+import org.jetbrains.kotlin.resolve.calls.inference.ForkPointData
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind.*
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
@@ -16,6 +18,8 @@ import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.TypeCheckerState
 import org.jetbrains.kotlin.types.model.*
 import org.jetbrains.kotlin.utils.SmartList
+import org.jetbrains.kotlin.utils.addIfNotNull
+import org.jetbrains.kotlin.utils.addToStdlib.popLast
 import kotlin.math.max
 
 class ConstraintInjector(
@@ -31,6 +35,8 @@ class ConstraintInjector(
         var maxTypeDepthFromInitialConstraints: Int
         val notFixedTypeVariables: MutableMap<TypeConstructorMarker, MutableVariableWithConstraints>
         val fixedTypeVariables: MutableMap<TypeConstructorMarker, KotlinTypeMarker>
+        val constraintsFromAllForkPoints: MutableList<Pair<IncorporationConstraintPosition, ForkPointData>>
+        val atCompletionState: Boolean
 
         fun addInitialConstraint(initialConstraint: InitialConstraint)
         fun addError(error: ConstraintSystemError)
@@ -39,6 +45,8 @@ class ConstraintInjector(
             position: IncorporationConstraintPosition,
             constraints: MutableList<Pair<TypeVariableMarker, Constraint>>
         )
+
+        fun processForkConstraints()
     }
 
     fun addInitialSubtypeConstraint(c: Context, lowerType: KotlinTypeMarker, upperType: KotlinTypeMarker, position: ConstraintPosition) {
@@ -133,10 +141,41 @@ class ConstraintInjector(
         processConstraints(c, typeCheckerState, skipProperEqualityConstraints = false)
     }
 
+    fun processForkConstraints(
+        c: Context,
+        constraintSet: Collection<Pair<TypeVariableMarker, Constraint>>,
+        position: IncorporationConstraintPosition
+    ) {
+        processGivenConstraints(
+            c,
+            TypeCheckerStateForConstraintInjector(c, position),
+            constraintSet,
+        )
+    }
+
     private fun processConstraints(
         c: Context,
         typeCheckerState: TypeCheckerStateForConstraintInjector,
         skipProperEqualityConstraints: Boolean = true
+    ): MutableList<Pair<TypeVariableMarker, Constraint>>? {
+        return processConstraintsIgnoringForksData(typeCheckerState, c, skipProperEqualityConstraints).also {
+            typeCheckerState.extractForkPointsData()?.let { allForkPointsData ->
+                allForkPointsData.mapTo(c.constraintsFromAllForkPoints) { forkPointData ->
+                    typeCheckerState.position to forkPointData
+                }
+
+                // During completion, we start processing fork constrains immediately
+                if (c.atCompletionState) {
+                    c.processForkConstraints()
+                }
+            }
+        }
+    }
+
+    private fun processConstraintsIgnoringForksData(
+        typeCheckerState: TypeCheckerStateForConstraintInjector,
+        c: Context,
+        skipProperEqualityConstraints: Boolean
     ): MutableList<Pair<TypeVariableMarker, Constraint>>? {
         val properConstraintsProcessingEnabled =
             languageVersionSettings.supportsFeature(LanguageFeature.ProperTypeInferenceConstraintsProcessing)
@@ -165,7 +204,7 @@ class ConstraintInjector(
     private fun processGivenConstraints(
         c: Context,
         typeCheckerState: TypeCheckerStateForConstraintInjector,
-        constraintsToProcess: MutableList<Pair<TypeVariableMarker, Constraint>>
+        constraintsToProcess: Collection<Pair<TypeVariableMarker, Constraint>>
     ) {
         for ((typeVariable, constraint) in constraintsToProcess) {
             if (c.shouldWeSkipConstraint(typeVariable, constraint)) continue
@@ -235,22 +274,94 @@ class ConstraintInjector(
         // We use `var` intentionally to avoid extra allocations as this property is quite "hot"
         private var possibleNewConstraints: MutableList<Pair<TypeVariableMarker, Constraint>>? = null
 
+        private var forkPointsData: MutableList<ForkPointData>? = null
+        private var stackForConstraintsSetsFromCurrentForkPoint: Stack<MutableList<ConstraintsFromSingleFork>>? = null
+        private var stackForCurrentConstraintSetFromSingleFork: Stack<MutableList<Pair<TypeVariableMarker, Constraint>>>? = null
+
         override val isInferenceCompatibilityEnabled = languageVersionSettings.supportsFeature(LanguageFeature.InferenceCompatibility)
 
+        private val allowForking: Boolean
+            get() = constraintIncorporator.utilContext.isForcedAllowForkingInferenceSystem
+
         private var baseLowerType = position.initialConstraint.a
         private var baseUpperType = position.initialConstraint.b
 
         private var isIncorporatingConstraintFromDeclaredUpperBound = false
 
         fun extractAllConstraints() = possibleNewConstraints.also { possibleNewConstraints = null }
+        fun extractForkPointsData() = forkPointsData.also { forkPointsData = null }
 
         fun addPossibleNewConstraint(variable: TypeVariableMarker, constraint: Constraint) {
+            val constraintsSetsFromCurrentFork = stackForConstraintsSetsFromCurrentForkPoint?.lastOrNull()
+            if (constraintsSetsFromCurrentFork != null) {
+                val currentConstraintSetForFork = stackForCurrentConstraintSetFromSingleFork?.lastOrNull()
+                require(currentConstraintSetForFork != null) { "Constraint has been added not under fork {...} call " }
+                currentConstraintSetForFork.add(variable to constraint)
+                return
+            }
+
             if (possibleNewConstraints == null) {
                 possibleNewConstraints = SmartList()
             }
             possibleNewConstraints!!.add(variable to constraint)
         }
 
+        override fun runForkingPoint(block: ForkPointContext.() -> Unit): Boolean {
+            if (!allowForking) {
+                return super.runForkingPoint(block)
+            }
+
+            if (stackForConstraintsSetsFromCurrentForkPoint == null) {
+                stackForConstraintsSetsFromCurrentForkPoint = SmartList()
+            }
+
+            stackForConstraintsSetsFromCurrentForkPoint!!.add(SmartList())
+            val isThereSuccessfulFork = with(MyForkCreationContext()) {
+                block()
+                anyForkSuccessful
+            }
+
+            val constraintSets = stackForConstraintsSetsFromCurrentForkPoint?.popLast()
+            if (constraintSets == null || constraintSets.isEmpty()) return isThereSuccessfulFork
+
+            if (constraintSets.size > 1) {
+                if (forkPointsData == null) {
+                    forkPointsData = SmartList()
+                }
+                forkPointsData!!.addIfNotNull(
+                    constraintSets
+                )
+                return true
+            } else if (constraintSets.size == 1) {
+                processForkConstraints(
+                    c,
+                    constraintSets.single(),
+                    position,
+                )
+            }
+
+            return isThereSuccessfulFork
+        }
+
+        private inner class MyForkCreationContext : ForkPointContext {
+            var anyForkSuccessful = false
+
+            override fun fork(block: () -> Boolean) {
+                if (stackForCurrentConstraintSetFromSingleFork == null) {
+                    stackForCurrentConstraintSetFromSingleFork = SmartList()
+                }
+
+                stackForCurrentConstraintSetFromSingleFork!!.add(SmartList())
+
+                block().also { anyForkSuccessful = anyForkSuccessful || it }
+
+                stackForConstraintsSetsFromCurrentForkPoint!!.last()
+                    .addIfNotNull(
+                        stackForCurrentConstraintSetFromSingleFork?.popLast()?.takeIf { it.isNotEmpty() }?.toSet()
+                    )
+            }
+        }
+
         fun hasConstraintsToProcess() = possibleNewConstraints != null
 
         fun setConstrainingTypesToPrintDebugInfo(lowerType: KotlinTypeMarker, upperType: KotlinTypeMarker) {
@@ -445,3 +556,5 @@ data class ConstraintContext(
     val inputTypePositionBeforeIncorporation: OnlyInputTypeConstraintPosition? = null,
     val isNullabilityConstraint: Boolean
 )
+
+private typealias Stack<E> = MutableList<E>

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

@@ -33,4 +33,5 @@ interface ConstraintSystemCompletionContext : VariableFixationFinder.Context, Re
     fun fixVariable(variable: TypeVariableMarker, resultType: KotlinTypeMarker, position: FixVariableConstraintPosition<*>)
 
     fun couldBeResolvedWithUnrestrictedBuilderInference(): Boolean
+    fun processForkConstraints()
 }

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

@@ -39,4 +39,6 @@ interface ConstraintSystemUtilContext {
     ): TypeVariableMarker
 
     val isForcedConsiderExtensionReceiverFromConstrainsInLambda get() = false
+
+    val isForcedAllowForkingInferenceSystem get() = false
 }

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

@@ -133,6 +133,8 @@ class ConstrainingTypeIsError(
     val position: IncorporationConstraintPosition
 ) : ConstraintSystemError(INAPPLICABLE)
 
+class NoSuccessfulFork(val position: IncorporationConstraintPosition) : ConstraintSystemError(INAPPLICABLE)
+
 class OnlyInputTypesDiagnostic(val typeVariable: TypeVariableMarker) : ConstraintSystemError(INAPPLICABLE)
 
 object LowerPriorityToPreserveCompatibility : ConstraintSystemError(RESOLVED_NEED_PRESERVE_COMPATIBILITY)

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

@@ -5,8 +5,12 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.model
 
+import org.jetbrains.kotlin.resolve.calls.inference.ForkPointData
 import org.jetbrains.kotlin.types.AbstractTypeChecker
-import org.jetbrains.kotlin.types.model.*
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+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 +46,7 @@ interface ConstraintStorage {
     val postponedTypeVariables: List<TypeVariableMarker>
     val builtFunctionalTypesForPostponedArgumentsByTopLevelTypeVariables: Map<Pair<TypeConstructorMarker, List<Pair<TypeConstructorMarker, Int>>>, KotlinTypeMarker>
     val builtFunctionalTypesForPostponedArgumentsByExpectedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker>
+    val constraintsFromAllForkPoints: List<Pair<IncorporationConstraintPosition, ForkPointData>>
 
     object Empty : ConstraintStorage {
         override val allTypeVariables: Map<TypeConstructorMarker, TypeVariableMarker> get() = emptyMap()
@@ -55,6 +60,7 @@ interface ConstraintStorage {
         override val postponedTypeVariables: List<TypeVariableMarker> get() = emptyList()
         override val builtFunctionalTypesForPostponedArgumentsByTopLevelTypeVariables: Map<Pair<TypeConstructorMarker, List<Pair<TypeConstructorMarker, Int>>>, KotlinTypeMarker> = emptyMap()
         override val builtFunctionalTypesForPostponedArgumentsByExpectedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker> = emptyMap()
+        override val constraintsFromAllForkPoints: List<Pair<IncorporationConstraintPosition, ForkPointData>> = emptyList()
     }
 }
 
@@ -113,14 +119,16 @@ class InitialConstraint(
     val constraintKind: ConstraintKind, // see [checkConstraint]
     val position: ConstraintPosition
 ) {
-    override fun toString(): String {
+    override fun toString(): String = "${asStringWithoutPosition()} from $position"
+
+    fun asStringWithoutPosition(): String {
         val sign =
             when (constraintKind) {
                 ConstraintKind.EQUALITY -> "=="
                 ConstraintKind.LOWER -> ":>"
                 ConstraintKind.UPPER -> "<:"
             }
-        return "$a $sign $b from $position"
+        return "$a $sign $b"
     }
 }
 

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.model
 
+import org.jetbrains.kotlin.resolve.calls.inference.ForkPointData
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemUtilContext
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 import org.jetbrains.kotlin.types.model.*
@@ -238,4 +239,6 @@ internal class MutableConstraintStorage : ConstraintStorage {
         LinkedHashMap()
     override val builtFunctionalTypesForPostponedArgumentsByExpectedTypeVariables: MutableMap<TypeConstructorMarker, KotlinTypeMarker> =
         LinkedHashMap()
+
+    override val constraintsFromAllForkPoints: MutableList<Pair<IncorporationConstraintPosition, ForkPointData>> = SmartList()
 }

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

@@ -40,6 +40,8 @@ class NewConstraintSystemImpl(
 
     private var couldBeResolvedWithUnrestrictedBuilderInference: Boolean = false
 
+    override var atCompletionState: Boolean = false
+
     private enum class State {
         BUILDING,
         TRANSACTION,
@@ -100,7 +102,11 @@ class NewConstraintSystemImpl(
         storage.missedConstraints.add(position to constraints)
     }
 
-    override fun asConstraintSystemCompleterContext() = apply { checkState(State.BUILDING) }
+    override fun asConstraintSystemCompleterContext() = apply {
+        checkState(State.BUILDING)
+
+        this.atCompletionState = true
+    }
 
     override fun asPostponedArgumentsAnalyzerContext() = apply { checkState(State.BUILDING) }
 
@@ -203,6 +209,7 @@ class NewConstraintSystemImpl(
         val beforeTypeVariablesTransactionSize = typeVariablesTransaction.size
         val beforeMissedConstraintsCount = storage.missedConstraints.size
         val beforeConstraintCountByVariables = storage.notFixedTypeVariables.mapValues { it.value.rawConstraintsCount }
+        val beforeConstraintsFromAllForks = storage.constraintsFromAllForkPoints.size
 
         state = State.TRANSACTION
         // typeVariablesTransaction is clear
@@ -218,6 +225,7 @@ class NewConstraintSystemImpl(
         storage.maxTypeDepthFromInitialConstraints = beforeMaxTypeDepthFromInitialConstraints
         storage.errors.trimToSize(beforeErrorsCount)
         storage.missedConstraints.trimToSize(beforeMissedConstraintsCount)
+        storage.constraintsFromAllForkPoints.trimToSize(beforeConstraintsFromAllForks)
 
         val addedInitialConstraints = storage.initialConstraints.subList(beforeInitialConstraintCount, storage.initialConstraints.size)
 
@@ -337,6 +345,54 @@ class NewConstraintSystemImpl(
             return storage.postponedTypeVariables
         }
 
+    override val constraintsFromAllForkPoints: MutableList<Pair<IncorporationConstraintPosition, ForkPointData>>
+        get() {
+            checkState(State.BUILDING, State.COMPLETION, State.TRANSACTION)
+            return storage.constraintsFromAllForkPoints
+        }
+
+    override fun processForkConstraints() {
+        if (constraintsFromAllForkPoints.isEmpty()) return
+        val allForkPointsData = constraintsFromAllForkPoints.toList()
+        constraintsFromAllForkPoints.clear()
+
+        // There may be multiple fork points:
+        // - One from subtyping A<Int> & A<T> <: A<Xv>
+        // - Another one from B<String> & B<F> <: B<Yv>
+        // Each of them defines two sets of constraints, e.g. for the first for point:
+        // 1. {Xv=Int} – is a one-element set (but potentially there might be more constraints in the set)
+        // 2. {Xv=T} – second constraints set
+        for ((position, forkPointData) in allForkPointsData) {
+            if (!processForkPointData(forkPointData, position)) {
+                addError(NoSuccessfulFork(position))
+            }
+        }
+    }
+
+    /**
+     * @return true if there is a successful constraints set for the fork
+     */
+    private fun processForkPointData(
+        forkPointData: ForkPointData,
+        position: IncorporationConstraintPosition
+    ): Boolean {
+        return forkPointData.any { constraintSetForSingleFork ->
+            runTransaction {
+                constraintInjector.processForkConstraints(
+                    this@NewConstraintSystemImpl.apply { checkState(State.BUILDING, State.COMPLETION, State.TRANSACTION) },
+                    constraintSetForSingleFork,
+                    position,
+                )
+
+                if (constraintsFromAllForkPoints.isNotEmpty()) {
+                    processForkConstraints()
+                }
+
+                !hasContradiction
+            }
+        }
+    }
+
     // ConstraintInjector.Context, KotlinConstraintSystemCompleter.Context
     override fun addError(error: ConstraintSystemError) {
         checkState(State.BUILDING, State.COMPLETION, State.TRANSACTION)