FIR: support don't use builder inference if possible

In this commit we upgrade FIR builder inference logic from
the compiler version to 1.7. FIR-based compiler now works with
"don't use builder inference" flag always ON and supports switching
the flag "use builder inference only if needed". To do it,
ContraintSystemCompleter (FIR) and KotlinConstraintSystemCompleter (FE 1.0)
are made similar with extracting some common parts into
ConstraintSystemCompletionContext.

Test status: one BB test fails after this commit (KT-49285).
Also we have a crush in DFA logic in FIR bootstrap test and somehow
questionable behavior in FIR diagnostic test. However,
two BB tests were fixed, the 3rd case from KT-49925 were also fixed.

#KT-49925 Fixed

compiler/fir/analysis-tests/testData/resolveWithStdlib/complexPostponedCfg.dot

@@ -126,6 +126,7 @@ digraph complexPostponedCfg_kt {
     35 -> {36} [color=green];
     36 -> {37};
     37 -> {38};
+    38 -> {10} [color=red];
     38 -> {22} [color=green];
     39 -> {40};
     40 -> {41};

compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/builderInferenceAndCoercionToUnit.fir.txt

@@ -9,7 +9,7 @@ FILE: builderInferenceAndCoercionToUnit.kt
 
     }
     public final fun test(strings: R|kotlin/collections/List<kotlin/String>|): R|kotlin/Unit| {
-        lval dropDown: R|DropDownComponent<kotlin/String>| = R|/DropDownComponent.DropDownComponent|<R|kotlin/String|>(initialValues = R|kotlin/collections/buildList|<R|kotlin/String|>(<L> = buildList@fun R|kotlin/collections/MutableList<kotlin/String>|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>  {
+        lval dropDown: R|DropDownComponent<kotlin/Any>| = R|/DropDownComponent.DropDownComponent|<R|kotlin/Any|>(initialValues = R|kotlin/collections/buildList|<R|kotlin/Any|>(<L> = buildList@fun R|kotlin/collections/MutableList<kotlin/Any>|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>  {
             this@R|special/anonymous|.R|SubstitutionOverride<kotlin/collections/MutableList.addAll: R|kotlin/Boolean|>|(R|<local>/strings|)
         }
         ))

compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/problems/buildListLazy.fir.txt

@@ -28,8 +28,8 @@ FILE: buildListLazy.kt
                 this@R|special/anonymous|.R|kotlin/collections/plusAssign|<R|NameAndSafeValue|>(R|/NameAndSafeValue.NameAndSafeValue|(R|<local>/name|, R|<local>/value|))
             }
             )
-            this@R|special/anonymous|.<Inapplicable(INAPPLICABLE): kotlin/collections/sortBy>#<R|NameAndSafeValue|, R|ERROR CLASS: Cannot infer argument for type parameter R|>(<L> = sortBy@fun <anonymous>(it: R|NameAndSafeValue|): <ERROR TYPE REF: Cannot infer argument for type parameter R> <inline=CrossInline, kind=UNKNOWN>  {
+            this@R|special/anonymous|.R|kotlin/collections/sortBy|<R|NameAndSafeValue|, R|kotlin/String|>(<L> = sortBy@fun <anonymous>(it: R|NameAndSafeValue|): R|kotlin/String?| <inline=CrossInline, kind=UNKNOWN>  {
-                ^ R|<local>/it|.<Unresolved name: name>#
+                ^ R|<local>/it|.R|/NameAndSafeValue.name|
             }
             )
         }

compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/problems/buildListLazy.kt

@@ -7,6 +7,6 @@ private val environment: List<NameAndSafeValue> by lazy {
         getEnv().forEach { (name, value) ->
             this += NameAndSafeValue(name, value)
         }
-        sortBy { <!ARGUMENT_TYPE_MISMATCH!>it.<!UNRESOLVED_REFERENCE!>name<!><!> }
+        sortBy { it.name }
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolutionStages.kt

@@ -497,6 +497,9 @@ internal object PostponedVariablesInitializerResolutionStage : ResolutionStage()
             if (!parameter.hasBuilderInferenceAnnotation()) continue
             val type = parameter.returnTypeRef.coneType
             val receiverType = type.receiverType(callInfo.session) ?: continue
+            val dontUseBuilderInferenceIfPossible =
+                context.session.languageVersionSettings.supportsFeature(LanguageFeature.UseBuilderInferenceOnlyIfNeeded)
+            if (dontUseBuilderInferenceIfPossible) continue
 
             for (freshVariable in candidate.freshVariables) {
                 if (candidate.csBuilder.isPostponedTypeVariable(freshVariable)) continue

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

@@ -5,10 +5,12 @@
 
 package org.jetbrains.kotlin.fir.resolve.inference
 
+import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
@@ -34,10 +36,11 @@ import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
-class ConstraintSystemCompleter(private val components: BodyResolveComponents, private val context: BodyResolveContext) {
+class ConstraintSystemCompleter(components: BodyResolveComponents, private val context: BodyResolveContext) {
     private val inferenceComponents = components.session.inferenceComponents
     val variableFixationFinder = inferenceComponents.variableFixationFinder
     private val postponedArgumentsInputTypesResolver = inferenceComponents.postponedArgumentInputTypesResolver
+    private val languageVersionSettings = components.session.languageVersionSettings
 
     fun complete(
         c: ConstraintSystemCompletionContext,
@@ -47,35 +50,56 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
         context: ResolutionContext,
         collectVariablesFromContext: Boolean = false,
         analyze: (PostponedResolvedAtom) -> Unit
-    ) = with(c) {
+    ) = c.runCompletion(completionMode, topLevelAtoms, candidateReturnType, context, collectVariablesFromContext, analyze)
-        val topLevelTypeVariables = candidateReturnType.extractTypeVariables()
 
-        c.processForkConstraints()
+    private fun ConstraintSystemCompletionContext.runCompletion(
+        completionMode: ConstraintSystemCompletionMode,
+        topLevelAtoms: List<FirStatement>,
+        topLevelType: ConeKotlinType,
+        context: ResolutionContext,
+        collectVariablesFromContext: Boolean = false,
+        analyze: (PostponedResolvedAtom) -> Unit
+    ) {
+        val topLevelTypeVariables = topLevelType.extractTypeVariables()
+
+        // NB: it's called in ConstraintSystemForks resolution stage by FE 1.0
+        processForkConstraints()
 
         completion@ while (true) {
-            val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms) // TODO: This is very slow
+            // TODO: This is very slow
+            val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
 
-            if (completionMode == ConstraintSystemCompletionMode.UNTIL_FIRST_LAMBDA && hasLambdaToAnalyze(postponedArguments)) return
+            if (completionMode == ConstraintSystemCompletionMode.UNTIL_FIRST_LAMBDA && hasLambdaToAnalyze(
+                    languageVersionSettings,
+                    postponedArguments
+                )
+            ) return
 
-            // Stage 1
+            // Stage 1: analyze postponed arguments with fixed parameter types
-            if (analyzeArgumentWithFixedParameterTypes(postponedArguments, analyze))
+            if (analyzeArgumentWithFixedParameterTypes(languageVersionSettings, postponedArguments, analyze))
                 continue
 
-            val someVariableIsReadyForFixation = isAnyVariableReadyForFixation(
+            val isThereAnyReadyForFixationVariable = variableFixationFinder.findFirstVariableForFixation(
-                completionMode, topLevelAtoms, candidateReturnType, collectVariablesFromContext, postponedArguments
+                this,
-            )
+                getOrderedAllTypeVariables(collectVariablesFromContext, topLevelAtoms),
+                postponedArguments,
+                completionMode,
+                topLevelType
+            ) != null
 
-            if (postponedArguments.isEmpty() && !someVariableIsReadyForFixation)
+            // If there aren't any postponed arguments and ready for fixation variables, then completion isn't needed: nothing to do
+            if (postponedArguments.isEmpty() && !isThereAnyReadyForFixationVariable)
                 break
 
             val postponedArgumentsWithRevisableType = postponedArguments
                 .filterIsInstance<PostponedAtomWithRevisableExpectedType>()
+                // NB: FE 1.0 does not perform this check
                 .filter { it.revisedExpectedType == null }
             val dependencyProvider =
-                TypeVariableDependencyInformationProvider(notFixedTypeVariables, postponedArguments, candidateReturnType, this)
+                TypeVariableDependencyInformationProvider(notFixedTypeVariables, postponedArguments, topLevelType, this)
 
-            // Stage 2
+            // Stage 2: collect parameter types for postponed arguments
-            val newExpectedTypeWasBuilt = postponedArgumentsInputTypesResolver.collectParameterTypesAndBuildNewExpectedTypes(
+            val wasBuiltNewExpectedTypeForSomeArgument = postponedArgumentsInputTypesResolver.collectParameterTypesAndBuildNewExpectedTypes(
                 this,
                 postponedArgumentsWithRevisableType,
                 completionMode,
@@ -83,19 +107,21 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
                 topLevelTypeVariables
             )
 
-            if (newExpectedTypeWasBuilt)
+            if (wasBuiltNewExpectedTypeForSomeArgument)
                 continue
 
             if (completionMode == ConstraintSystemCompletionMode.FULL) {
-                // Stage 3
+                // Stage 3: fix variables for parameter types of all postponed arguments
                 for (argument in postponedArguments) {
                     val variableWasFixed = postponedArgumentsInputTypesResolver.fixNextReadyVariableForParameterTypeIfNeeded(
                         this,
                         argument,
                         postponedArguments,
-                        candidateReturnType,
+                        topLevelType,
                         dependencyProvider,
-                    ) { // atom provided here is used only inside constraint positions, omitting right now
+                    ) {
+                        // NB: FE 1.0 calls findResolvedAtomBy here
+                        // atom provided here is used only inside constraint positions, omitting right now
                         null
                     }
 
@@ -103,10 +129,11 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
                         continue@completion
                 }
 
-                // Stage 4
+                // Stage 4: create atoms with revised expected types if needed
                 for (argument in postponedArgumentsWithRevisableType) {
-                    val argumentWasTransformed =
+                    val argumentWasTransformed = transformToAtomWithNewFunctionalExpectedType(
-                        transformToAtomWithNewFunctionalExpectedType(this, context, argument)
+                        this, context, argument
+                    )
 
                     if (argumentWasTransformed)
                         continue@completion
@@ -114,17 +141,27 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
             }
 
             // Stage 5: analyze the next ready postponed argument
-            if (analyzeNextReadyPostponedArgument(postponedArguments, completionMode, analyze))
+            if (analyzeNextReadyPostponedArgument(languageVersionSettings, postponedArguments, completionMode, analyze))
                 continue
 
-            // Stage 6: fix type variables – fix if possible or report not enough information (if completion mode is full)
+            // Stage 6: fix next ready type variable with proper constraints
-            val variableWasFixed = fixVariablesOrReportNotEnoughInformation(
+            if (fixNextReadyVariable(completionMode, topLevelAtoms, topLevelType, collectVariablesFromContext, postponedArguments))
-                completionMode, topLevelAtoms, candidateReturnType, collectVariablesFromContext, postponedArguments
+                continue
+
+            // Stage 7: try to complete call with the builder inference if there are uninferred type variables
+            val areThereAppearedProperConstraintsForSomeVariable = tryToCompleteWithBuilderInference(
+                completionMode, topLevelAtoms, topLevelType, postponedArguments, collectVariablesFromContext, analyze
             )
-            if (variableWasFixed)
+
+            if (areThereAppearedProperConstraintsForSomeVariable)
                 continue
 
-            // Stage 7: force analysis of remaining not analyzed postponed arguments and rerun stages if there are
+            // Stage 8: report "not enough information" for uninferred type variables
+            reportNotEnoughTypeInformation(
+                completionMode, topLevelAtoms, topLevelType, collectVariablesFromContext, postponedArguments
+            )
+
+            // Stage 9: force analysis of remaining not analyzed postponed arguments and rerun stages if there are
             if (completionMode == ConstraintSystemCompletionMode.FULL) {
                 if (analyzeRemainingNotAnalyzedPostponedArgument(postponedArguments, analyze))
                     continue
@@ -134,37 +171,44 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
         }
     }
 
-    private fun ConstraintSystemCompletionContext.analyzeArgumentWithFixedParameterTypes(
+    private fun ConstraintSystemCompletionContext.tryToCompleteWithBuilderInference(
-        postponedArguments: List<PostponedResolvedAtom>,
-        analyze: (PostponedResolvedAtom) -> Unit
-    ): Boolean {
-        val argumentWithFixedOrPostponedInputTypes = findPostponedArgumentWithFixedOrPostponedInputTypes(postponedArguments)
-
-        if (argumentWithFixedOrPostponedInputTypes != null) {
-            analyze(argumentWithFixedOrPostponedInputTypes)
-            return true
-        }
-
-        return false
-    }
-
-    private fun ConstraintSystemCompletionContext.findPostponedArgumentWithFixedOrPostponedInputTypes(postponedArguments: List<PostponedResolvedAtom>) =
-        postponedArguments.firstOrNull { argument -> argument.inputTypes.all { containsOnlyFixedOrPostponedVariables(it) } }
-
-    private fun ConstraintSystemCompletionContext.isAnyVariableReadyForFixation(
         completionMode: ConstraintSystemCompletionMode,
         topLevelAtoms: List<FirStatement>,
         topLevelType: ConeKotlinType,
-        collectVariablesFromContext: Boolean,
         postponedArguments: List<PostponedResolvedAtom>,
+        collectVariablesFromContext: Boolean,
+        analyze: (PostponedResolvedAtom) -> Unit
     ): Boolean {
-        return variableFixationFinder.findFirstVariableForFixation(
+        if (completionMode == ConstraintSystemCompletionMode.PARTIAL) return false
-            this,
+
-            getOrderedAllTypeVariables(this, topLevelAtoms, collectVariablesFromContext),
+        // If we use the builder inference anyway (if the annotation is presented), then we are already analysed builder inference lambdas
-            postponedArguments,
+        if (!languageVersionSettings.supportsFeature(LanguageFeature.UseBuilderInferenceOnlyIfNeeded)) return false
-            completionMode,
+
-            topLevelType
+        val lambdaArguments = postponedArguments.filterIsInstance<ResolvedLambdaAtom>().takeIf { it.isNotEmpty() } ?: return false
-        ) != null
+
+        // We assume useBuilderInferenceWithoutAnnotation = true for FIR
+
+        for (argument in lambdaArguments) {
+            val notFixedInputTypeVariables = argument.inputTypes
+                .map { it.extractTypeVariables() }.flatten().filter { it !in fixedTypeVariables }
+
+            if (notFixedInputTypeVariables.isEmpty()) continue
+
+            for (variable in notFixedInputTypeVariables) {
+                getBuilder().markPostponedVariable(notFixedTypeVariables.getValue(variable).typeVariable)
+            }
+
+            analyze(argument)
+        }
+
+        val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
+            this, getOrderedAllTypeVariables(collectVariablesFromContext, topLevelAtoms), postponedArguments, completionMode, topLevelType
+        )
+
+        // continue completion (rerun stages) only if ready for fixation variables with proper constraints have appeared
+        // (after analysing a lambda with the builder inference)
+        // otherwise we don't continue and report "not enough type information" error
+        return variableForFixation?.hasProperConstraint == true
     }
 
     private fun transformToAtomWithNewFunctionalExpectedType(
@@ -186,81 +230,73 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
         return true
     }
 
-    private fun ConstraintSystemCompletionContext.analyzeNextReadyPostponedArgument(
+    private fun ConstraintSystemCompletionContext.fixNextReadyVariable(
-        postponedArguments: List<PostponedResolvedAtom>,
-        completionMode: ConstraintSystemCompletionMode,
-        analyze: (PostponedResolvedAtom) -> Unit,
-    ): Boolean {
-        if (completionMode == ConstraintSystemCompletionMode.FULL) {
-            val argumentWithTypeVariableAsExpectedType = findPostponedArgumentWithRevisableExpectedType(postponedArguments)
-
-            if (argumentWithTypeVariableAsExpectedType != null) {
-                analyze(argumentWithTypeVariableAsExpectedType)
-                return true
-            }
-        }
-
-        return analyzeArgumentWithFixedParameterTypes(postponedArguments, analyze)
-    }
-
-    // Avoiding smart cast from filterIsInstanceOrNull looks dirty
-    private fun findPostponedArgumentWithRevisableExpectedType(postponedArguments: List<PostponedResolvedAtom>): PostponedResolvedAtom? =
-        postponedArguments.firstOrNull { argument -> argument is PostponedAtomWithRevisableExpectedType }
-
-    private fun ConstraintSystemCompletionContext.fixVariablesOrReportNotEnoughInformation(
         completionMode: ConstraintSystemCompletionMode,
         topLevelAtoms: List<FirStatement>,
         topLevelType: ConeKotlinType,
         collectVariablesFromContext: Boolean,
         postponedArguments: List<PostponedResolvedAtom>,
     ): Boolean {
-        while (true) {
+        val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
-            val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
+            this,
-                this,
+            getOrderedAllTypeVariables(collectVariablesFromContext, topLevelAtoms),
-                getOrderedAllTypeVariables(this, topLevelAtoms, collectVariablesFromContext),
+            postponedArguments,
-                postponedArguments,
+            completionMode,
-                completionMode,
+            topLevelType
-                topLevelType
+        ) ?: return false
-            ) ?: break
 
-            if (!variableForFixation.hasProperConstraint && completionMode == ConstraintSystemCompletionMode.PARTIAL)
+        val variableWithConstraints = notFixedTypeVariables.getValue(variableForFixation.variable)
-                break
+        if (!variableForFixation.hasProperConstraint) {
-
+            if (context.inferenceSession.isSyntheticTypeVariable(variableWithConstraints.typeVariable)) {
-            val variableWithConstraints = notFixedTypeVariables.getValue(variableForFixation.variable)
+                context.inferenceSession.fixSyntheticTypeVariableWithNotEnoughInformation(
-
+                    variableWithConstraints.typeVariable as ConeTypeVariable,
-            when {
+                    this
-                variableForFixation.hasProperConstraint -> {
+                )
-                    fixVariable(this, topLevelType, variableWithConstraints, postponedArguments)
+                return true
-                    return true
-                }
-                context.inferenceSession.isSyntheticTypeVariable(variableWithConstraints.typeVariable) -> {
-                    context.inferenceSession.fixSyntheticTypeVariableWithNotEnoughInformation(
-                        variableWithConstraints.typeVariable as ConeTypeVariable,
-                        this
-                    )
-                    return true
-                }
-                else -> {
-                    processVariableWhenNotEnoughInformation(this, variableWithConstraints, topLevelAtoms)
-                }
             }
+
+            return false
         }
 
-        return false
+        fixVariable(this, topLevelType, variableWithConstraints, postponedArguments)
+
+        return true
     }
 
-    private fun processVariableWhenNotEnoughInformation(
+    private fun ConstraintSystemCompletionContext.reportNotEnoughTypeInformation(
-        c: ConstraintSystemCompletionContext,
+        completionMode: ConstraintSystemCompletionMode,
+        topLevelAtoms: List<FirStatement>,
+        topLevelType: ConeKotlinType,
+        collectVariablesFromContext: Boolean,
+        postponedArguments: List<PostponedResolvedAtom>,
+    ) {
+        while (true) {
+            val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
+                this, getOrderedAllTypeVariables(collectVariablesFromContext, topLevelAtoms),
+                postponedArguments, completionMode, topLevelType,
+            ) ?: break
+
+            assert(!variableForFixation.hasProperConstraint) {
+                "At this stage there should be no remaining variables with proper constraints"
+            }
+
+            if (completionMode == ConstraintSystemCompletionMode.PARTIAL) break
+
+            val variableWithConstraints = notFixedTypeVariables.getValue(variableForFixation.variable)
+            processVariableWhenNotEnoughInformation(variableWithConstraints, topLevelAtoms)
+        }
+    }
+
+    private fun ConstraintSystemCompletionContext.processVariableWhenNotEnoughInformation(
         variableWithConstraints: VariableWithConstraints,
         topLevelAtoms: List<FirStatement>,
     ) {
         val typeVariable = variableWithConstraints.typeVariable
-        val resolvedAtom =
+        val resolvedAtom = findResolvedAtomBy(typeVariable, topLevelAtoms) ?: topLevelAtoms.firstOrNull()
-            findResolvedAtomBy(typeVariable, topLevelAtoms) ?: topLevelAtoms.firstOrNull()
 
         if (resolvedAtom != null) {
-            c.addError(
+            addError(
-                NotEnoughInformationForTypeParameter(typeVariable, resolvedAtom, c.couldBeResolvedWithUnrestrictedBuilderInference())
+                NotEnoughInformationForTypeParameter(typeVariable, resolvedAtom, couldBeResolvedWithUnrestrictedBuilderInference())
             )
         }
 
@@ -274,84 +310,19 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
             else -> createCannotInferErrorType("Cannot infer type variable $typeVariable")
         }
 
-        c.fixVariable(typeVariable, resultErrorType, ConeFixVariableConstraintPosition(typeVariable))
+        fixVariable(typeVariable, resultErrorType, ConeFixVariableConstraintPosition(typeVariable))
     }
 
-    private fun createCannotInferErrorType(message: String, isUninferredParameter: Boolean = false) =
+    private fun ConstraintSystemCompletionContext.getOrderedAllTypeVariables(
-        ConeClassErrorType(
+        collectVariablesFromContext: Boolean,
-            ConeSimpleDiagnostic(
-                message,
-                DiagnosticKind.CannotInferParameterType,
-            ),
-            isUninferredParameter,
-        )
-
-    private fun findResolvedAtomBy(
-        typeVariable: TypeVariableMarker,
         topLevelAtoms: List<FirStatement>
-    ): FirStatement? {
+    ): List<TypeConstructorMarker> {
-
-        fun FirStatement.findFirstAtomContainingVariable(): FirStatement? {
-
-            var result: FirStatement? = null
-
-            fun suggestElement(element: FirElement) {
-                if (result == null && element is FirStatement) {
-                    result = element
-                }
-            }
-
-            this@findFirstAtomContainingVariable.processAllContainingCallCandidates(processBlocks = true) { candidate ->
-                if (typeVariable in candidate.freshVariables) {
-                    suggestElement(candidate.callInfo.callSite)
-                }
-
-                for (postponedAtom in candidate.postponedAtoms) {
-                    if (postponedAtom is ResolvedLambdaAtom) {
-                        if (postponedAtom.typeVariableForLambdaReturnType == typeVariable) {
-                            suggestElement(postponedAtom.atom)
-                        }
-                    }
-                }
-            }
-
-            return result
-        }
-
-        return topLevelAtoms.firstNotNullOfOrNull(FirStatement::findFirstAtomContainingVariable)
-    }
-
-    private fun analyzeRemainingNotAnalyzedPostponedArgument(
-        postponedArguments: List<PostponedResolvedAtom>,
-        analyze: (PostponedResolvedAtom) -> Unit
-    ): Boolean {
-        val remainingNotAnalyzedPostponedArgument = postponedArguments.firstOrNull { !it.analyzed }
-
-        if (remainingNotAnalyzedPostponedArgument != null) {
-            analyze(remainingNotAnalyzedPostponedArgument)
-            return true
-        }
-
-        return false
-    }
-
-    private fun ConstraintSystemCompletionContext.hasLambdaToAnalyze(
-        postponedArguments: List<PostponedResolvedAtom>
-    ): Boolean {
-        return analyzeArgumentWithFixedParameterTypes(postponedArguments) {}
-    }
-
-    private fun getOrderedAllTypeVariables(
-        c: ConstraintSystemCompletionContext,
-        topLevelAtoms: List<FirStatement>,
-        collectVariablesFromContext: Boolean
-    ): List<TypeConstructorMarker> = with(c) {
         if (collectVariablesFromContext) {
-            return c.notFixedTypeVariables.keys.toList()
+            return notFixedTypeVariables.keys.toList()
         }
-        val result = LinkedHashSet<TypeConstructorMarker>(c.notFixedTypeVariables.size)
+        val result = LinkedHashSet<TypeConstructorMarker>(notFixedTypeVariables.size)
         fun ConeTypeVariable?.toTypeConstructor(): TypeConstructorMarker? =
-            this?.typeConstructor?.takeIf { it in c.notFixedTypeVariables.keys }
+            this?.typeConstructor?.takeIf { it in notFixedTypeVariables.keys }
 
         // TODO: non-top-level variables?
         fun PostponedAtomWithRevisableExpectedType.collectNotFixedVariables() {
@@ -372,33 +343,35 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
                 }
 
                 for (postponedAtom in candidate.postponedAtoms) {
-                    when {
+                    when (postponedAtom) {
-                        postponedAtom is ResolvedLambdaAtom -> {
+                        is ResolvedLambdaAtom -> {
                             result.addIfNotNull(postponedAtom.typeVariableForLambdaReturnType.toTypeConstructor())
                         }
-                        postponedAtom is LambdaWithTypeVariableAsExpectedTypeAtom -> {
+                        is LambdaWithTypeVariableAsExpectedTypeAtom -> {
                             postponedAtom.collectNotFixedVariables()
                         }
-                        postponedAtom is ResolvedCallableReferenceAtom -> {
+                        is ResolvedCallableReferenceAtom -> {
                             if (postponedAtom.mightNeedAdditionalResolution) {
                                 postponedAtom.collectNotFixedVariables()
                             }
                         }
+                        // ResolvedCallAtom?
+                        // ResolvedCallableReferenceArgumentAtom?
                     }
                 }
             }
         }
 
-        for (topLevel in topLevelAtoms) {
+        for (topLevelAtom in topLevelAtoms) {
-            topLevel.collectAllTypeVariables()
+            topLevelAtom.collectAllTypeVariables()
         }
 
         if (context.inferenceSession.hasSyntheticTypeVariables()) {
-            result.addAll(c.notFixedTypeVariables.filter { context.inferenceSession.isSyntheticTypeVariable(it.value.typeVariable) }.keys.asIterable())
+            result.addAll(notFixedTypeVariables.filter { context.inferenceSession.isSyntheticTypeVariable(it.value.typeVariable) }.keys.asIterable())
         }
 
-        require(result.size == c.notFixedTypeVariables.size) {
+        require(result.size == notFixedTypeVariables.size) {
-            val notFoundTypeVariables = c.notFixedTypeVariables.keys.toMutableSet().apply { removeAll(result) }
+            val notFoundTypeVariables = notFixedTypeVariables.keys.toMutableSet().apply { removeAll(result) }
             "Not all type variables found: $notFoundTypeVariables"
         }
 
@@ -417,21 +390,69 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents, p
         c.fixVariable(variable, resultType, ConeFixVariableConstraintPosition(variable)) // TODO: obtain atom for diagnostics
     }
 
-    private fun getOrderedNotAnalyzedPostponedArguments(topLevelAtoms: List<FirStatement>): List<PostponedResolvedAtom> {
+    companion object {
-        val notAnalyzedArguments = arrayListOf<PostponedResolvedAtom>()
+        private fun getOrderedNotAnalyzedPostponedArguments(topLevelAtoms: List<FirStatement>): List<PostponedResolvedAtom> {
-        for (primitive in topLevelAtoms) {
+            val notAnalyzedArguments = arrayListOf<PostponedResolvedAtom>()
-            primitive.processAllContainingCallCandidates(
+            for (primitive in topLevelAtoms) {
-                // TODO: remove this argument and relevant parameter
+                primitive.processAllContainingCallCandidates(
-                // Currently, it's used because otherwise problem happens with a lambda in a try-block (see tryWithLambdaInside test)
+                    // TODO: remove this argument and relevant parameter
-                processBlocks = true
+                    // Currently, it's used because otherwise problem happens with a lambda in a try-block (see tryWithLambdaInside test)
-            ) { candidate ->
+                    processBlocks = true
-                candidate.postponedAtoms.forEach {
+                ) { candidate ->
-                    notAnalyzedArguments.addIfNotNull(it.safeAs<PostponedResolvedAtom>()?.takeUnless { it.analyzed })
+                    candidate.postponedAtoms.forEach { atom ->
+                        notAnalyzedArguments.addIfNotNull(atom.safeAs<PostponedResolvedAtom>()?.takeUnless { it.analyzed })
+                    }
                 }
             }
+
+            return notAnalyzedArguments
         }
 
-        return notAnalyzedArguments
+        private fun findResolvedAtomBy(
+            typeVariable: TypeVariableMarker,
+            topLevelAtoms: List<FirStatement>
+        ): FirStatement? {
+
+            fun FirStatement.findFirstAtomContainingVariable(): FirStatement? {
+
+                var result: FirStatement? = null
+
+                fun suggestElement(element: FirElement) {
+                    if (result == null && element is FirStatement) {
+                        result = element
+                    }
+                }
+
+                this@findFirstAtomContainingVariable.processAllContainingCallCandidates(processBlocks = true) { candidate ->
+                    if (typeVariable in candidate.freshVariables) {
+                        suggestElement(candidate.callInfo.callSite)
+                    }
+
+                    for (postponedAtom in candidate.postponedAtoms) {
+                        if (postponedAtom is ResolvedLambdaAtom) {
+                            if (postponedAtom.typeVariableForLambdaReturnType == typeVariable) {
+                                suggestElement(postponedAtom.atom)
+                            }
+                        }
+                    }
+                }
+
+                return result
+            }
+
+            return topLevelAtoms.firstNotNullOfOrNull(FirStatement::findFirstAtomContainingVariable)
+        }
+
+        private fun createCannotInferErrorType(message: String, isUninferredParameter: Boolean = false) =
+            ConeClassErrorType(
+                ConeSimpleDiagnostic(
+                    message,
+                    DiagnosticKind.CannotInferParameterType,
+                ),
+                isUninferredParameter,
+            )
+
+
     }
 }
 

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

@@ -5,33 +5,109 @@
 
 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.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintSystemError
 import org.jetbrains.kotlin.resolve.calls.inference.model.FixVariableConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
+import org.jetbrains.kotlin.resolve.calls.model.PostponedAtomWithRevisableExpectedType
+import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
-interface ConstraintSystemCompletionContext : VariableFixationFinder.Context, ResultTypeResolver.Context {
+abstract class ConstraintSystemCompletionContext : VariableFixationFinder.Context, ResultTypeResolver.Context {
-    val allTypeVariables: Map<TypeConstructorMarker, TypeVariableMarker>
+    abstract val allTypeVariables: Map<TypeConstructorMarker, TypeVariableMarker>
-    override val notFixedTypeVariables: Map<TypeConstructorMarker, VariableWithConstraints>
+    abstract override val notFixedTypeVariables: Map<TypeConstructorMarker, VariableWithConstraints>
-    override val fixedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker>
+    abstract override val fixedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker>
-    override val postponedTypeVariables: List<TypeVariableMarker>
+    abstract override val postponedTypeVariables: List<TypeVariableMarker>
 
-    fun getBuilder(): ConstraintSystemBuilder
+    abstract fun getBuilder(): ConstraintSystemBuilder
 
     // type can be proper if it not contains not fixed type variables
-    fun canBeProper(type: KotlinTypeMarker): Boolean
+    abstract fun canBeProper(type: KotlinTypeMarker): Boolean
 
-    fun containsOnlyFixedOrPostponedVariables(type: KotlinTypeMarker): Boolean
+    abstract fun containsOnlyFixedOrPostponedVariables(type: KotlinTypeMarker): Boolean
-    fun containsOnlyFixedVariables(type: KotlinTypeMarker): Boolean
+    abstract fun containsOnlyFixedVariables(type: KotlinTypeMarker): Boolean
 
     // mutable operations
-    fun addError(error: ConstraintSystemError)
+    abstract fun addError(error: ConstraintSystemError)
 
-    fun fixVariable(variable: TypeVariableMarker, resultType: KotlinTypeMarker, position: FixVariableConstraintPosition<*>)
+    abstract fun fixVariable(variable: TypeVariableMarker, resultType: KotlinTypeMarker, position: FixVariableConstraintPosition<*>)
 
-    fun couldBeResolvedWithUnrestrictedBuilderInference(): Boolean
+    abstract fun couldBeResolvedWithUnrestrictedBuilderInference(): Boolean
-    fun processForkConstraints()
+    abstract fun processForkConstraints()
+
+    fun <A : PostponedResolvedAtomMarker> analyzeArgumentWithFixedParameterTypes(
+        languageVersionSettings: LanguageVersionSettings,
+        postponedArguments: List<A>,
+        analyze: (A) -> Unit
+    ): Boolean {
+        val useBuilderInferenceOnlyIfNeeded =
+            languageVersionSettings.supportsFeature(LanguageFeature.UseBuilderInferenceOnlyIfNeeded)
+        val argumentToAnalyze = if (useBuilderInferenceOnlyIfNeeded) {
+            findPostponedArgumentWithFixedInputTypes(postponedArguments)
+        } else {
+            findPostponedArgumentWithFixedOrPostponedInputTypes(postponedArguments)
+        }
+
+        if (argumentToAnalyze != null) {
+            analyze(argumentToAnalyze)
+            return true
+        }
+
+        return false
+    }
+
+    fun <A : PostponedResolvedAtomMarker> analyzeNextReadyPostponedArgument(
+        languageVersionSettings: LanguageVersionSettings,
+        postponedArguments: List<A>,
+        completionMode: ConstraintSystemCompletionMode,
+        analyze: (A) -> Unit
+    ): Boolean {
+        if (completionMode == ConstraintSystemCompletionMode.FULL) {
+            val argumentWithTypeVariableAsExpectedType = findPostponedArgumentWithRevisableExpectedType(postponedArguments)
+
+            if (argumentWithTypeVariableAsExpectedType != null) {
+                analyze(argumentWithTypeVariableAsExpectedType)
+                return true
+            }
+        }
+
+        return analyzeArgumentWithFixedParameterTypes(languageVersionSettings, postponedArguments, analyze)
+    }
+
+    fun <A : PostponedResolvedAtomMarker> analyzeRemainingNotAnalyzedPostponedArgument(
+        postponedArguments: List<A>,
+        analyze: (A) -> Unit
+    ): Boolean {
+        val remainingNotAnalyzedPostponedArgument = postponedArguments.firstOrNull { !it.analyzed }
+
+        if (remainingNotAnalyzedPostponedArgument != null) {
+            analyze(remainingNotAnalyzedPostponedArgument)
+            return true
+        }
+
+        return false
+    }
+
+    fun <A : PostponedResolvedAtomMarker> hasLambdaToAnalyze(
+        languageVersionSettings: LanguageVersionSettings,
+        postponedArguments: List<A>
+    ): Boolean {
+        return analyzeArgumentWithFixedParameterTypes(languageVersionSettings, postponedArguments) {}
+    }
+
+    // Avoiding smart cast from filterIsInstanceOrNull looks dirty
+    private fun <A : PostponedResolvedAtomMarker> findPostponedArgumentWithRevisableExpectedType(postponedArguments: List<A>): A? =
+        postponedArguments.firstOrNull { argument -> argument is PostponedAtomWithRevisableExpectedType }
+
+    private fun <T : PostponedResolvedAtomMarker> findPostponedArgumentWithFixedOrPostponedInputTypes(
+        postponedArguments: List<T>
+    ) = postponedArguments.firstOrNull { argument -> argument.inputTypes.all { containsOnlyFixedOrPostponedVariables(it) } }
+
+    private fun <T : PostponedResolvedAtomMarker> findPostponedArgumentWithFixedInputTypes(
+        postponedArguments: List<T>
+    ) = postponedArguments.firstOrNull { argument -> argument.inputTypes.all { containsOnlyFixedVariables(it) } }
 }

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

@@ -20,12 +20,12 @@ import kotlin.math.max
 class NewConstraintSystemImpl(
     private val constraintInjector: ConstraintInjector,
     val typeSystemContext: TypeSystemInferenceExtensionContext
-) : TypeSystemInferenceExtensionContext by typeSystemContext,
+) : ConstraintSystemCompletionContext(),
+    TypeSystemInferenceExtensionContext by typeSystemContext,
     NewConstraintSystem,
     ConstraintSystemBuilder,
     ConstraintInjector.Context,
     ResultTypeResolver.Context,
-    ConstraintSystemCompletionContext,
     PostponedArgumentsAnalyzerContext
 {
     private val utilContext = constraintInjector.constraintIncorporator.utilContext

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

@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
-import org.jetbrains.kotlin.types.model.TypeVariableTypeConstructorMarker
 import org.jetbrains.kotlin.types.model.safeSubstitute
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -24,7 +23,7 @@ import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 class KotlinConstraintSystemCompleter(
     private val resultTypeResolver: ResultTypeResolver,
     val variableFixationFinder: VariableFixationFinder,
-    private val postponedArgumentInputTypesResolver: PostponedArgumentInputTypesResolver,
+    private val postponedArgumentsInputTypesResolver: PostponedArgumentInputTypesResolver,
     private val languageVersionSettings: LanguageVersionSettings
 ) {
     fun runCompletion(
@@ -74,28 +73,38 @@ class KotlinConstraintSystemCompleter(
         val topLevelTypeVariables = topLevelType.extractTypeVariables()
 
         completion@ while (true) {
-            // TODO
+            // TODO: This is very slow
             val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
-            if (completionMode == ConstraintSystemCompletionMode.UNTIL_FIRST_LAMBDA && hasLambdaToAnalyze(postponedArguments)) return
+
+            if (completionMode == ConstraintSystemCompletionMode.UNTIL_FIRST_LAMBDA && hasLambdaToAnalyze(
+                    languageVersionSettings,
+                    postponedArguments
+                )
+            ) return
 
             // Stage 1: analyze postponed arguments with fixed parameter types
-            if (analyzeArgumentWithFixedParameterTypes(postponedArguments, analyze))
+            if (analyzeArgumentWithFixedParameterTypes(languageVersionSettings, postponedArguments, analyze))
                 continue
 
-            val isThereAnyReadyForFixationVariable = isThereAnyReadyForFixationVariable(
+            val isThereAnyReadyForFixationVariable = variableFixationFinder.findFirstVariableForFixation(
-                completionMode, topLevelAtoms, topLevelType, collectVariablesFromContext, postponedArguments
+                this,
-            )
+                getOrderedAllTypeVariables(collectVariablesFromContext, topLevelAtoms),
+                postponedArguments,
+                completionMode,
+                topLevelType
+            ) != null
 
             // If there aren't any postponed arguments and ready for fixation variables, then completion isn't needed: nothing to do
             if (postponedArguments.isEmpty() && !isThereAnyReadyForFixationVariable)
                 break
 
-            val postponedArgumentsWithRevisableType = postponedArguments.filterIsInstance<PostponedAtomWithRevisableExpectedType>()
+            val postponedArgumentsWithRevisableType = postponedArguments
+                .filterIsInstance<PostponedAtomWithRevisableExpectedType>()
             val dependencyProvider =
                 TypeVariableDependencyInformationProvider(notFixedTypeVariables, postponedArguments, topLevelType, this)
 
             // Stage 2: collect parameter types for postponed arguments
-            val wasBuiltNewExpectedTypeForSomeArgument = postponedArgumentInputTypesResolver.collectParameterTypesAndBuildNewExpectedTypes(
+            val wasBuiltNewExpectedTypeForSomeArgument = postponedArgumentsInputTypesResolver.collectParameterTypesAndBuildNewExpectedTypes(
                 this,
                 postponedArgumentsWithRevisableType,
                 completionMode,
@@ -109,7 +118,7 @@ class KotlinConstraintSystemCompleter(
             if (completionMode == ConstraintSystemCompletionMode.FULL) {
                 // Stage 3: fix variables for parameter types of all postponed arguments
                 for (argument in postponedArguments) {
-                    val wasFixedSomeVariable = postponedArgumentInputTypesResolver.fixNextReadyVariableForParameterTypeIfNeeded(
+                    val variableWasFixed = postponedArgumentsInputTypesResolver.fixNextReadyVariableForParameterTypeIfNeeded(
                         this,
                         argument,
                         postponedArguments,
@@ -119,23 +128,23 @@ class KotlinConstraintSystemCompleter(
                         findResolvedAtomBy(it, topLevelAtoms) ?: topLevelAtoms.firstOrNull()
                     }
 
-                    if (wasFixedSomeVariable)
+                    if (variableWasFixed)
                         continue@completion
                 }
 
                 // Stage 4: create atoms with revised expected types if needed
                 for (argument in postponedArgumentsWithRevisableType) {
-                    val wasTransformedSomeArgument = transformToAtomWithNewFunctionalExpectedType(
+                    val argumentWasTransformed = transformToAtomWithNewFunctionalExpectedType(
                         this, argument, diagnosticsHolder
                     )
 
-                    if (wasTransformedSomeArgument)
+                    if (argumentWasTransformed)
                         continue@completion
                 }
             }
 
             // Stage 5: analyze the next ready postponed argument
-            if (analyzeNextReadyPostponedArgument(postponedArguments, completionMode, analyze))
+            if (analyzeNextReadyPostponedArgument(languageVersionSettings, postponedArguments, completionMode, analyze))
                 continue
 
             // Stage 6: fix next ready type variable with proper constraints
@@ -165,6 +174,29 @@ class KotlinConstraintSystemCompleter(
         }
     }
 
+    fun prepareLambdaAtomForFactoryPattern(
+        atom: ResolvedLambdaAtom,
+        candidate: SimpleResolutionCandidate,
+        diagnosticsHolder: KotlinDiagnosticsHolder,
+    ): ResolvedLambdaAtom {
+        val returnVariable = TypeVariableForLambdaReturnType(candidate.callComponents.builtIns, "_R")
+        val csBuilder = candidate.getSystem().getBuilder()
+        csBuilder.registerVariable(returnVariable)
+        val functionalType: KotlinType = csBuilder.buildCurrentSubstitutor().safeSubstitute(candidate.getSystem().asConstraintSystemCompleterContext(), atom.expectedType!!) as KotlinType
+        val expectedType = KotlinTypeFactory.simpleType(
+            functionalType.attributes,
+            functionalType.constructor,
+            functionalType.arguments.dropLast(1) + returnVariable.defaultType.asTypeProjection(),
+            functionalType.isMarkedNullable
+        )
+        csBuilder.addSubtypeConstraint(
+            expectedType,
+            functionalType,
+            ArgumentConstraintPositionImpl(atom.atom)
+        )
+        return atom.transformToResolvedLambda(csBuilder, diagnosticsHolder, expectedType, returnVariable)
+    }
+
     private fun ConstraintSystemCompletionContext.tryToCompleteWithBuilderInference(
         completionMode: ConstraintSystemCompletionMode,
         topLevelAtoms: List<ResolvedAtom>,
@@ -231,115 +263,6 @@ class KotlinConstraintSystemCompleter(
         return true
     }
 
-    private fun ConstraintSystemCompletionContext.analyzeArgumentWithFixedParameterTypes(
-        postponedArguments: List<PostponedResolvedAtom>,
-        analyze: (PostponedResolvedAtom) -> Unit
-    ): Boolean {
-        val useBuilderInferenceOnlyIfNeeded = languageVersionSettings.supportsFeature(LanguageFeature.UseBuilderInferenceOnlyIfNeeded)
-        val argumentToAnalyze = if (useBuilderInferenceOnlyIfNeeded) {
-            findPostponedArgumentWithFixedInputTypes(postponedArguments)
-        } else {
-            findPostponedArgumentWithFixedOrPostponedInputTypes(postponedArguments)
-        }
-
-        if (argumentToAnalyze != null) {
-            analyze(argumentToAnalyze)
-            return true
-        }
-
-        return false
-    }
-
-    private fun ConstraintSystemCompletionContext.analyzeNextReadyPostponedArgument(
-        postponedArguments: List<PostponedResolvedAtom>,
-        completionMode: ConstraintSystemCompletionMode,
-        analyze: (PostponedResolvedAtom) -> Unit
-    ): Boolean {
-        if (completionMode == ConstraintSystemCompletionMode.FULL) {
-            val argumentWithTypeVariableAsExpectedType = findPostponedArgumentWithRevisableExpectedType(postponedArguments)
-
-            if (argumentWithTypeVariableAsExpectedType != null) {
-                analyze(argumentWithTypeVariableAsExpectedType)
-                return true
-            }
-        }
-
-        return analyzeArgumentWithFixedParameterTypes(postponedArguments, analyze)
-    }
-
-    private fun analyzeRemainingNotAnalyzedPostponedArgument(
-        postponedArguments: List<PostponedResolvedAtom>,
-        analyze: (PostponedResolvedAtom) -> Unit
-    ): Boolean {
-        val remainingNotAnalyzedPostponedArgument = postponedArguments.firstOrNull { !it.analyzed }
-
-        if (remainingNotAnalyzedPostponedArgument != null) {
-            analyze(remainingNotAnalyzedPostponedArgument)
-            return true
-        }
-
-        return false
-    }
-
-    fun prepareLambdaAtomForFactoryPattern(
-        atom: ResolvedLambdaAtom,
-        candidate: SimpleResolutionCandidate,
-        diagnosticsHolder: KotlinDiagnosticsHolder,
-    ): ResolvedLambdaAtom {
-        val returnVariable = TypeVariableForLambdaReturnType(candidate.callComponents.builtIns, "_R")
-        val csBuilder = candidate.getSystem().getBuilder()
-        csBuilder.registerVariable(returnVariable)
-        val functionalType: KotlinType = csBuilder.buildCurrentSubstitutor().safeSubstitute(candidate.getSystem().asConstraintSystemCompleterContext(), atom.expectedType!!) as KotlinType
-        val expectedType = KotlinTypeFactory.simpleType(
-            functionalType.attributes,
-            functionalType.constructor,
-            functionalType.arguments.dropLast(1) + returnVariable.defaultType.asTypeProjection(),
-            functionalType.isMarkedNullable
-        )
-        csBuilder.addSubtypeConstraint(
-            expectedType,
-            functionalType,
-            ArgumentConstraintPositionImpl(atom.atom)
-        )
-        return atom.transformToResolvedLambda(csBuilder, diagnosticsHolder, expectedType, returnVariable)
-    }
-
-    private fun ConstraintSystemCompletionContext.hasLambdaToAnalyze(
-        postponedArguments: List<PostponedResolvedAtom>
-    ): Boolean {
-        return analyzeArgumentWithFixedParameterTypes(postponedArguments) {}
-    }
-
-    private fun findPostponedArgumentWithRevisableExpectedType(postponedArguments: List<PostponedResolvedAtom>) =
-        postponedArguments.firstOrNull { argument -> argument is PostponedAtomWithRevisableExpectedType }
-
-    private fun ConstraintSystemCompletionContext.findPostponedArgumentWithFixedOrPostponedInputTypes(
-        postponedArguments: List<PostponedResolvedAtom>
-    ) = postponedArguments.firstOrNull { argument -> argument.inputTypes.all { containsOnlyFixedOrPostponedVariables(it) } }
-
-    private fun ConstraintSystemCompletionContext.findPostponedArgumentWithFixedInputTypes(postponedArguments: List<PostponedResolvedAtom>) =
-        postponedArguments.firstOrNull { argument -> argument.inputTypes.all { containsOnlyFixedVariables(it) } }
-
-    private fun fixVariable(
-        c: ConstraintSystemCompletionContext,
-        variableWithConstraints: VariableWithConstraints,
-        topLevelAtoms: List<ResolvedAtom>
-    ) {
-        fixVariable(c, variableWithConstraints, TypeVariableDirectionCalculator.ResolveDirection.UNKNOWN, topLevelAtoms)
-    }
-
-    private fun fixVariable(
-        c: ConstraintSystemCompletionContext,
-        variableWithConstraints: VariableWithConstraints,
-        direction: TypeVariableDirectionCalculator.ResolveDirection,
-        topLevelAtoms: List<ResolvedAtom>
-    ) {
-        val resultType = resultTypeResolver.findResultType(c, variableWithConstraints, direction)
-        val variable = variableWithConstraints.typeVariable
-        val resolvedAtom = findResolvedAtomBy(variable, topLevelAtoms) ?: topLevelAtoms.firstOrNull()
-        c.fixVariable(variable, resultType, FixVariableConstraintPositionImpl(variable, resolvedAtom))
-    }
-
     private fun ConstraintSystemCompletionContext.fixNextReadyVariable(
         completionMode: ConstraintSystemCompletionMode,
         topLevelAtoms: List<ResolvedAtom>,
@@ -348,8 +271,11 @@ class KotlinConstraintSystemCompleter(
         postponedArguments: List<PostponedResolvedAtom>,
     ): Boolean {
         val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
-            this, getOrderedAllTypeVariables(collectVariablesFromContext, topLevelAtoms),
+            this,
-            postponedArguments, completionMode, topLevelType,
+            getOrderedAllTypeVariables(collectVariablesFromContext, topLevelAtoms),
+            postponedArguments,
+            completionMode,
+            topLevelType
         ) ?: return false
 
         if (!variableForFixation.hasProperConstraint) return false
@@ -380,32 +306,41 @@ class KotlinConstraintSystemCompleter(
             if (completionMode == ConstraintSystemCompletionMode.PARTIAL) break
 
             val variableWithConstraints = notFixedTypeVariables.getValue(variableForFixation.variable)
-            val typeVariable = variableWithConstraints.typeVariable
+            processVariableWhenNotEnoughInformation(variableWithConstraints, topLevelAtoms, diagnosticsHolder)
-            val resolvedAtom = findResolvedAtomBy(typeVariable, topLevelAtoms) ?: topLevelAtoms.firstOrNull()
-
-            if (resolvedAtom != null) {
-                addError(
-                    NotEnoughInformationForTypeParameterImpl(typeVariable, resolvedAtom, couldBeResolvedWithUnrestrictedBuilderInference())
-                )
-            }
-
-            val resultErrorType = when {
-                typeVariable is TypeVariableFromCallableDescriptor -> {
-                    ErrorUtils.createUninferredParameterType(typeVariable.originalTypeParameter)
-                }
-                typeVariable is TypeVariableForLambdaParameterType && typeVariable.atom is LambdaKotlinCallArgument -> {
-                    diagnosticsHolder.addDiagnostic(
-                        NotEnoughInformationForLambdaParameter(typeVariable.atom, typeVariable.index)
-                    )
-                    ErrorUtils.createErrorType("Cannot infer lambda parameter type")
-                }
-                else -> ErrorUtils.createErrorType("Cannot infer type variable $typeVariable")
-            }
-
-            fixVariable(typeVariable, resultErrorType, FixVariableConstraintPositionImpl(typeVariable, resolvedAtom))
         }
     }
 
+    private fun ConstraintSystemCompletionContext.processVariableWhenNotEnoughInformation(
+        variableWithConstraints: VariableWithConstraints,
+        topLevelAtoms: List<ResolvedAtom>,
+        diagnosticsHolder: KotlinDiagnosticsHolder
+    ) {
+        val typeVariable = variableWithConstraints.typeVariable
+        val resolvedAtom = findResolvedAtomBy(typeVariable, topLevelAtoms) ?: topLevelAtoms.firstOrNull()
+
+        if (resolvedAtom != null) {
+            addError(
+                NotEnoughInformationForTypeParameterImpl(typeVariable, resolvedAtom, couldBeResolvedWithUnrestrictedBuilderInference())
+            )
+        }
+
+        val resultErrorType = when {
+            typeVariable is TypeVariableFromCallableDescriptor -> {
+                ErrorUtils.createUninferredParameterType(typeVariable.originalTypeParameter)
+            }
+            typeVariable is TypeVariableForLambdaParameterType && typeVariable.atom is LambdaKotlinCallArgument -> {
+                diagnosticsHolder.addDiagnostic(
+                    NotEnoughInformationForLambdaParameter(typeVariable.atom, typeVariable.index)
+                )
+                ErrorUtils.createErrorType("Cannot infer lambda parameter type")
+            }
+            else -> ErrorUtils.createErrorType("Cannot infer type variable $typeVariable")
+        }
+
+        fixVariable(typeVariable, resultErrorType, FixVariableConstraintPositionImpl(typeVariable, resolvedAtom))
+
+    }
+
     private fun ConstraintSystemCompletionContext.getOrderedAllTypeVariables(
         collectVariablesFromContext: Boolean,
         topLevelAtoms: List<ResolvedAtom>
@@ -416,19 +351,27 @@ class KotlinConstraintSystemCompleter(
         fun getVariablesFromRevisedExpectedType(revisedExpectedType: KotlinType?) =
             revisedExpectedType?.arguments?.map { it.type.constructor }?.filterIsInstance<TypeVariableTypeConstructor>()
 
-        fun ResolvedAtom.process(to: LinkedHashSet<TypeConstructor>) {
+        // Note that it's important to use Set here, because several atoms can share the same type variable
+        val result = linkedSetOf<TypeConstructor>()
+
+        fun ResolvedAtom.collectAllTypeVariables() {
             val typeVariables = when (this) {
-                is LambdaWithTypeVariableAsExpectedTypeAtom -> getVariablesFromRevisedExpectedType(revisedExpectedType).orEmpty()
+                is ResolvedLambdaAtom -> {
-                is ResolvedCallAtom -> freshVariablesSubstitutor.freshVariables.map { it.freshTypeConstructor }
+                    listOfNotNull(typeVariableForLambdaReturnType?.freshTypeConstructor)
-                is PostponedCallableReferenceAtom ->
+                }
+                is LambdaWithTypeVariableAsExpectedTypeAtom -> {
+                    getVariablesFromRevisedExpectedType(revisedExpectedType).orEmpty()
+                }
+                is PostponedCallableReferenceAtom -> {
                     getVariablesFromRevisedExpectedType(revisedExpectedType).orEmpty() +
                             candidate?.freshVariablesSubstitutor?.freshVariables?.map { it.freshTypeConstructor }.orEmpty()
+                }
+                is ResolvedCallAtom -> freshVariablesSubstitutor.freshVariables.map { it.freshTypeConstructor }
                 is ResolvedCallableReferenceArgumentAtom -> candidate?.freshVariablesSubstitutor?.freshVariables?.map { it.freshTypeConstructor }.orEmpty()
-                is ResolvedLambdaAtom -> listOfNotNull(typeVariableForLambdaReturnType?.freshTypeConstructor)
                 else -> emptyList()
             }
 
-            typeVariables.mapNotNullTo(to) {
+            typeVariables.mapNotNullTo(result) {
                 it.takeIf { notFixedTypeVariables.containsKey(it) }
             }
 
@@ -436,21 +379,19 @@ class KotlinConstraintSystemCompleter(
              * Hack for completing error candidates in delegate resolve
              */
             if (this is StubResolvedAtom && typeVariable in notFixedTypeVariables) {
-                to += typeVariable
+                result += typeVariable
             }
 
             if (analyzed) {
-                subResolvedAtoms?.forEach { it.process(to) }
+                subResolvedAtoms?.forEach { it.collectAllTypeVariables() }
             }
         }
 
-        // Note that it's important to use Set here, because several atoms can share the same type variable
+        for (topLevelAtom in topLevelAtoms) {
-        val result = linkedSetOf<TypeConstructor>()
+            topLevelAtom.collectAllTypeVariables()
-        for (primitive in topLevelAtoms) {
-            primitive.process(result)
         }
 
-        assert(result.size == notFixedTypeVariables.size) {
+        require(result.size == notFixedTypeVariables.size) {
             val notFoundTypeVariables = notFixedTypeVariables.keys.toMutableSet().apply { removeAll(result) }
             "Not all type variables found: $notFoundTypeVariables"
         }
@@ -458,19 +399,25 @@ class KotlinConstraintSystemCompleter(
         return result.toList()
     }
 
-    private fun ConstraintSystemCompletionContext.isThereAnyReadyForFixationVariable(
+    private fun fixVariable(
-        completionMode: ConstraintSystemCompletionMode,
+        c: ConstraintSystemCompletionContext,
-        topLevelAtoms: List<ResolvedAtom>,
+        variableWithConstraints: VariableWithConstraints,
-        topLevelType: UnwrappedType,
+        topLevelAtoms: List<ResolvedAtom>
-        collectVariablesFromContext: Boolean,
+    ) {
-        postponedArguments: List<PostponedResolvedAtom>
+        fixVariable(c, variableWithConstraints, TypeVariableDirectionCalculator.ResolveDirection.UNKNOWN, topLevelAtoms)
-    ) = variableFixationFinder.findFirstVariableForFixation(
+    }
-        this,
+
-        getOrderedAllTypeVariables(collectVariablesFromContext, topLevelAtoms),
+    private fun fixVariable(
-        postponedArguments,
+        c: ConstraintSystemCompletionContext,
-        completionMode,
+        variableWithConstraints: VariableWithConstraints,
-        topLevelType,
+        direction: TypeVariableDirectionCalculator.ResolveDirection,
-    ) != null
+        topLevelAtoms: List<ResolvedAtom>
+    ) {
+        val resultType = resultTypeResolver.findResultType(c, variableWithConstraints, direction)
+        val variable = variableWithConstraints.typeVariable
+        val resolvedAtom = findResolvedAtomBy(variable, topLevelAtoms) ?: topLevelAtoms.firstOrNull()
+        c.fixVariable(variable, resultType, FixVariableConstraintPositionImpl(variable, resolvedAtom))
+    }
 
     companion object {
         fun getOrderedNotAnalyzedPostponedArguments(topLevelAtoms: List<ResolvedAtom>): List<PostponedResolvedAtom> {

compiler/testData/codegen/box/inference/builderInference/kt42139.kt

@@ -1,7 +1,5 @@
 // !LANGUAGE: +UnrestrictedBuilderInference
 // WITH_STDLIB
-// IGNORE_BACKEND_FIR: JVM_IR
-// FIR status: NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER on select call (K)
 
 fun <R> select(vararg x: R) = x[0]
 fun <K> myEmptyList(): List<K> = emptyList()

compiler/testData/codegen/box/inference/builderInference/withoutAnnotation.kt

@@ -1,6 +1,4 @@
 // WITH_STDLIB
-// IGNORE_BACKEND_FIR: JVM_IR
-// FIR status: NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER on buildMap call (K)
 // !LANGUAGE: +UseBuilderInferenceWithoutAnnotation
 
 fun <K, V> buildMap(builderAction: MutableMap<K, V>.() -> Unit): Map<K, V> = mapOf()

compiler/testData/diagnostics/tests/callableReference/resolve/ambiguityWithBoundExtensionReceiver.fir.kt

@@ -10,6 +10,6 @@ fun A.foo() = ""
 
 class A {
     fun main() {
-        <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>bar<!>(::<!OVERLOAD_RESOLUTION_AMBIGUITY!>foo<!>) <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>checkType<!> { _<String>() }
+        <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>bar<!>(::<!OVERLOAD_RESOLUTION_AMBIGUITY!>foo<!>) checkType { _<String>() }
     }
 }

compiler/testData/diagnostics/tests/inference/builderInference/changingResolveIfDontUseBuilderInference.fir.kt

@@ -20,10 +20,10 @@ fun <K> id(x: K): K = x
 fun main() {
     val x: Map<in String, String> = buildMap {
         put("", "")
-        swap(foo())
+        swap(<!ARGUMENT_TYPE_MISMATCH!>foo()<!>)
     } // `Map<CharSequence, String>` if we use builder inference, `Map<String, String>` if we don't
 
     val y: MutableMap<String, CharSequence> = build7 {
-        <!ARGUMENT_TYPE_MISMATCH!>id(run { this })<!>
+        <!ARGUMENT_TYPE_MISMATCH, TYPE_MISMATCH, TYPE_MISMATCH!>id(run { this })<!>
     }
-}
+}

compiler/testData/diagnostics/tests/inference/builderInference/kt49828.fir.kt

@@ -204,7 +204,7 @@ interface Foo2<K, V> {
 fun <L, K, V> twoBuilderLambdas(@BuilderInference block: Foo<L>.() -> Unit, @BuilderInference block2: Foo2<K, V>.() -> Unit) {}
 
 fun test() {
-    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER, NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>twoBuilderLambdas<!>(
+    twoBuilderLambdas(
         {
             add("")
             with (get()) {

compiler/testData/diagnostics/tests/inference/builderInference/simpleLambdaInCallWithAnotherLambdaWithBuilderInference.fir.kt

@@ -1,20 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_PARAMETER -DEPRECATION -OPT_IN_IS_NOT_ENABLED -UNUSED_VARIABLE
-// WITH_STDLIB
-
-import kotlin.experimental.ExperimentalTypeInference
-
-@OptIn(ExperimentalTypeInference::class)
-fun <R> combined(
-    check: () -> Unit,
-    @BuilderInference block: TestInterface<R>.() -> Unit
-): R = TODO()
-
-interface TestInterface<R> {
-    fun emit(r: R)
-}
-
-fun test() {
-    val ret = <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>combined<!>({ }) {
-        emit(1)
-    }
-}

compiler/testData/diagnostics/tests/inference/builderInference/simpleLambdaInCallWithAnotherLambdaWithBuilderInference.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER -DEPRECATION -OPT_IN_IS_NOT_ENABLED -UNUSED_VARIABLE
 // WITH_STDLIB
 

compiler/testData/diagnostics/tests/inference/builderInference/withoutAnnotationDisabledFeature.fir.kt

@@ -5,8 +5,8 @@
 fun <K, V> buildMap(builderAction: MutableMap<K, V>.() -> Unit): Map<K, V> = mapOf()
 
 fun box(): String {
-    val x = <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER, NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>buildMap<!> {
+    val x = buildMap {
         put("", "")
     }
     return "OK"
-}
+}

compiler/testData/diagnostics/tests/inference/kt6175.fir.kt

@@ -26,10 +26,10 @@ fun test2() {
 fun <T, R> baz(body: (List<R>) -> T): T = fail()
 
 fun test3() {
-    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER, NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>baz<!> {
+    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>baz<!> {
         true
     }
-    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER, NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>baz<!> { x ->
+    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>baz<!> { x ->
         true
     }
 }
@@ -37,10 +37,10 @@ fun test3() {
 fun <T, R : Any> brr(body: (List<R?>) -> T): T = fail()
 
 fun test4() {
-    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER, NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>brr<!> {
+    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>brr<!> {
         true
     }
-    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER, NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>brr<!> { x ->
+    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>brr<!> { x ->
         true
     }
 }

compiler/testData/diagnostics/testsWithStdLib/builderInference/inferCoroutineTypeInOldVersion.fir.kt

@@ -0,0 +1,29 @@
+// !LANGUAGE: +UnrestrictedBuilderInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+class Builder<T> {
+    suspend fun add(t: T) {}
+}
+
+fun <S> build(g: suspend Builder<S>.() -> Unit): List<S> = TODO()
+fun <S> wrongBuild(g: Builder<S>.() -> Unit): List<S> = TODO()
+
+fun <S> Builder<S>.extensionAdd(s: S) {}
+
+suspend fun <S> Builder<S>.safeExtensionAdd(s: S) {}
+
+val member = build {
+    add(42)
+}
+
+val memberWithoutAnn = wrongBuild {
+    <!ILLEGAL_SUSPEND_FUNCTION_CALL!>add<!>(42)
+}
+
+val extension = build {
+    extensionAdd("foo")
+}
+
+val safeExtension = build {
+    safeExtensionAdd("foo")
+}

compiler/testData/diagnostics/testsWithStdLib/builderInference/inferCoroutineTypeInOldVersion.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !LANGUAGE: +UnrestrictedBuilderInference
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 

compiler/testData/diagnostics/testsWithStdLib/builderInference/resolveUsualCallWithBuilderInference.fir.kt

@@ -0,0 +1,38 @@
+// ALLOW_KOTLIN_PACKAGE
+// !LANGUAGE: +UnrestrictedBuilderInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// FILE: annotation.kt
+
+package kotlin
+
+annotation class BuilderInference
+
+// FILE: test.kt
+
+class Builder<T> {
+    fun add(t: T) {}
+}
+
+fun <S> build(@BuilderInference g: Builder<S>.() -> Unit): List<S> = TODO()
+fun <S> wrongBuild(g: Builder<S>.() -> Unit): List<S> = TODO()
+
+fun <S> Builder<S>.extensionAdd(s: S) {}
+
+@BuilderInference
+fun <S> Builder<S>.safeExtensionAdd(s: S) {}
+
+val member = build {
+    add(42)
+}
+
+val memberWithoutAnn = wrongBuild {
+    add(42)
+}
+
+val extension = build {
+    extensionAdd("foo")
+}
+
+val safeExtension = build {
+    safeExtensionAdd("foo")
+}

compiler/testData/diagnostics/testsWithStdLib/builderInference/resolveUsualCallWithBuilderInference.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // ALLOW_KOTLIN_PACKAGE
 // !LANGUAGE: +UnrestrictedBuilderInference
 // !DIAGNOSTICS: -UNUSED_PARAMETER

compiler/testData/diagnostics/testsWithStdLib/builderInference/resolveUsualCallWithBuilderInferenceWithRestrictions.fir.kt

@@ -0,0 +1,38 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// ALLOW_KOTLIN_PACKAGE
+// !WITH_NEW_INFERENCE
+// FILE: annotation.kt
+
+package kotlin
+
+annotation class BuilderInference
+
+// FILE: test.kt
+
+class Builder<T> {
+    fun add(t: T) {}
+}
+
+fun <S> build(@BuilderInference g: Builder<S>.() -> Unit): List<S> = TODO()
+fun <S> wrongBuild(g: Builder<S>.() -> Unit): List<S> = TODO()
+
+fun <S> Builder<S>.extensionAdd(s: S) {}
+
+@BuilderInference
+fun <S> Builder<S>.safeExtensionAdd(s: S) {}
+
+val member = build {
+    add(42)
+}
+
+val memberWithoutAnn = wrongBuild {
+    add(42)
+}
+
+val extension = build {
+    extensionAdd("foo")
+}
+
+val safeExtension = build {
+    safeExtensionAdd("foo")
+}

compiler/testData/diagnostics/testsWithStdLib/builderInference/resolveUsualCallWithBuilderInferenceWithRestrictions.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // ALLOW_KOTLIN_PACKAGE
 // !WITH_NEW_INFERENCE

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/applyInsideCoroutine.fir.kt

@@ -7,25 +7,25 @@ class Controller<T> {
 
 fun <S> generate(g: suspend Controller<S>.() -> Unit): S = TODO()
 
-val test1 = <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
+val test1 = generate {
-    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>apply<!> {
+    apply {
         yield(4)
     }
 }
 
-val test2 = <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
+val test2 = generate {
     yield(B)
-    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>apply<!> {
+    apply {
         yield(C)
     }
 }
 
-val test3 = <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
+val test3 = generate {
-    this.<!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER, NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>let<!> {
+    this.let {
         yield(B)
     }
 
-    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>apply<!> {
+    apply {
         yield(C)
     }
 }

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/extensionsWithNonValuableConstraints_1_2.fir.kt

@@ -1,17 +0,0 @@
-// !LANGUAGE: -ExperimentalBuilderInference
-// !DIAGNOSTICS: -UNUSED_PARAMETER
-
-interface Base
-
-interface Controller<T> : Base {
-    suspend fun yield(t: T) {}
-}
-
-fun <S> generate(g: suspend Controller<S>.() -> Unit): S = TODO()
-
-suspend fun Base.baseExtension() {}
-
-val test1 = <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
-    yield("foo")
-    baseExtension()
-}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/extensionsWithNonValuableConstraints_1_2.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: -ExperimentalBuilderInference
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/recursiveGenerators2.fir.kt

@@ -0,0 +1,23 @@
+// !OPT_IN: kotlin.RequiresOptIn
+// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
+// NI_EXPECTED_FILE
+
+@file:OptIn(ExperimentalTypeInference::class)
+
+import kotlin.experimental.ExperimentalTypeInference
+
+class GenericController<T> {
+    suspend fun yield(t: T) {}
+}
+
+fun <S> generate(@BuilderInference g: suspend GenericController<S>.() -> Unit): List<S> = TODO()
+
+@BuilderInference
+suspend fun <S> GenericController<List<S>>.yieldGenerate(g: suspend GenericController<S>.() -> Unit): Unit = TODO()
+
+val test1 = generate {
+    // TODO: KT-15185
+    yieldGenerate {
+        yield(4)
+    }
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/recursiveGenerators2.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !OPT_IN: kotlin.RequiresOptIn
 // !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
 // NI_EXPECTED_FILE

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/suspendCallsWithErrors.fir.kt

@@ -18,6 +18,6 @@ val test1 = generate {
     yield(<!NO_COMPANION_OBJECT!>A<!>)
 }
 
-val test2: Int = <!INITIALIZER_TYPE_MISMATCH, NEW_INFERENCE_ERROR!>generate {
+val test2: Int = generate {
-    yield(A())
+    yield(<!ARGUMENT_TYPE_MISMATCH!>A()<!>)
-}<!>
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/suspendCallsWrongUpperBound.fir.kt

@@ -0,0 +1,12 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// NI_EXPECTED_FILE
+
+class Controller<T : Number> {
+    suspend fun yield(t: T) {}
+}
+
+fun <S : Number> generate(g: suspend Controller<S>.() -> Unit): S = TODO()
+
+val test = <!NEW_INFERENCE_ERROR!>generate {
+    yield("foo")
+}<!>

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/suspendCallsWrongUpperBound.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // NI_EXPECTED_FILE
 

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/withUninferredParameter.fir.kt

@@ -0,0 +1,20 @@
+// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_ANONYMOUS_PARAMETER -UNUSED_VARIABLE
+// NI_EXPECTED_FILE
+
+class GenericController<T> {
+    suspend fun yield(t: T) {}
+}
+
+fun <S> generate(g: suspend GenericController<S>.(S) -> Unit): S = TODO()
+
+val test1 = generate {
+    yield(4)
+}
+
+val test2 = generate<Int> {
+    yield(4)
+}
+
+val test3 = generate { bar: Int ->
+    yield(4)
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/withUninferredParameter.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_ANONYMOUS_PARAMETER -UNUSED_VARIABLE
 // NI_EXPECTED_FILE
 

compiler/testData/ir/irText/types/castsInsideCoroutineInference.fir.ir.txt

@@ -94,8 +94,8 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
               BLOCK_BODY
                 VAR name:channel type:<root>.ChannelCoroutine<kotlin.Any> [val]
                   TYPE_OP type=<root>.ChannelCoroutine<kotlin.Any> origin=CAST typeOperand=<root>.ChannelCoroutine<kotlin.Any>
-                    CALL 'public abstract fun <get-channel> (): <root>.SendChannel<E of <root>.ProducerScope> declared in <root>.ProducerScope' type=<root>.SendChannel<kotlin.Nothing> origin=GET_PROPERTY
+                    CALL 'public abstract fun <get-channel> (): <root>.SendChannel<E of <root>.ProducerScope> declared in <root>.ProducerScope' type=<root>.SendChannel<kotlin.Any> origin=GET_PROPERTY
-                      $this: GET_VAR '$this$produce: <root>.ProducerScope<kotlin.Any> declared in <root>.asFairChannel.<anonymous>' type=<root>.ProducerScope<kotlin.Nothing> origin=null
+                      $this: GET_VAR '$this$produce: <root>.ProducerScope<kotlin.Any> declared in <root>.asFairChannel.<anonymous>' type=<root>.ProducerScope<kotlin.Any> origin=null
                 CALL 'public final fun collect <T> (action: kotlin.coroutines.SuspendFunction1<@[ParameterName(name = 'value')] T of <root>.collect, kotlin.Unit>): kotlin.Unit [inline,suspend] declared in <root>' type=kotlin.Unit origin=null
                   <T>: kotlin.Any?
                   $receiver: GET_VAR 'flow: <root>.Flow<*> declared in <root>.asFairChannel' type=<root>.Flow<*> origin=null
@@ -123,12 +123,12 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
     VALUE_PARAMETER name:flow index:0 type:<root>.Flow<*>
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='private final fun asChannel (flow: <root>.Flow<*>): <root>.ReceiveChannel<kotlin.Any> declared in <root>'
-        CALL 'public final fun produce <E> (block: @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction1<<root>.ProducerScope<E of <root>.produce>, kotlin.Unit>): <root>.ReceiveChannel<E of <root>.produce> declared in <root>' type=<root>.ReceiveChannel<@[ParameterName(name = 'value')] kotlin.Any> origin=null
+        CALL 'public final fun produce <E> (block: @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction1<<root>.ProducerScope<E of <root>.produce>, kotlin.Unit>): <root>.ReceiveChannel<E of <root>.produce> declared in <root>' type=<root>.ReceiveChannel<kotlin.Any> origin=null
-          <E>: @[ParameterName(name = 'value')] kotlin.Any
+          <E>: kotlin.Any
           $receiver: GET_VAR '<this>: <root>.CoroutineScope declared in <root>.asChannel' type=<root>.CoroutineScope origin=null
-          block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.coroutines.SuspendFunction1<<root>.ProducerScope<@[ParameterName(name = 'value')] kotlin.Any>, kotlin.Unit> origin=LAMBDA
+          block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.coroutines.SuspendFunction1<<root>.ProducerScope<kotlin.Any>, kotlin.Unit> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.ProducerScope<@[ParameterName(name = 'value')] kotlin.Any>) returnType:kotlin.Unit [suspend]
+            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.ProducerScope<kotlin.Any>) returnType:kotlin.Unit [suspend]
-              $receiver: VALUE_PARAMETER name:$this$produce type:<root>.ProducerScope<@[ParameterName(name = 'value')] kotlin.Any>
+              $receiver: VALUE_PARAMETER name:$this$produce type:<root>.ProducerScope<kotlin.Any>
               BLOCK_BODY
                 CALL 'public final fun collect <T> (action: kotlin.coroutines.SuspendFunction1<@[ParameterName(name = 'value')] T of <root>.collect, kotlin.Unit>): kotlin.Unit [inline,suspend] declared in <root>' type=kotlin.Unit origin=null
                   <T>: kotlin.Any?
@@ -139,8 +139,8 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
                       BLOCK_BODY
                         RETURN type=kotlin.Nothing from='local final fun <anonymous> (value: @[ParameterName(name = 'value')] kotlin.Any?): kotlin.Unit [suspend] declared in <root>.asChannel.<anonymous>'
                           CALL 'public abstract fun send (e: E of <root>.SendChannel): kotlin.Unit [suspend] declared in <root>.SendChannel' type=kotlin.Unit origin=null
-                            $this: CALL 'public abstract fun <get-channel> (): <root>.SendChannel<E of <root>.ProducerScope> declared in <root>.ProducerScope' type=<root>.SendChannel<@[ParameterName(name = 'value')] kotlin.Any> origin=GET_PROPERTY
+                            $this: CALL 'public abstract fun <get-channel> (): <root>.SendChannel<E of <root>.ProducerScope> declared in <root>.ProducerScope' type=<root>.SendChannel<kotlin.Any> origin=GET_PROPERTY
-                              $this: GET_VAR '$this$produce: <root>.ProducerScope<@[ParameterName(name = 'value')] kotlin.Any> declared in <root>.asChannel.<anonymous>' type=<root>.ProducerScope<@[ParameterName(name = 'value')] kotlin.Any> origin=null
+                              $this: GET_VAR '$this$produce: <root>.ProducerScope<kotlin.Any> declared in <root>.asChannel.<anonymous>' type=<root>.ProducerScope<kotlin.Any> origin=null
                             e: BLOCK type=@[ParameterName(name = 'value')] kotlin.Any origin=ELVIS
                               VAR IR_TEMPORARY_VARIABLE name:tmp_1 type:@[ParameterName(name = 'value')] kotlin.Any? [val]
                                 GET_VAR 'value: @[ParameterName(name = 'value')] kotlin.Any? declared in <root>.asChannel.<anonymous>.<anonymous>' type=@[ParameterName(name = 'value')] kotlin.Any? origin=null

compiler/testData/ir/irText/types/castsInsideCoroutineInference.fir.kt.txt

@@ -52,7 +52,7 @@ private fun CoroutineScope.asFairChannel(flow: Flow<*>): ReceiveChannel<Any> {
 }
 
 private fun CoroutineScope.asChannel(flow: Flow<*>): ReceiveChannel<Any> {
-  return <this>.produce<@ParameterName(name = "value") Any>(block = local suspend fun ProducerScope<@ParameterName(name = "value") Any>.<anonymous>() {
+  return <this>.produce<Any>(block = local suspend fun ProducerScope<Any>.<anonymous>() {
     flow.collect<Any?>(action = local suspend fun <anonymous>(value: @ParameterName(name = "value") Any?) {
       return $this$produce.<get-channel>().send(e = { // BLOCK
         val <elvis>: @ParameterName(name = "value") Any? = value