K2: Implement partially constrained lambda analysis (PCLA)

It's expected to partially mimic the behavior of what previously was called builder inference, but with more clear contracts (documentation is in progress, though) See a lot of fixed issues in the later commits with test data, especially [red-to-green] ^KT-59791 In Progress
2023-08-21 17:51:54 +02:00
parent 9ae5287df8
commit 276f5b26d8
44 changed files with 1356 additions and 1086 deletions
@@ -7,10 +7,13 @@ package org.jetbrains.kotlin.resolve.calls.components

 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
+import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
 import org.jetbrains.kotlin.types.model.*

 interface PostponedArgumentsAnalyzerContext : TypeSystemInferenceExtensionContext {
-    fun buildCurrentSubstitutor(additionalBindings: Map<TypeConstructorMarker, StubTypeMarker>): TypeSubstitutorMarker
+    val notFixedTypeVariables: Map<TypeConstructorMarker, VariableWithConstraints>
+
+    fun buildCurrentSubstitutor(additionalBindings: Map<TypeConstructorMarker, KotlinTypeMarker>): TypeSubstitutorMarker
    fun buildNotFixedVariablesToStubTypesSubstitutor(): TypeSubstitutorMarker
    fun bindingStubsForPostponedVariables(): Map<TypeVariableMarker, StubTypeMarker>

@@ -11,21 +11,16 @@ import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.PostponedAtomWithRevisableExpectedType
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
+import org.jetbrains.kotlin.types.model.K2Only
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker

 abstract class ConstraintSystemCompletionContext : VariableFixationFinder.Context, ResultTypeResolver.Context {
-    abstract val allTypeVariables: Map<TypeConstructorMarker, TypeVariableMarker>
    abstract override val notFixedTypeVariables: Map<TypeConstructorMarker, VariableWithConstraints>
    abstract override val fixedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker>
    abstract override val postponedTypeVariables: List<TypeVariableMarker>

-    /**
-     * See [org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage.outerSystemVariablesPrefixSize]
-     */
-    abstract val outerSystemVariablesPrefixSize: Int
-
    abstract fun getBuilder(): ConstraintSystemBuilder

    // type can be proper if it not contains not fixed type variables
@@ -69,7 +64,7 @@ abstract class ConstraintSystemCompletionContext : VariableFixationFinder.Contex
        completionMode: ConstraintSystemCompletionMode,
        analyze: (A) -> Unit
    ): Boolean {
-        if (completionMode == ConstraintSystemCompletionMode.FULL) {
+        if (completionMode.allLambdasShouldBeAnalyzed) {
            val argumentWithTypeVariableAsExpectedType = findPostponedArgumentWithRevisableExpectedType(postponedArguments)

            if (argumentWithTypeVariableAsExpectedType != null) {
@@ -120,4 +115,11 @@ abstract class ConstraintSystemCompletionContext : VariableFixationFinder.Contex
                !it.typeConstructor().isClassTypeConstructor() && !it.typeConstructor().isTypeParameterTypeConstructor()
            }
        }.map { it.type }
+
+    /**
+     * @see [org.jetbrains.kotlin.resolve.calls.inference.components.VariableFixationFinder.Context.typeVariablesThatAreNotCountedAsProperTypes]
+     * @see [org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer.fixInnerVariablesForProvideDelegateIfNeeded]
+     */
+    @K2Only
+    abstract fun <R> withTypeVariablesThatAreCountedAsProperTypes(typeVariables: Set<TypeConstructorMarker>, block: () -> R): R
 }
@@ -5,8 +5,9 @@

 package org.jetbrains.kotlin.resolve.calls.inference.components

-enum class ConstraintSystemCompletionMode {
-    FULL,
+enum class ConstraintSystemCompletionMode(val allLambdasShouldBeAnalyzed: Boolean) {
+    FULL(true),
+    PCLA_POSTPONED_CALL(true),

    /**
     * This mode allows us to infer variables in calls, which have enough type-info to be completed right-away
@@ -18,6 +19,6 @@ enum class ConstraintSystemCompletionMode {
     * x.plus(run { x }) // Here, to select plus overload we need to analyze lambda
     * ```
     */
-    PARTIAL,
-    UNTIL_FIRST_LAMBDA
+    PARTIAL(false),
+    UNTIL_FIRST_LAMBDA(false),
 }
@@ -9,8 +9,13 @@ import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind
 import org.jetbrains.kotlin.builtins.functions.isBasicFunctionOrKFunction
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
-import org.jetbrains.kotlin.resolve.calls.inference.model.*
-import org.jetbrains.kotlin.resolve.calls.model.*
+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.model.LambdaWithTypeVariableAsExpectedTypeMarker
+import org.jetbrains.kotlin.resolve.calls.model.PostponedAtomWithRevisableExpectedType
+import org.jetbrains.kotlin.resolve.calls.model.PostponedCallableReferenceMarker
+import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.types.model.*
 import org.jetbrains.kotlin.utils.SmartSet
 import java.util.*
@@ -531,24 +536,40 @@ class PostponedArgumentInputTypesResolver(
        postponedArguments: List<PostponedResolvedAtomMarker>,
        topLevelType: KotlinTypeMarker,
        dependencyProvider: TypeVariableDependencyInformationProvider,
-        resolvedAtomProvider: ResolvedAtomProvider
-    ): Boolean = with(c) {
-        val expectedType = argument.run { (this as? PostponedAtomWithRevisableExpectedType)?.revisedExpectedType ?: expectedType }
+        resolvedAtomProvider: ResolvedAtomProvider,
+    ): Boolean {
+        val expectedType = argument.expectedFunctionType(c) ?: return false

-        if (expectedType != null && expectedType.isFunctionOrKFunctionWithAnySuspendability()) {
-            val wasFixedSomeVariable = c.fixNextReadyVariableForParameterType(
-                expectedType,
-                postponedArguments,
-                topLevelType,
-                dependencyProvider,
-                resolvedAtomProvider
-            )
+        return c.fixNextReadyVariableForParameterType(
+            expectedType,
+            postponedArguments,
+            topLevelType,
+            dependencyProvider,
+            resolvedAtomProvider,
+        )
+    }

-            if (wasFixedSomeVariable)
-                return true
-        }
+    @K2Only
+    fun findNextVariableForReportingNotInferredInputType(
+        c: Context,
+        argument: PostponedResolvedAtomMarker,
+        postponedArguments: List<PostponedResolvedAtomMarker>,
+        topLevelType: KotlinTypeMarker,
+        dependencyProvider: TypeVariableDependencyInformationProvider,
+    ): VariableFixationFinder.VariableForFixation? {
+        val expectedType = argument.expectedFunctionType(c) ?: return null

-        return false
+        return c.findNextVariableForParameterType(
+            expectedType,
+            dependencyProvider,
+            postponedArguments,
+            topLevelType,
+        )
+    }
+
+    private fun PostponedResolvedAtomMarker.expectedFunctionType(c: Context): KotlinTypeMarker? = with(c) {
+        val expectedType = (this@expectedFunctionType as? PostponedAtomWithRevisableExpectedType)?.revisedExpectedType ?: expectedType
+        expectedType?.takeIf { it.isFunctionOrKFunctionWithAnySuspendability() }
    }

    private fun Context.fixNextReadyVariableForParameterType(
@@ -558,17 +579,9 @@ class PostponedArgumentInputTypesResolver(
        dependencyProvider: TypeVariableDependencyInformationProvider,
        resolvedAtomByTypeVariableProvider: ResolvedAtomProvider,
    ): Boolean = with(resolutionTypeSystemContext) {
-        val relatedVariables = type.extractArgumentsForFunctionTypeOrSubtype()
-            .flatMap { getAllDeeplyRelatedTypeVariables(it, dependencyProvider) }
-        val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
-            this@fixNextReadyVariableForParameterType,
-            relatedVariables,
-            postponedArguments,
-            ConstraintSystemCompletionMode.FULL,
-            topLevelType
-        )
+        val variableForFixation = findNextVariableForParameterType(type, dependencyProvider, postponedArguments, topLevelType)

-        if (variableForFixation == null || !variableForFixation.hasProperConstraint)
+        if (variableForFixation == null || !variableForFixation.isReady)
            return false

        val variableWithConstraints = notFixedTypeVariables.getValue(variableForFixation.variable)
@@ -589,6 +602,26 @@ class PostponedArgumentInputTypesResolver(
        return true
    }

+    private fun Context.findNextVariableForParameterType(
+        type: KotlinTypeMarker,
+        dependencyProvider: TypeVariableDependencyInformationProvider,
+        postponedArguments: List<PostponedResolvedAtomMarker>,
+        topLevelType: KotlinTypeMarker,
+    ): VariableFixationFinder.VariableForFixation? {
+        val outerTypeVariables = outerTypeVariables.orEmpty()
+        val relatedVariables = type.extractArgumentsForFunctionTypeOrSubtype()
+            .flatMap { getAllDeeplyRelatedTypeVariables(it, dependencyProvider) }
+            .filter { it !in outerTypeVariables }
+
+        return variableFixationFinder.findFirstVariableForFixation(
+            this,
+            relatedVariables,
+            postponedArguments,
+            ConstraintSystemCompletionMode.FULL,
+            topLevelType,
+        )
+    }
+
    private fun KotlinTypeMarker?.wrapToTypeWithKind() = this?.let { TypeWithKind(it) }

    companion object {
@@ -22,6 +22,8 @@ class ResultTypeResolver(
    private val languageVersionSettings: LanguageVersionSettings
 ) {
    interface Context : TypeSystemInferenceExtensionContext {
+        val notFixedTypeVariables: Map<TypeConstructorMarker, VariableWithConstraints>
+        val outerSystemVariablesPrefixSize: Int
        fun isProperType(type: KotlinTypeMarker): Boolean
        fun buildNotFixedVariablesToStubTypesSubstitutor(): TypeSubstitutorMarker
        fun isReified(variable: TypeVariableMarker): Boolean
@@ -220,6 +222,10 @@ class ResultTypeResolver(

                if (typesWithoutStubs.isNotEmpty()) {
                    commonSuperType = computeCommonSuperType(typesWithoutStubs)
+                } else if (outerSystemVariablesPrefixSize > 0) {
+                    // outerSystemVariablesPrefixSize > 0 only for PCLA (K2)
+                    @OptIn(K2Only::class)
+                    commonSuperType = createSubstitutionFromSubtypingStubTypesToTypeVariables().safeSubstitute(commonSuperType)
                }
            }

@@ -273,6 +279,14 @@ class ResultTypeResolver(
        }

        if (!atLeastOneProper) return emptyList()
+
+        // PCLA slow path
+        // We only allow using TVs fixation for nested PCLA calls
+        if (outerSystemVariablesPrefixSize > 0) {
+            val notFixedToStubTypesSubstitutor = buildNotFixedVariablesToStubTypesSubstitutor()
+            return lowerConstraintTypes.map { notFixedToStubTypesSubstitutor.safeSubstitute(it) }
+        }
+
        if (!atLeastOneNonProper) return lowerConstraintTypes

        val notFixedToStubTypesSubstitutor = buildNotFixedVariablesToStubTypesSubstitutor()
@@ -336,7 +350,7 @@ class ResultTypeResolver(
    }

    private fun Context.isProperTypeForFixation(type: KotlinTypeMarker): Boolean =
-        isProperTypeForFixation(type) { isProperType(it) }
+        isProperTypeForFixation(type, notFixedTypeVariables.keys) { isProperType(it) }

    private fun findResultIfThereIsEqualsConstraint(c: Context, variableWithConstraints: VariableWithConstraints): KotlinTypeMarker? =
        with(c) {
@@ -7,15 +7,22 @@ package org.jetbrains.kotlin.resolve.calls.inference.components

 import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
-import org.jetbrains.kotlin.types.model.*
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.model.TypeConstructorMarker
+import org.jetbrains.kotlin.types.model.freshTypeConstructor
+import org.jetbrains.kotlin.types.model.typeConstructor
 import org.jetbrains.kotlin.utils.SmartSet

 class TypeVariableDependencyInformationProvider(
    private val notFixedTypeVariables: Map<TypeConstructorMarker, VariableWithConstraints>,
    private val postponedKtPrimitives: List<PostponedResolvedAtomMarker>,
    private val topLevelType: KotlinTypeMarker?,
-    private val typeSystemContext: TypeSystemInferenceExtensionContext
+    private val typeSystemContext: VariableFixationFinder.Context
 ) {
+
+    private val outerTypeVariables: Set<TypeConstructorMarker>? =
+        typeSystemContext.outerTypeVariables
+
    /*
     * Not oriented edges
     * TypeVariable(A) has UPPER(Function1<TypeVariable(B), R>) => A and B are related deeply
@@ -41,7 +48,16 @@ class TypeVariableDependencyInformationProvider(
        computeRelatedToTopLevelType()
    }

-    fun isVariableRelatedToTopLevelType(variable: TypeConstructorMarker) = relatedToTopLevelType.contains(variable)
+    fun isVariableRelatedToTopLevelType(variable: TypeConstructorMarker) =
+        relatedToTopLevelType.contains(variable)
+
+
+    fun isRelatedToOuterTypeVariable(variable: TypeConstructorMarker): Boolean {
+        val outerTypeVariables = outerTypeVariables ?: return false
+        val myDependent = getDeeplyDependentVariables(variable) ?: return false
+        return myDependent.any { it in outerTypeVariables }
+    }
+
    fun isVariableRelatedToAnyOutputType(variable: TypeConstructorMarker) = relatedToAllOutputTypes.contains(variable)

    fun getDeeplyDependentVariables(variable: TypeConstructorMarker) = deepTypeVariableDependencies[variable]
@@ -27,25 +27,40 @@ class VariableFixationFinder(
        val fixedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker>
        val postponedTypeVariables: List<TypeVariableMarker>
        val constraintsFromAllForkPoints: MutableList<Pair<IncorporationConstraintPosition, ForkPointData>>
+        val allTypeVariables: Map<TypeConstructorMarker, TypeVariableMarker>

        /**
-         * If not null, that property means that we should assume temporary
-         * `allTypeVariables.keys.minus(typeVariablesThatAreNotCountedAsProperTypes)` as proper types when fixating some variables.
+         * See [org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage.outerSystemVariablesPrefixSize]
+         */
+        val outerSystemVariablesPrefixSize: Int
+
+        val outerTypeVariables: Set<TypeConstructorMarker>?
+            get() =
+                when {
+                    outerSystemVariablesPrefixSize > 0 -> allTypeVariables.keys.take(outerSystemVariablesPrefixSize).toSet()
+                    else -> null
+                }
+
+        /**
+         * If not null, that property means that we should assume temporary them all as proper types when fixating some variables.
         *
         * By default, if that property is null, we assume all `allTypeVariables` as not proper.
         *
         * Currently, that is only used for `provideDelegate` resolution, see
         * [org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer.fixInnerVariablesForProvideDelegateIfNeeded]
         */
-        val typeVariablesThatAreNotCountedAsProperTypes: Set<TypeConstructorMarker>?
+        val typeVariablesThatAreCountedAsProperTypes: Set<TypeConstructorMarker>?

        fun isReified(variable: TypeVariableMarker): Boolean
    }

-    data class VariableForFixation(
+    class VariableForFixation(
        val variable: TypeConstructorMarker,
-        val hasProperConstraint: Boolean,
-    )
+        private val hasProperConstraint: Boolean,
+        private val hasDependencyOnOuterTypeVariable: Boolean = false,
+    ) {
+        val isReady: Boolean get() = hasProperConstraint && !hasDependencyOnOuterTypeVariable
+    }

    fun findFirstVariableForFixation(
        c: Context,
@@ -53,11 +68,13 @@ class VariableFixationFinder(
        postponedKtPrimitives: List<PostponedResolvedAtomMarker>,
        completionMode: ConstraintSystemCompletionMode,
        topLevelType: KotlinTypeMarker,
-    ): VariableForFixation? = c.findTypeVariableForFixation(allTypeVariables, postponedKtPrimitives, completionMode, topLevelType)
+    ): VariableForFixation? =
+        c.findTypeVariableForFixation(allTypeVariables, postponedKtPrimitives, completionMode, topLevelType)

    enum class TypeVariableFixationReadiness {
        FORBIDDEN,
        WITHOUT_PROPER_ARGUMENT_CONSTRAINT, // proper constraint from arguments -- not from upper bound for type parameters
+        OUTER_TYPE_VARIABLE_DEPENDENCY,
        READY_FOR_FIXATION_DECLARED_UPPER_BOUND_WITH_SELF_TYPES,
        WITH_COMPLEX_DEPENDENCY, // if type variable T has constraint with non fixed type variable inside (non-top-level): T <: Foo<S>
        ALL_CONSTRAINTS_TRIVIAL_OR_NON_PROPER, // proper trivial constraint from arguments, Nothing <: T
@@ -85,6 +102,7 @@ class VariableFixationFinder(
        isTypeInferenceForSelfTypesSupported && areAllProperConstraintsSelfTypeBased(variable) ->
            TypeVariableFixationReadiness.READY_FOR_FIXATION_DECLARED_UPPER_BOUND_WITH_SELF_TYPES
        !variableHasProperArgumentConstraints(variable) -> TypeVariableFixationReadiness.WITHOUT_PROPER_ARGUMENT_CONSTRAINT
+        dependencyProvider.isRelatedToOuterTypeVariable(variable) -> TypeVariableFixationReadiness.OUTER_TYPE_VARIABLE_DEPENDENCY
        hasDependencyToOtherTypeVariables(variable) -> TypeVariableFixationReadiness.WITH_COMPLEX_DEPENDENCY
        // TODO: Consider removing this kind of readiness, see KT-63032
        allConstraintsTrivialOrNonProper(variable) -> TypeVariableFixationReadiness.ALL_CONSTRAINTS_TRIVIAL_OR_NON_PROPER
@@ -152,7 +170,7 @@ class VariableFixationFinder(
        if (allTypeVariables.isEmpty()) return null

        val dependencyProvider = TypeVariableDependencyInformationProvider(
-            notFixedTypeVariables, postponedArguments, topLevelType.takeIf { completionMode == PARTIAL }, this
+            notFixedTypeVariables, postponedArguments, topLevelType.takeIf { completionMode == PARTIAL }, this,
        )

        val candidate =
@@ -161,6 +179,8 @@ class VariableFixationFinder(
        return when (getTypeVariableReadiness(candidate, dependencyProvider)) {
            TypeVariableFixationReadiness.FORBIDDEN -> null
            TypeVariableFixationReadiness.WITHOUT_PROPER_ARGUMENT_CONSTRAINT -> VariableForFixation(candidate, false)
+            TypeVariableFixationReadiness.OUTER_TYPE_VARIABLE_DEPENDENCY ->
+                VariableForFixation(candidate, hasProperConstraint = true, hasDependencyOnOuterTypeVariable = true)

            else -> VariableForFixation(candidate, true)
        }
@@ -190,12 +210,13 @@ class VariableFixationFinder(
                && !c.isNullabilityConstraint

    private fun Context.isProperType(type: KotlinTypeMarker): Boolean =
-        isProperTypeForFixation(type) { t -> !t.contains { isNotFixedRelevantVariable(it) } }
+        isProperTypeForFixation(type, notFixedTypeVariables.keys) { t -> !t.contains { isNotFixedRelevantVariable(it) } }

    private fun Context.isNotFixedRelevantVariable(it: KotlinTypeMarker): Boolean {
-        if (!notFixedTypeVariables.containsKey(it.typeConstructor())) return false
-        if (typeVariablesThatAreNotCountedAsProperTypes == null) return true
-        return typeVariablesThatAreNotCountedAsProperTypes!!.contains(it.typeConstructor())
+        val key = it.typeConstructor()
+        if (!notFixedTypeVariables.containsKey(key)) return false
+        if (typeVariablesThatAreCountedAsProperTypes?.contains(key) == true) return false
+        return true
    }

    private fun Context.isReified(variable: TypeConstructorMarker): Boolean =
@@ -235,8 +256,25 @@ class VariableFixationFinder(
    }
 }

-inline fun TypeSystemInferenceExtensionContext.isProperTypeForFixation(type: KotlinTypeMarker, isProper: (KotlinTypeMarker) -> Boolean) =
-    isProper(type) && extractProjectionsForAllCapturedTypes(type).all(isProper)
+/**
+ * Returns `false` for fixed type variables types even if `isProper(type) == true`
+ * Thus allowing only non-TVs types to be used for fixation on top level.
+ * While this limitation is important, it doesn't really limit final results because when we have a constraint like T <: E or E <: T
+ * and we're going to fix T into E, we assume that if E has some other constraints, they are being incorporated to T, so we would choose
+ * them instead of E itself.
+ */
+inline fun TypeSystemInferenceExtensionContext.isProperTypeForFixation(
+    type: KotlinTypeMarker,
+    notFixedTypeVariables: Set<TypeConstructorMarker>,
+    isProper: (KotlinTypeMarker) -> Boolean
+): Boolean {
+    // We don't allow fixing T into any top-level TV type, like T := F or T := F & Any
+    // Even if F is considered as a proper by `isProper` (e.g., it belongs to an outer CS)
+    // But at the same time, we don't forbid fixing into T := MutableList<F>
+    if (type.typeConstructor() in notFixedTypeVariables) return false
+
+    return isProper(type) && extractProjectionsForAllCapturedTypes(type).all(isProper)
+}

 fun TypeSystemInferenceExtensionContext.extractProjectionsForAllCapturedTypes(baseType: KotlinTypeMarker): Set<KotlinTypeMarker> {
    if (baseType.isFlexible()) {
@@ -49,6 +49,8 @@ interface ConstraintStorage {
    val constraintsFromAllForkPoints: List<Pair<IncorporationConstraintPosition, ForkPointData>>

    /**
+     *  Outer system for a call means some set of variables defined beside it/its arguments
+     *
     *  In case some candidate's CS is built in the context of some outer CS, first [outerSystemVariablesPrefixSize] in the list
     *  of [allTypeVariables] belong to the outer CS.
     *
@@ -62,6 +64,8 @@ interface ConstraintStorage {
     */
    val outerSystemVariablesPrefixSize: Int

+    val usesOuterCs: Boolean
+
    object Empty : ConstraintStorage {
        override val allTypeVariables: Map<TypeConstructorMarker, TypeVariableMarker> get() = emptyMap()
        override val notFixedTypeVariables: Map<TypeConstructorMarker, VariableWithConstraints> get() = emptyMap()
@@ -77,6 +81,8 @@ interface ConstraintStorage {
        override val constraintsFromAllForkPoints: List<Pair<IncorporationConstraintPosition, ForkPointData>> = emptyList()

        override val outerSystemVariablesPrefixSize: Int get() = 0
+
+        override val usesOuterCs: Boolean get() = false
    }
 }

@@ -221,6 +221,14 @@ class MutableVariableWithConstraints private constructor(
        }
    }

+    fun runConstraintsSimplification() {
+        val currentState = constraints.toList()
+        mutableConstraints.apply {
+            clear()
+            addAll(currentState)
+        }
+    }
+
    private fun isUsefulConstraint(constraint: Constraint, equalityConstraints: Map<Int, List<Constraint>>): Boolean {
        if (constraint.kind == ConstraintKind.EQUALITY) return true
        return equalityConstraints[constraint.typeHashCode]?.none { it.type == constraint.type } ?: true
@@ -251,4 +259,15 @@ internal class MutableConstraintStorage : ConstraintStorage {
    override val constraintsFromAllForkPoints: MutableList<Pair<IncorporationConstraintPosition, ForkPointData>> = SmartList()

    override var outerSystemVariablesPrefixSize: Int = 0
+
+    override var usesOuterCs: Boolean = false
+
+    @AssertionsOnly
+    internal var outerCS: ConstraintStorage? = null
 }
+
+/**
+ * Annotated member is used only for assertion purposes and does not affect semantics
+ */
+@RequiresOptIn
+annotation class AssertionsOnly
@@ -41,7 +41,7 @@ class NewConstraintSystemImpl(
    private val properTypesCache: MutableSet<KotlinTypeMarker> = SmartSet.create()
    private val notProperTypesCache: MutableSet<KotlinTypeMarker> = SmartSet.create()
    private val intersectionTypesCache: MutableMap<Collection<KotlinTypeMarker>, EmptyIntersectionTypeInfo?> = mutableMapOf()
-    override var typeVariablesThatAreNotCountedAsProperTypes: Set<TypeConstructorMarker>? = null
+    override var typeVariablesThatAreCountedAsProperTypes: Set<TypeConstructorMarker>? = null

    private var couldBeResolvedWithUnrestrictedBuilderInference: Boolean = false

@@ -51,23 +51,26 @@ class NewConstraintSystemImpl(
     * @see [org.jetbrains.kotlin.resolve.calls.inference.components.VariableFixationFinder.Context.typeVariablesThatAreNotCountedAsProperTypes]
     * @see [org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer.fixInnerVariablesForProvideDelegateIfNeeded]
     */
-    fun withTypeVariablesThatAreNotCountedAsProperTypes(typeVariables: Set<TypeConstructorMarker>, block: () -> Unit) {
+    @K2Only
+    override fun <R> withTypeVariablesThatAreCountedAsProperTypes(typeVariables: Set<TypeConstructorMarker>, block: () -> R): R {
        checkState(State.BUILDING)
        // Cleaning cache is necessary because temporarily we change the meaning of what does "proper type" mean
        properTypesCache.clear()
        notProperTypesCache.clear()

-        require(typeVariablesThatAreNotCountedAsProperTypes == null) {
+        require(typeVariablesThatAreCountedAsProperTypes == null) {
            "Currently there should be no nested withDisallowingOnlyThisTypeVariablesForProperTypes calls"
        }

-        typeVariablesThatAreNotCountedAsProperTypes = typeVariables
+        typeVariablesThatAreCountedAsProperTypes = typeVariables

-        block()
+        val result = block()

-        typeVariablesThatAreNotCountedAsProperTypes = null
+        typeVariablesThatAreCountedAsProperTypes = null
        properTypesCache.clear()
        notProperTypesCache.clear()
+
+        return result
    }

    private enum class State {
@@ -221,6 +224,7 @@ class NewConstraintSystemImpl(
    // ConstraintSystemBuilder
    private fun transactionRegisterVariable(variable: TypeVariableMarker) {
        if (state != State.TRANSACTION) return
+        if (variable.freshTypeConstructor() in storage.allTypeVariables) return
        typeVariablesTransaction.add(variable)
    }

@@ -302,13 +306,25 @@ class NewConstraintSystemImpl(
        }

    fun addOuterSystem(outerSystem: ConstraintStorage) {
-        addOtherSystem(outerSystem)
+        require(!storage.usesOuterCs)
+
+        storage.usesOuterCs = true
        storage.outerSystemVariablesPrefixSize = outerSystem.allTypeVariables.size
+        @OptIn(AssertionsOnly::class)
+        storage.outerCS = outerSystem
+
+        addOtherSystem(outerSystem, isAddingOuter = true)
    }

-    fun setBaseSystem(outerSystem: ConstraintStorage) {
-        addOtherSystem(outerSystem)
-        storage.outerSystemVariablesPrefixSize = outerSystem.outerSystemVariablesPrefixSize
+    @K2Only
+    fun setBaseSystem(baseSystem: ConstraintStorage) {
+        require(storage.allTypeVariables.isEmpty())
+        storage.usesOuterCs = baseSystem.usesOuterCs
+        storage.outerSystemVariablesPrefixSize = baseSystem.outerSystemVariablesPrefixSize
+        @OptIn(AssertionsOnly::class)
+        storage.outerCS = (baseSystem as? MutableConstraintStorage)?.outerCS
+
+        addOtherSystem(baseSystem)
    }

    fun prepareForGlobalCompletion() {
@@ -317,6 +333,17 @@ class NewConstraintSystemImpl(
    }

    override fun addOtherSystem(otherSystem: ConstraintStorage) {
+        addOtherSystem(otherSystem, isAddingOuter = false)
+    }
+
+    fun replaceContentWith(otherSystem: ConstraintStorage) {
+        addOtherSystem(otherSystem, isAddingOuter = false, clearNotFixedTypeVariables = true)
+    }
+
+    private fun addOtherSystem(otherSystem: ConstraintStorage, isAddingOuter: Boolean, clearNotFixedTypeVariables: Boolean = false) {
+        @OptIn(AssertionsOnly::class)
+        runOuterCSRelatedAssertions(otherSystem, isAddingOuter)
+
        if (otherSystem.allTypeVariables.isNotEmpty()) {
            otherSystem.allTypeVariables.forEach {
                transactionRegisterVariable(it.value)
@@ -324,16 +351,47 @@ class NewConstraintSystemImpl(
            storage.allTypeVariables.putAll(otherSystem.allTypeVariables)
            notProperTypesCache.clear()
        }
+
+        // `clearNotFixedTypeVariables` means that we're mostly replacing the content, thus we need to remove variables that have been fixed
+        // in `otherSystem` from `this.notFixedTypeVariables`, too
+        if (clearNotFixedTypeVariables) {
+            notFixedTypeVariables.clear()
+        }
+
        for ((variable, constraints) in otherSystem.notFixedTypeVariables) {
            notFixedTypeVariables[variable] = MutableVariableWithConstraints(this, constraints)
        }
-        storage.initialConstraints.addAll(otherSystem.initialConstraints)
+
+        val currentInitialConstraints = storage.initialConstraints.toSet()
+
+        otherSystem.initialConstraints.filterTo(storage.initialConstraints) {
+            it !in currentInitialConstraints
+        }
+
        storage.maxTypeDepthFromInitialConstraints =
            max(storage.maxTypeDepthFromInitialConstraints, otherSystem.maxTypeDepthFromInitialConstraints)
        storage.errors.addAll(otherSystem.errors)
        storage.fixedTypeVariables.putAll(otherSystem.fixedTypeVariables)
        storage.postponedTypeVariables.addAll(otherSystem.postponedTypeVariables)
        storage.constraintsFromAllForkPoints.addAll(otherSystem.constraintsFromAllForkPoints)
+
+    }
+
+
+    @AssertionsOnly
+    private fun runOuterCSRelatedAssertions(otherSystem: ConstraintStorage, isAddingOuter: Boolean) {
+        if (!otherSystem.usesOuterCs) return
+
+        // When integrating a child system back, it's ok that for root CS, `storage.usesOuterCs == false`
+        if ((otherSystem as? MutableConstraintStorage)?.outerCS === storage) return
+
+        require(storage.usesOuterCs)
+
+        if (!isAddingOuter) {
+            require(storage.outerSystemVariablesPrefixSize == otherSystem.outerSystemVariablesPrefixSize) {
+                "Expected to be ${otherSystem.outerSystemVariablesPrefixSize}, but ${storage.outerSystemVariablesPrefixSize} found"
+            }
+        }
    }

    // ResultTypeResolver.Context, ConstraintSystemBuilder
@@ -357,11 +415,11 @@ class NewConstraintSystemImpl(
                it

            if (typeToCheck == null) return@contains false
-            if (typeVariablesThatAreNotCountedAsProperTypes != null) {
-                return@contains typeVariablesThatAreNotCountedAsProperTypes!!.contains(typeToCheck.typeConstructor())
+            if (typeVariablesThatAreCountedAsProperTypes?.contains(typeToCheck.typeConstructor()) == true) {
+                return@contains false
            }

-            storage.allTypeVariables.containsKey(typeToCheck.typeConstructor())
+            return@contains storage.allTypeVariables.containsKey(typeToCheck.typeConstructor())
        }

    override fun isTypeVariable(type: KotlinTypeMarker): Boolean {
@@ -689,7 +747,7 @@ class NewConstraintSystemImpl(
        return buildCurrentSubstitutor(emptyMap())
    }

-    override fun buildCurrentSubstitutor(additionalBindings: Map<TypeConstructorMarker, StubTypeMarker>): TypeSubstitutorMarker {
+    override fun buildCurrentSubstitutor(additionalBindings: Map<TypeConstructorMarker, KotlinTypeMarker>): TypeSubstitutorMarker {
        checkState(State.BUILDING, State.COMPLETION, State.TRANSACTION)
        return storage.buildCurrentSubstitutor(this, additionalBindings)
    }
@@ -715,6 +773,8 @@ class NewConstraintSystemImpl(
        return storage
    }

+    val usesOuterCs: Boolean get() = storage.usesOuterCs
+
    // PostponedArgumentsAnalyzer.Context
    override fun hasUpperOrEqualUnitConstraint(type: KotlinTypeMarker): Boolean {
        checkState(State.BUILDING, State.COMPLETION, State.FREEZED)