[FIR] Support builder (coroutine) inference

#KT-37431 Fixed

compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/builderInference.kt

@@ -0,0 +1,10 @@
+fun <T> foo(@BuilderInference block: MutableList<T>.() -> Unit): T = null!!
+
+fun takeString(s: String) {}
+
+fun test() {
+    val s = foo {
+        this.add("")
+    }
+    takeString(s)
+}

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

@@ -0,0 +1,13 @@
+FILE: builderInference.kt
+    public final fun <T> foo(@R|kotlin/BuilderInference|() block: R|kotlin/collections/MutableList<T>.() -> kotlin/Unit|): R|T| {
+        ^foo Null(null)!!
+    }
+    public final fun takeString(s: R|kotlin/String|): R|kotlin/Unit| {
+    }
+    public final fun test(): R|kotlin/Unit| {
+        lval s: R|kotlin/String| = R|/foo|<R|kotlin/String|>(<L> = foo@fun R|kotlin/collections/MutableList<kotlin/String>|.<anonymous>(): R|kotlin/Unit| {
+            ^ this@R|special/anonymous|.R|FakeOverride<kotlin/collections/MutableList.add: R|kotlin/Boolean|>|(String())
+        }
+        )
+        R|/takeString|(R|<local>/s|)
+    }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithStdlibTestGenerated.java

@@ -687,6 +687,11 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
             KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolveWithStdlib/inference"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
         }
 
+        @TestMetadata("builderInference.kt")
+        public void testBuilderInference() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/builderInference.kt");
+        }
+
         @TestMetadata("complexConstraintSystem.kt")
         public void testComplexConstraintSystem() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/complexConstraintSystem.kt");

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeTypeUtils.kt

@@ -6,10 +6,31 @@
 package org.jetbrains.kotlin.fir.types
 
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.utils.SmartSet
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 val ConeKotlinType.isNullable: Boolean get() = nullability != ConeNullability.NOT_NULL
 
 val ConeKotlinType.isMarkedNullable: Boolean get() = nullability == ConeNullability.NULLABLE
 
-val ConeKotlinType.classId: ClassId? get() = this.safeAs<ConeClassLikeType>()?.lookupTag?.classId
+val ConeKotlinType.classId: ClassId? get() = this.safeAs<ConeClassLikeType>()?.lookupTag?.classId
+
+fun ConeKotlinType.contains(predicate: (ConeKotlinType) -> Boolean): Boolean {
+    return contains(predicate, null)
+}
+
+private fun ConeKotlinType.contains(predicate: (ConeKotlinType) -> Boolean, visited: SmartSet<ConeKotlinType>?): Boolean {
+    if (visited?.contains(this) == true) return false
+    if (predicate(this)) return true
+
+    @Suppress("NAME_SHADOWING")
+    val visited = visited ?: SmartSet.create()
+    visited += this
+
+    return when (this) {
+        is ConeFlexibleType -> lowerBound.contains(predicate, visited) || upperBound.contains(predicate, visited)
+        is ConeDefinitelyNotNullType -> original.contains(predicate, visited)
+        is ConeIntersectionType -> intersectedTypes.any { it.contains(predicate, visited) }
+        else -> typeArguments.any { it is ConeKotlinTypeProjection && it.type.contains(predicate, visited) }
+    }
+}

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

@@ -14,7 +14,8 @@ enum class CallKind(vararg resolutionSequence: ResolutionStage) {
         CreateFreshTypeVariableSubstitutorStage,
         CheckReceivers.Dispatch,
         CheckReceivers.Extension,
-        CheckLowPriorityInOverloadResolution
+        CheckLowPriorityInOverloadResolution,
+        PostponedVariablesInitializerResolutionStage
     ),
     SyntheticSelect(
         MapArguments,
@@ -34,7 +35,8 @@ enum class CallKind(vararg resolutionSequence: ResolutionStage) {
         CheckReceivers.Extension,
         CheckArguments,
         EagerResolveOfCallableReferences,
-        CheckLowPriorityInOverloadResolution
+        CheckLowPriorityInOverloadResolution,
+        PostponedVariablesInitializerResolutionStage
     ),
     DelegatingConstructorCall(
         CheckVisibility,

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

@@ -15,8 +15,7 @@ import org.jetbrains.kotlin.fir.inferenceContext
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.inference.ResolvedCallableReferenceAtom
+import org.jetbrains.kotlin.fir.resolve.inference.*
-import org.jetbrains.kotlin.fir.resolve.inference.csBuilder
 import org.jetbrains.kotlin.fir.resolve.inference.extractInputOutputTypesFromCallableReferenceExpectedType
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.SyntheticSymbol
@@ -534,3 +533,35 @@ internal object CheckLowPriorityInOverloadResolution : CheckerStage() {
         }
     }
 }
+
+internal object PostponedVariablesInitializerResolutionStage : ResolutionStage() {
+    val BUILDER_INFERENCE_CLASS_ID: ClassId = ClassId.fromString("kotlin/BuilderInference")
+
+    override suspend fun check(candidate: Candidate, sink: CheckerSink, callInfo: CallInfo) {
+        val argumentMapping = candidate.argumentMapping ?: return
+        // TODO: convert type argument mapping to map [FirTypeParameterSymbol, FirTypedProjection?]
+        if (candidate.typeArgumentMapping is TypeArgumentMapping.Mapped) return
+        for (parameter in argumentMapping.values) {
+            if (!parameter.hasBuilderInferenceMarker()) continue
+            val type = parameter.returnTypeRef.coneTypeSafe<ConeKotlinType>() ?: continue
+            val receiverType = type.receiverType(callInfo.session) ?: continue
+
+            for (freshVariable in candidate.freshVariables) {
+                candidate.typeArgumentMapping
+                if (candidate.csBuilder.isPostponedTypeVariable(freshVariable)) continue
+                if (freshVariable !is TypeParameterBasedTypeVariable) continue
+                val typeParameterSymbol = freshVariable.typeParameterSymbol
+                val typeHasVariable = receiverType.contains {
+                    (it as? ConeTypeParameterType)?.lookupTag?.typeParameterSymbol == typeParameterSymbol
+                }
+                if (typeHasVariable) {
+                    candidate.csBuilder.markPostponedVariable(freshVariable)
+                }
+            }
+        }
+    }
+
+    private fun FirValueParameter.hasBuilderInferenceMarker(): Boolean {
+        return this.hasAnnotation(BUILDER_INFERENCE_CLASS_ID)
+    }
+}

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

@@ -5,12 +5,10 @@
 
 package org.jetbrains.kotlin.fir.resolve.inference
 
-import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
-import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.candidate
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -19,50 +17,32 @@ import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
-import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 abstract class AbstractManyCandidatesInferenceSession(
     protected val components: BodyResolveComponents,
-    initialCall: FirExpression,
     private val postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
 ) : FirInferenceSession() {
     private val errorCalls: MutableList<FirResolvable> = mutableListOf()
-    private val partiallyResolvedCalls: MutableList<FirResolvable> = mutableListOf()
+    protected val partiallyResolvedCalls: MutableList<Pair<FirResolvable, Candidate>> = mutableListOf()
     private val completedCalls: MutableSet<FirResolvable> = mutableSetOf()
 
-    init {
-        val initialCandidate = (initialCall as? FirResolvable)
-            ?.calleeReference
-            ?.safeAs<FirNamedReferenceWithCandidate>()
-            ?.candidate
-        if (initialCandidate != null) {
-            partiallyResolvedCalls += initialCall as FirResolvable
-        }
-    }
-
     private val unitType: ConeKotlinType = components.session.builtinTypes.unitType.coneTypeUnsafe()
 
     override val currentConstraintSystem: ConstraintStorage
         get() = partiallyResolvedCalls.lastOrNull()
-            ?.calleeReference
+            ?.second
-            ?.safeAs<FirNamedReferenceWithCandidate>()
-            ?.candidate
             ?.system
             ?.currentStorage()
             ?: ConstraintStorage.Empty
 
     private lateinit var resultingConstraintSystem: NewConstraintSystem
 
-    override fun shouldRunCompletion(candidate: Candidate): Boolean {
+    override fun <T> addCompetedCall(call: T, candidate: Candidate) where T : FirResolvable, T : FirStatement {
-        return false
-    }
-
-    override fun <T> addCompetedCall(call: T) where T : FirResolvable, T : FirStatement {
         // do nothing
     }
 
     final override fun <T> addPartiallyResolvedCall(call: T) where T : FirResolvable, T : FirStatement {
-        partiallyResolvedCalls += call
+        partiallyResolvedCalls += call to call.candidate
     }
 
     final override fun <T> addErrorCall(call: T) where T : FirResolvable, T : FirStatement {
@@ -97,7 +77,7 @@ abstract class AbstractManyCandidatesInferenceSession(
         }
 
         @Suppress("UNCHECKED_CAST")
-        val resolvedCalls = partiallyResolvedCalls as List<FirResolvable>
+        val resolvedCalls = partiallyResolvedCalls.map { it.first }
         val commonSystem = components.inferenceComponents.createConstraintSystem().apply {
             addOtherSystem(currentConstraintSystem)
         }

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

@@ -0,0 +1,220 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.resolve.inference
+
+import org.jetbrains.kotlin.fir.expressions.FirResolvable
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeStubType
+import org.jetbrains.kotlin.fir.types.ConeTypeVariable
+import org.jetbrains.kotlin.fir.types.ConeTypeVariableTypeConstructor
+import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
+import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
+import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind
+import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
+import org.jetbrains.kotlin.resolve.calls.inference.model.CoroutinePosition
+import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintSystemImpl
+import org.jetbrains.kotlin.types.model.TypeConstructorMarker
+
+class FirBuilderInferenceSession(
+    components: BodyResolveComponents,
+    postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
+    private val stubsForPostponedVariables: Map<ConeTypeVariable, ConeStubType>,
+) : AbstractManyCandidatesInferenceSession(components, postponedArgumentsAnalyzer) {
+    private val commonCalls: MutableList<Pair<FirStatement, Candidate>> = mutableListOf()
+
+    override fun <T> shouldRunCompletion(call: T): Boolean where T : FirResolvable, T : FirStatement {
+        val candidate = call.candidate
+        val system = candidate.system
+
+        if (system.hasContradiction) return true
+
+        val storage = system.getBuilder().currentStorage()
+
+        return !storage.notFixedTypeVariables.keys.any {
+            val variable = storage.allTypeVariables[it]
+            val isPostponed = variable != null && variable in storage.postponedTypeVariables
+            !isPostponed && !components.callCompleter.completer.variableFixationFinder.isTypeVariableHasProperConstraint(system, it)
+        } || call.hasPostponed()
+    }
+
+    private fun FirStatement.hasPostponed(): Boolean {
+        var result = false
+        processAllContainingCallCandidates(processBlocks = false) {
+            result = result || it.hasPostponed()
+        }
+        return result
+    }
+
+    private fun Candidate.hasPostponed(): Boolean {
+        return postponedAtoms.any { !it.analyzed }
+    }
+
+    override fun <T> addCompetedCall(call: T, candidate: Candidate) where T : FirResolvable, T : FirStatement {
+        if (skipCall(call)) return
+        commonCalls += call to candidate
+    }
+
+    override fun <T> writeOnlyStubs(call: T): Boolean where T : FirResolvable, T : FirStatement {
+        return !skipCall(call)
+    }
+
+    private fun <T> skipCall(call: T): Boolean where T : FirResolvable, T : FirStatement {
+        // TODO: what is FIR analog?
+        // if (descriptor is FakeCallableDescriptorForObject) return true
+        // if (!DescriptorUtils.isObject(descriptor) && isInLHSOfDoubleColonExpression(callInfo)) return true
+
+        return false
+    }
+
+    override val currentConstraintSystem: ConstraintStorage
+        get() = ConstraintStorage.Empty
+
+    override fun <T> shouldCompleteResolvedSubAtomsOf(call: T): Boolean where T : FirResolvable, T : FirStatement {
+        return true
+    }
+
+    override fun inferPostponedVariables(
+        lambda: ResolvedLambdaAtom,
+        initialStorage: ConstraintStorage
+    ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? {
+        val (commonSystem, effectivelyEmptyConstraintSystem) = buildCommonSystem(initialStorage)
+        if (effectivelyEmptyConstraintSystem) {
+            updateCalls(commonSystem, lambda)
+            return null
+        }
+
+        val context = commonSystem.asConstraintSystemCompleterContext()
+        @Suppress("UNCHECKED_CAST")
+        components.callCompleter.completer.complete(
+            context,
+            KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL,
+            partiallyResolvedCalls.map { it.first as FirStatement },
+            components.session.builtinTypes.unitType.type,
+            collectVariablesFromContext = true
+        ) {
+            error("Shouldn't be called in complete constraint system mode")
+        }
+
+        updateCalls(commonSystem, lambda)
+
+        @Suppress("UNCHECKED_CAST")
+        return commonSystem.fixedTypeVariables as Map<ConeTypeVariableTypeConstructor, ConeKotlinType>
+    }
+
+    private fun buildCommonSystem(initialStorage: ConstraintStorage): Pair<NewConstraintSystemImpl, Boolean> {
+        val commonSystem = components.inferenceComponents.createConstraintSystem()
+        val nonFixedToVariablesSubstitutor = createNonFixedTypeToVariableSubstitutor()
+
+        integrateConstraints(commonSystem, initialStorage, nonFixedToVariablesSubstitutor, false)
+
+        var effectivelyEmptyCommonSystem = true
+
+        for ((_, candidate) in commonCalls) {
+            val hasConstraints =
+                integrateConstraints(commonSystem, candidate.system.asReadOnlyStorage(), nonFixedToVariablesSubstitutor, false)
+            if (hasConstraints) effectivelyEmptyCommonSystem = false
+        }
+        for ((_, candidate) in partiallyResolvedCalls) {
+            val hasConstraints =
+                integrateConstraints(commonSystem, candidate.system.asReadOnlyStorage(), nonFixedToVariablesSubstitutor, true)
+            if (hasConstraints) effectivelyEmptyCommonSystem = false
+        }
+
+        // TODO: add diagnostics holder
+//        for (diagnostic in diagnostics) {
+//            commonSystem.addError(diagnostic)
+//        }
+
+        return commonSystem to effectivelyEmptyCommonSystem
+    }
+
+    private fun createNonFixedTypeToVariableSubstitutor(): ConeSubstitutor {
+        val ctx = components.inferenceComponents.ctx
+
+        val bindings = mutableMapOf<TypeConstructorMarker, ConeKotlinType>()
+        for ((variable, nonFixedType) in stubsForPostponedVariables) {
+            bindings[nonFixedType.variable.typeConstructor] = variable.defaultType
+        }
+
+        return ctx.typeSubstitutorByTypeConstructor(bindings)
+    }
+
+    private fun integrateConstraints(
+        commonSystem: NewConstraintSystemImpl,
+        storage: ConstraintStorage,
+        nonFixedToVariablesSubstitutor: ConeSubstitutor,
+        shouldIntegrateAllConstraints: Boolean
+    ): Boolean {
+        storage.notFixedTypeVariables.values.forEach { commonSystem.registerVariable(it.typeVariable) }
+
+        /*
+        * storage can contain the following substitutions:
+        *  TypeVariable(A) -> ProperType
+        *  TypeVariable(B) -> Special-Non-Fixed-Type
+        *
+        * while substitutor from parameter map non-fixed types to the original type variable
+        * */
+        val callSubstitutor = storage.buildAbstractResultingSubstitutor(commonSystem, transformTypeVariablesToErrorTypes = false) as ConeSubstitutor
+
+        var introducedConstraint = false
+
+        for (initialConstraint in storage.initialConstraints) {
+            val lower = nonFixedToVariablesSubstitutor.substituteOrSelf(callSubstitutor.substituteOrSelf(initialConstraint.a as ConeKotlinType)) // TODO: SUB
+            val upper = nonFixedToVariablesSubstitutor.substituteOrSelf(callSubstitutor.substituteOrSelf(initialConstraint.b as ConeKotlinType)) // TODO: SUB
+
+            if (commonSystem.isProperType(lower) && commonSystem.isProperType(upper)) continue
+
+            introducedConstraint = true
+
+            when (initialConstraint.constraintKind) {
+                ConstraintKind.LOWER -> error("LOWER constraint shouldn't be used, please use UPPER")
+
+                ConstraintKind.UPPER -> commonSystem.addSubtypeConstraint(lower, upper, initialConstraint.position)
+
+                ConstraintKind.EQUALITY ->
+                    with(commonSystem) {
+                        addSubtypeConstraint(lower, upper, initialConstraint.position)
+                        addSubtypeConstraint(upper, lower, initialConstraint.position)
+                    }
+            }
+        }
+
+        if (shouldIntegrateAllConstraints) {
+            for ((variableConstructor, type) in storage.fixedTypeVariables) {
+                val typeVariable = storage.allTypeVariables.getValue(variableConstructor)
+                commonSystem.registerVariable(typeVariable)
+                commonSystem.addEqualityConstraint((typeVariable as ConeTypeVariable).defaultType, type, CoroutinePosition())
+                introducedConstraint = true
+            }
+        }
+
+        return introducedConstraint
+    }
+
+    private fun updateCalls(commonSystem: NewConstraintSystemImpl, lambda: ResolvedLambdaAtom) {
+        val nonFixedToVariablesSubstitutor = createNonFixedTypeToVariableSubstitutor()
+        val commonSystemSubstitutor = commonSystem.buildCurrentSubstitutor() as ConeSubstitutor
+        val nonFixedTypesToResultSubstitutor = ConeComposedSubstitutor(commonSystemSubstitutor, nonFixedToVariablesSubstitutor)
+        val completionResultsWriter = components.callCompleter.createCompletionResultsWriter(nonFixedTypesToResultSubstitutor)
+        for ((call, _) in partiallyResolvedCalls) {
+            call.transformSingle(completionResultsWriter, null)
+            // TODO: support diagnostics, see CoroutineInferenceSession.kt:286
+        }
+    }
+}
+
+class ConeComposedSubstitutor(val left: ConeSubstitutor, val right: ConeSubstitutor) : ConeSubstitutor() {
+    override fun substituteOrNull(type: ConeKotlinType): ConeKotlinType? {
+        val rightSubstitution = right.substituteOrNull(type)
+        return left.substituteOrNull(rightSubstitution ?: type)
+    }
+}

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

@@ -9,13 +9,11 @@ import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
-import org.jetbrains.kotlin.fir.resolve.calls.candidate
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
@@ -25,10 +23,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolve.typeFromCallee
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
-import org.jetbrains.kotlin.fir.types.ConeKotlinErrorType
+import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.Name
@@ -38,10 +33,11 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystem
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.model.StubTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class FirCallCompleter(
     private val transformer: FirBodyResolveTransformer,
-    components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents
+    private val components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents
 ) : BodyResolveComponents by components {
     val completer = ConstraintSystemCompleter(components)
     private val inferenceSession
@@ -74,7 +70,7 @@ class FirCallCompleter(
 
         return when (completionMode) {
             ConstraintSystemCompletionMode.FULL -> {
-                if (inferenceSession.shouldRunCompletion(candidate)) {
+                if (inferenceSession.shouldRunCompletion(call)) {
                     completer.complete(candidate.system.asConstraintSystemCompleterContext(), completionMode, listOf(call), initialType) {
                         analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it, candidate)
                     }
@@ -89,7 +85,7 @@ class FirCallCompleter(
                         ),
                         null
                     )
-                    inferenceSession.addCompetedCall(completedCall)
+                    inferenceSession.addCompetedCall(completedCall, candidate)
                     CompletionResult(completedCall, true)
                 } else {
                     inferenceSession.addPartiallyResolvedCall(call)
@@ -124,13 +120,22 @@ class FirCallCompleter(
     }
 
     fun createPostponedArgumentsAnalyzer(): PostponedArgumentsAnalyzer {
+        val lambdaAnalyzer = LambdaAnalyzerImpl()
         return PostponedArgumentsAnalyzer(
-            LambdaAnalyzerImpl(), inferenceComponents,
+            lambdaAnalyzer, inferenceComponents,
             transformer.components.callResolver
-        )
+        ).also {
+            lambdaAnalyzer.initAnalyzer(it)
+        }
     }
 
     private inner class LambdaAnalyzerImpl : LambdaAnalyzer {
+        private lateinit var postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer
+
+        fun initAnalyzer(postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer) {
+            this.postponedArgumentsAnalyzer = postponedArgumentsAnalyzer
+        }
+
         override fun analyzeAndGetLambdaReturnArguments(
             lambdaAtom: ResolvedLambdaAtom,
             receiverType: ConeKotlinType?,
@@ -178,16 +183,26 @@ class FirCallCompleter(
             lambdaArgument.replaceValueParameters(lambdaArgument.valueParameters + listOfNotNull(itParam))
             lambdaArgument.replaceReturnTypeRef(expectedReturnTypeRef ?: noExpectedType)
 
+            val builderInferenceSession = runIf(stubsForPostponedVariables.isNotEmpty()) {
+                @Suppress("UNCHECKED_CAST")
+                FirBuilderInferenceSession(components, postponedArgumentsAnalyzer, stubsForPostponedVariables as Map<ConeTypeVariable, ConeStubType>)
+            }
+
             val localContext = towerDataContextForAnonymousFunctions.getValue(lambdaArgument.symbol)
             transformer.context.withTowerDataContext(localContext) {
-                lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
+                if (builderInferenceSession != null) {
+                    components.inferenceComponents.withInferenceSession(builderInferenceSession) {
+                        lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
+                    }
+                } else {
+                    lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
+                }
             }
             transformer.context.dropContextForAnonymousFunction(lambdaArgument)
 
             val returnArguments = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambdaArgument)
 
-            // TODO: add detecting of coroutine inference session
+            return ReturnArgumentsAnalysisResult(returnArguments, builderInferenceSession)
-            return ReturnArgumentsAnalysisResult(returnArguments, null)
         }
     }
 

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -26,19 +27,32 @@ import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class FirDelegatedPropertyInferenceSession(
     val property: FirProperty,
     initialCall: FirExpression,
     components: BodyResolveComponents,
     postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
-) : AbstractManyCandidatesInferenceSession(components, initialCall, postponedArgumentsAnalyzer) {
+) : AbstractManyCandidatesInferenceSession(components, postponedArgumentsAnalyzer) {
+    init {
+        val initialCandidate = (initialCall as? FirResolvable)
+            ?.calleeReference
+            ?.safeAs<FirNamedReferenceWithCandidate>()
+            ?.candidate
+        if (initialCandidate != null) {
+            addPartiallyResolvedCall(initialCall)
+        }
+    }
+
     val expectedType: ConeKotlinType? by lazy { property.returnTypeRef.coneTypeSafe() }
 
+    override fun <T> shouldRunCompletion(call: T): Boolean where T : FirResolvable, T : FirStatement = false
+
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
         initialStorage: ConstraintStorage
-    ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType> = emptyMap()
+    ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? = null
 
     override fun <T> shouldCompleteResolvedSubAtomsOf(call: T): Boolean where T : FirResolvable, T : FirStatement = true
 
@@ -89,4 +103,6 @@ class FirDelegatedPropertyInferenceSession(
         val substitutedType = substitutor.substituteOrSelf(valueParameterForThis.returnTypeRef.coneTypeUnsafe<ConeKotlinType>())
         commonSystem.addSubtypeConstraint(typeOfThis, substitutedType, SimpleConstraintSystemConstraintPosition)
     }
+
+    override fun <T> writeOnlyStubs(call: T): Boolean where T : FirResolvable, T : FirStatement = false
 }

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

@@ -15,40 +15,40 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 abstract class FirInferenceSession {
     companion object {
         val DEFAULT: FirInferenceSession = object : FirInferenceSession() {
-            override fun shouldRunCompletion(candidate: Candidate): Boolean = true
+            override fun <T> shouldRunCompletion(call: T): Boolean where T : FirResolvable, T : FirStatement = true
 
             override val currentConstraintSystem: ConstraintStorage
                 get() = ConstraintStorage.Empty
 
             override fun <T> addPartiallyResolvedCall(call: T) where T : FirResolvable, T : FirStatement {}
             override fun <T> addErrorCall(call: T) where T : FirResolvable, T : FirStatement {}
-            override fun <T> addCompetedCall(call: T) where T : FirResolvable, T : FirStatement {}
+            override fun <T> addCompetedCall(call: T, candidate: Candidate) where T : FirResolvable, T : FirStatement {}
 
             override fun inferPostponedVariables(
                 lambda: ResolvedLambdaAtom,
                 initialStorage: ConstraintStorage
-            ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType> = emptyMap()
+            ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? = null
 
+            override fun <T> writeOnlyStubs(call: T): Boolean where T : FirResolvable, T : FirStatement = false
             override fun <T> callCompleted(call: T): Boolean where T : FirResolvable, T : FirStatement = false
             override fun <T> shouldCompleteResolvedSubAtomsOf(call: T): Boolean where T : FirResolvable, T : FirStatement = true
         }
     }
 
-    abstract fun shouldRunCompletion(candidate: Candidate): Boolean
+    abstract fun <T> shouldRunCompletion(call: T): Boolean where T : FirResolvable, T : FirStatement
     abstract val currentConstraintSystem: ConstraintStorage
 
     abstract fun <T> addPartiallyResolvedCall(call: T) where T : FirResolvable, T : FirStatement
     abstract fun <T> addErrorCall(call: T) where T : FirResolvable, T : FirStatement
-    abstract fun <T> addCompetedCall(call: T) where T : FirResolvable, T : FirStatement
+    abstract fun <T> addCompetedCall(call: T, candidate: Candidate) where T : FirResolvable, T : FirStatement
 
     abstract fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
         initialStorage: ConstraintStorage,
         // TODO: diagnostic holder
-    ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>
+    ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>?
 
-//    TODO: do we need this?
+    abstract fun <T> writeOnlyStubs(call: T): Boolean where T : FirResolvable, T : FirStatement
-//    abstract fun writeOnlyStubs(): Boolean
     abstract fun <T> callCompleted(call: T): Boolean where T : FirResolvable, T : FirStatement
     abstract fun <T> shouldCompleteResolvedSubAtomsOf(call: T): Boolean where T : FirResolvable, T : FirStatement
 }

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

@@ -31,20 +31,21 @@ val Candidate.csBuilder: NewConstraintSystemImpl get() = system.getBuilder()
 
 
 class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
-    private val variableFixationFinder = VariableFixationFinder(components.inferenceComponents.trivialConstraintTypeInferenceOracle)
+    val variableFixationFinder = VariableFixationFinder(components.inferenceComponents.trivialConstraintTypeInferenceOracle)
 
     fun complete(
         c: KotlinConstraintSystemCompleter.Context,
         completionMode: ConstraintSystemCompletionMode,
         topLevelAtoms: List<FirStatement>,
         candidateReturnType: ConeKotlinType,
+        collectVariablesFromContext: Boolean = false,
         analyze: (PostponedResolvedAtom) -> Unit
     ) {
 
         while (true) {
             if (analyzePostponeArgumentIfPossible(c, topLevelAtoms, analyze)) continue
 
-            val allTypeVariables = getOrderedAllTypeVariables(c, topLevelAtoms)
+            val allTypeVariables = getOrderedAllTypeVariables(c, topLevelAtoms, collectVariablesFromContext)
             val postponedAtoms = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
             val variableForFixation =
                 variableFixationFinder.findFirstVariableForFixation(
@@ -167,8 +168,12 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
 
     private fun getOrderedAllTypeVariables(
         c: KotlinConstraintSystemCompleter.Context,
-        topLevelAtoms: List<FirStatement>
+        topLevelAtoms: List<FirStatement>,
+        collectVariablesFromContext: Boolean
     ): List<TypeConstructorMarker> {
+        if (collectVariablesFromContext) {
+            return c.notFixedTypeVariables.keys.toList()
+        }
         val result = LinkedHashSet<TypeConstructorMarker>(c.notFixedTypeVariables.size)
         fun ConeTypeVariable?.toTypeConstructor(): TypeConstructorMarker? =
             this?.typeConstructor?.takeIf { it in c.notFixedTypeVariables.keys }
@@ -179,14 +184,11 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
                     typeVariable.toTypeConstructor()
                 }
 
-                for (postponedAtom in candidate.postponedAtoms) {
+                for (lambdaAtom in candidate.postponedAtoms) {
-                    when (postponedAtom) {
+                    if (lambdaAtom is ResolvedLambdaAtom) {
-                        is ResolvedLambdaAtom -> postponedAtom.typeVariableForLambdaReturnType
+                        result.addIfNotNull(lambdaAtom.typeVariableForLambdaReturnType.toTypeConstructor())
                     }
                 }
-                for (lambdaAtom in candidate.postponedAtoms.filterIsInstance<ResolvedLambdaAtom>()) {
-                    result.addIfNotNull(lambdaAtom.typeVariableForLambdaReturnType.toTypeConstructor())
-                }
             }
         }
 
@@ -252,76 +254,74 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
         return notAnalyzedArguments
     }
 
-    private fun FirStatement.processAllContainingCallCandidates(processBlocks: Boolean, processor: (Candidate) -> Unit) {
-        when (this) {
-            is FirFunctionCall -> {
-                processCandidateIfApplicable(processor, processBlocks)
-                this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
-            }
-
-            is FirSafeCallExpression -> {
-                this.regularQualifiedAccess.processAllContainingCallCandidates(processBlocks, processor)
-            }
-
-            is FirWhenExpression -> {
-                processCandidateIfApplicable(processor, processBlocks)
-                this.branches.forEach { it.result.processAllContainingCallCandidates(processBlocks, processor) }
-            }
-
-            is FirTryExpression -> {
-                processCandidateIfApplicable(processor, processBlocks)
-                tryBlock.processAllContainingCallCandidates(processBlocks, processor)
-                catches.forEach { it.block.processAllContainingCallCandidates(processBlocks, processor) }
-            }
-
-            is FirCheckNotNullCall -> {
-                processCandidateIfApplicable(processor, processBlocks)
-                this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
-            }
-
-            is FirQualifiedAccessExpression -> {
-                processCandidateIfApplicable(processor, processBlocks)
-            }
-
-            is FirVariableAssignment -> {
-                processCandidateIfApplicable(processor, processBlocks)
-                rValue.processAllContainingCallCandidates(processBlocks, processor)
-            }
-
-            is FirWrappedArgumentExpression -> this.expression.processAllContainingCallCandidates(processBlocks, processor)
-            is FirBlock -> {
-                if (processBlocks) {
-                    this.returnExpressions().forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
-                }
-            }
-
-            is FirDelegatedConstructorCall -> {
-                processCandidateIfApplicable(processor, processBlocks)
-                this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
-            }
-        }
-    }
-
-    private fun FirResolvable.processCandidateIfApplicable(
-        processor: (Candidate) -> Unit,
-        processBlocks: Boolean
-    ) {
-        val candidate = (calleeReference as? FirNamedReferenceWithCandidate)?.candidate ?: return
-        processor(candidate)
-
-        for (atom in candidate.postponedAtoms) {
-            if (atom !is ResolvedLambdaAtom || !atom.analyzed) continue
-
-            atom.returnStatements.forEach {
-                it.processAllContainingCallCandidates(processBlocks, processor)
-            }
-        }
-    }
-
     private fun canWeAnalyzeIt(c: KotlinConstraintSystemCompleter.Context, argument: PostponedResolvedAtomMarker): Boolean {
         if (argument.analyzed) return false
-
         return argument.inputTypes.all { c.containsOnlyFixedOrPostponedVariables(it) }
     }
-
+}
+
+fun FirStatement.processAllContainingCallCandidates(processBlocks: Boolean, processor: (Candidate) -> Unit) {
+    when (this) {
+        is FirFunctionCall -> {
+            processCandidateIfApplicable(processor, processBlocks)
+            this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
+        }
+
+        is FirSafeCallExpression -> {
+            this.regularQualifiedAccess.processAllContainingCallCandidates(processBlocks, processor)
+        }
+
+        is FirWhenExpression -> {
+            processCandidateIfApplicable(processor, processBlocks)
+            this.branches.forEach { it.result.processAllContainingCallCandidates(processBlocks, processor) }
+        }
+
+        is FirTryExpression -> {
+            processCandidateIfApplicable(processor, processBlocks)
+            tryBlock.processAllContainingCallCandidates(processBlocks, processor)
+            catches.forEach { it.block.processAllContainingCallCandidates(processBlocks, processor) }
+        }
+
+        is FirCheckNotNullCall -> {
+            processCandidateIfApplicable(processor, processBlocks)
+            this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
+        }
+
+        is FirQualifiedAccessExpression -> {
+            processCandidateIfApplicable(processor, processBlocks)
+        }
+
+        is FirVariableAssignment -> {
+            processCandidateIfApplicable(processor, processBlocks)
+            rValue.processAllContainingCallCandidates(processBlocks, processor)
+        }
+
+        is FirWrappedArgumentExpression -> this.expression.processAllContainingCallCandidates(processBlocks, processor)
+        is FirBlock -> {
+            if (processBlocks) {
+                this.returnExpressions().forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
+            }
+        }
+
+        is FirDelegatedConstructorCall -> {
+            processCandidateIfApplicable(processor, processBlocks)
+            this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
+        }
+    }
+}
+
+private fun FirResolvable.processCandidateIfApplicable(
+    processor: (Candidate) -> Unit,
+    processBlocks: Boolean
+) {
+    val candidate = (calleeReference as? FirNamedReferenceWithCandidate)?.candidate ?: return
+    processor(candidate)
+
+    for (atom in candidate.postponedAtoms) {
+        if (atom !is ResolvedLambdaAtom || !atom.analyzed) continue
+
+        atom.returnStatements.forEach {
+            it.processAllContainingCallCandidates(processBlocks, processor)
+        }
+    }
 }

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

@@ -45,6 +45,13 @@ fun ConeKotlinType.receiverType(expectedTypeRef: FirTypeRef?, session: FirSessio
     return null
 }
 
+fun ConeKotlinType.receiverType(session: FirSession): ConeKotlinType? {
+    if (isBuiltinFunctionalType(session)) {
+        return ((this as ConeClassLikeType).fullyExpandedType(session).typeArguments.first() as ConeKotlinTypeProjection).type
+    }
+    return null
+}
+
 fun ConeKotlinType.returnType(session: FirSession): ConeKotlinType? {
     require(this is ConeClassLikeType)
     val projection = fullyExpandedType(session).typeArguments.last()

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

@@ -130,6 +130,24 @@ class PostponedArgumentsAnalyzer(
             stubsForPostponedVariables
         )
 
+        if (inferenceSession != null) {
+            val storageSnapshot = c.getBuilder().currentStorage()
+
+            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot)
+
+            if (postponedVariables == null) {
+                c.getBuilder().removePostponedVariables()
+            } else {
+                for ((constructor, resultType) in postponedVariables) {
+                    val variableWithConstraints = storageSnapshot.notFixedTypeVariables[constructor] ?: continue
+                    val variable = variableWithConstraints.typeVariable as ConeTypeVariable
+
+                    c.getBuilder().unmarkPostponedVariable(variable)
+                    c.getBuilder().addEqualityConstraint(variable.defaultType, resultType, CoroutinePosition())
+                }
+            }
+        }
+
         returnArguments.forEach { c.addSubsystemFromExpression(it) }
 
         val checkerSink: CheckerSink = CheckerSinkImpl(components)
@@ -157,20 +175,6 @@ class PostponedArgumentsAnalyzer(
 
         lambda.analyzed = true
         lambda.returnStatements = returnArguments
-
-        if (inferenceSession != null) {
-            val storageSnapshot = c.getBuilder().currentStorage()
-
-            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot)
-
-            for ((constructor, resultType) in postponedVariables) {
-                val variableWithConstraints = storageSnapshot.notFixedTypeVariables[constructor] ?: continue
-                val variable = variableWithConstraints.typeVariable as ConeTypeVariable
-
-                c.getBuilder().unmarkPostponedVariable(variable)
-                c.getBuilder().addEqualityConstraint(variable.defaultType, resultType, CoroutinePosition())
-            }
-        }
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -371,20 +371,31 @@ class FirCallCompletionResultsWriterTransformer(
         data: ExpectedArgumentType?,
     ): CompositeTransformResult<FirStatement> {
         val expectedType = data?.getExpectedType(anonymousFunction)?.takeIf { it.isBuiltinFunctionalType(session) }
-        val expectedReturnType = expectedType?.returnType(session) as? ConeClassLikeType
 
+        var needUpdateLambdaType = false
+
+        val initialReceiverType = anonymousFunction.receiverTypeRef?.coneTypeSafe<ConeKotlinType>()
+        val resultReceiverType = initialReceiverType?.let { finalSubstitutor.substituteOrNull(it) }
+        if (resultReceiverType != null) {
+            anonymousFunction.replaceReceiverTypeRef(anonymousFunction.receiverTypeRef!!.resolvedTypeFromPrototype(resultReceiverType))
+            needUpdateLambdaType = true
+        }
+
+        val expectedReturnType = expectedType?.returnType(session) as? ConeClassLikeType
         val initialType = anonymousFunction.returnTypeRef.coneTypeSafe<ConeKotlinType>()
         if (initialType != null) {
             val finalType = expectedReturnType ?: finalSubstitutor.substituteOrNull(initialType)
-
             val resultType = anonymousFunction.returnTypeRef.withReplacedConeType(finalType)
-
             anonymousFunction.transformReturnTypeRef(StoreType, resultType)
+            needUpdateLambdaType = true
+        }
 
+        if (needUpdateLambdaType) {
             anonymousFunction.replaceTypeRef(
                 anonymousFunction.constructFunctionalTypeRef(session, isSuspend = expectedType?.isSuspendFunctionType(session) == true)
             )
         }
+
         val result = transformElement(anonymousFunction, null)
         val resultFunction = result.single
         if (resultFunction.returnTypeRef.coneTypeSafe<ConeIntegerLiteralType>() != null) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeInferenceContext.kt

@@ -295,19 +295,19 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
         return this.original as SimpleTypeMarker
     }
 
-    override fun typeSubstitutorByTypeConstructor(map: Map<TypeConstructorMarker, KotlinTypeMarker>): TypeSubstitutorMarker {
+    override fun typeSubstitutorByTypeConstructor(map: Map<TypeConstructorMarker, KotlinTypeMarker>): ConeSubstitutor {
         if (map.isEmpty()) return createEmptySubstitutor()
         return object : AbstractConeSubstitutor(),
             TypeSubstitutorMarker {
             override fun substituteType(type: ConeKotlinType): ConeKotlinType? {
-                if (type !is ConeLookupTagBasedType) return null
+                if (type !is ConeLookupTagBasedType && type !is ConeStubType) return null
                 val new = map[type.typeConstructor()] ?: return null
                 return (new as ConeKotlinType).approximateIntegerLiteralType().updateNullabilityIfNeeded(type)
             }
         }
     }
 
-    override fun createEmptySubstitutor(): TypeSubstitutorMarker {
+    override fun createEmptySubstitutor(): ConeSubstitutor {
         return ConeSubstitutor.Empty
     }
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationUtil.kt

@@ -14,10 +14,9 @@ import org.jetbrains.kotlin.fir.declarations.impl.FirFileImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousObjectSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.ConeFlexibleType
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import org.jetbrains.kotlin.name.ClassId
 
 fun FirTypeParameterBuilder.addDefaultBoundIfNecessary(isFlexible: Boolean = false) {
     if (bounds.isEmpty()) {
@@ -114,3 +113,7 @@ fun FirRegularClass.addDeclaration(declaration: FirDeclaration) {
 
 private object IsFromVarargKey: FirDeclarationDataKey()
 var FirProperty.isFromVararg: Boolean? by FirDeclarationDataRegistry.data(IsFromVarargKey)
+
+fun FirAnnotatedDeclaration.hasAnnotation(classId: ClassId): Boolean {
+    return annotations.any { it.annotationTypeRef.coneTypeSafe<ConeClassLikeType>()?.classId == classId }
+}

compiler/testData/codegen/box/coroutines/asyncIteratorNullMerge_1_3.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 

compiler/testData/codegen/box/coroutines/asyncIteratorToList_1_3.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // KJS_WITH_FULL_RUNTIME
 // WITH_RUNTIME
 // WITH_COROUTINES

compiler/testData/codegen/box/coroutines/asyncIterator_1_3.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 

compiler/testData/codegen/box/coroutines/emptyCommonConstraintSystemForCoroutineInferenceCall.kt

@@ -1,7 +1,6 @@
 // WITH_RUNTIME
 // WITH_COROUTINES
 // KJS_WITH_FULL_RUNTIME
-// IGNORE_BACKEND_FIR: JVM_IR
 
 import helpers.*
 import kotlin.coroutines.*

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

@@ -1,19 +0,0 @@
-// !LANGUAGE: +NewInference
-// !DIAGNOSTICS: -UNUSED_PARAMETER
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-
-import kotlin.experimental.ExperimentalTypeInference
-
-suspend fun main() {
-    iFlow { <!INAPPLICABLE_CANDIDATE!>emit<!>(1) }
-}
-
-
-@OptIn(ExperimentalTypeInference::class)
-fun <K> iFlow(@BuilderInference block: suspend iFlowCollector<in K>.() -> Unit): iFlow<K> = TODO()
-
-interface iFlowCollector<S> {
-    suspend fun emit(value: S)
-}
-
-interface iFlow<out V>

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +NewInference
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // !USE_EXPERIMENTAL: kotlin.RequiresOptIn

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

@@ -1,25 +0,0 @@
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
-
-@file:OptIn(ExperimentalTypeInference::class)
-
-import kotlin.experimental.ExperimentalTypeInference
-
-interface Controller<T> {
-    suspend fun yield(t: T) {}
-
-    fun justString(): String = ""
-
-    fun <Z> generidFun(t: Z) = t
-}
-
-fun <S> generate(@BuilderInference g: suspend Controller<S>.() -> Unit): S = TODO()
-
-val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(justString())
-}
-
-val test2 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(generidFun(2))
-}
-

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

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

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

@@ -1,38 +0,0 @@
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// !DIAGNOSTICS: -UNUSED_VARIABLE -UNUSED_PARAMETER -UNUSED_EXPRESSION
-
-@file:OptIn(ExperimentalTypeInference::class)
-
-package a.b
-
-import kotlin.experimental.ExperimentalTypeInference
-
-class BatchInfo1(val batchSize: Int)
-class BatchInfo2<T>(val data: T)
-
-object Obj
-
-fun test1() {
-    val a: Sequence<String> = sequence {
-        val x = BatchInfo1::class
-        val y = a.b.BatchInfo1::class
-        val z = Obj::class
-
-        val x1 = BatchInfo1::batchSize
-        val y1 = a.b.BatchInfo1::class
-    }
-}
-
-interface Scope<T> {
-    fun yield(t: T) {}
-}
-
-fun <S> generate(@BuilderInference g: Scope<S>.() -> Unit): S = TODO()
-
-val test2 = generate {
-    { <!INAPPLICABLE_CANDIDATE!>yield<!>("foo") }::class
-}
-
-val test3 = generate {
-    ({ <!INAPPLICABLE_CANDIDATE!>yield<!>("foo") })::class
-}

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

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

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

@@ -1,43 +0,0 @@
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
-
-@file:OptIn(ExperimentalTypeInference::class)
-
-import kotlin.experimental.ExperimentalTypeInference
-
-class GenericController<T> {
-    suspend fun yield(t: T) {}
-}
-
-@BuilderInference
-suspend fun <K> GenericController<K>.yieldAll(s: Collection<K>) {}
-
-fun <S> generate(@BuilderInference g: suspend GenericController<S>.() -> Unit): S = TODO()
-
-val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(4)
-    <!INAPPLICABLE_CANDIDATE!>yieldAll<!>(setOf(4, 5))
-}
-
-val test2 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yieldAll<!>(setOf(B))
-}
-
-val test3 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yieldAll<!>(setOf(B, C))
-}
-
-val test4 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yieldAll<!>(setOf(B))
-
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(C)
-}
-
-
-
-// Utils
-fun <X> setOf(vararg x: X): Set<X> = TODO()
-
-interface A
-object B : A
-object C : A

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

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

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

@@ -25,7 +25,7 @@ fun <S> Controller<S>.genericExtension() {}
 fun Controller<String>.safeExtension() {}
 
 val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>("foo")
+    yield("foo")
     baseExtension()
 }
 
@@ -34,7 +34,7 @@ val test2 = generate {
 }
 
 val test3 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(42)
+    yield(42)
     outNullableAnyExtension()
 }
 
@@ -43,20 +43,20 @@ val test4 = generate {
 }
 
 val test5 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(42)
+    yield(42)
     outAnyExtension()
 }
 
 val test6 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>("bar")
+    yield("bar")
-    <!INAPPLICABLE_CANDIDATE!>invNullableAnyExtension<!>()
+    invNullableAnyExtension()
 }
 
 val test7 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>("baz")
+    yield("baz")
-    <!INAPPLICABLE_CANDIDATE!>genericExtension<!><Int>()
+    genericExtension<Int>()
 }
 
 val test8 = generate {
-    <!INAPPLICABLE_CANDIDATE!>safeExtension<!>()
+    safeExtension()
 }

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

@@ -25,7 +25,7 @@ fun Base<String>.stringBase() {}
 
 val test1 = generate {
     starBase()
-    <!INAPPLICABLE_CANDIDATE!>yield<!>("foo")
+    yield("foo")
 }
 
 val test2 = generate {
@@ -33,17 +33,17 @@ val test2 = generate {
 }
 
 val test3 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>("bar")
+    yield("bar")
-    <!INAPPLICABLE_CANDIDATE!>stringBase<!>()
+    stringBase()
 }
 
 val test4 = generateSpecific {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(42)
+    yield(42)
     starBase()
 }
 
 val test5 = generateSpecific {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(42)
+    yield(42)
     stringBase()
 }
 

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

@@ -17,20 +17,20 @@ class GenericController<T> {
 fun <S> generate(@BuilderInference g: suspend GenericController<S>.() -> Unit): List<S> = TODO()
 
 val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(3)
+    yield(3)
 }
 
 val test2 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(3)
+    yield(3)
-    <!INAPPLICABLE_CANDIDATE!>notYield<!>(3)
+    notYield(3)
 }
 
 val test3 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(3)
+    yield(3)
-    <!INAPPLICABLE_CANDIDATE!>yieldBarReturnType<!>(3)
+    yieldBarReturnType(3)
 }
 
 val test4 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(3)
+    yield(3)
     barReturnType()
 }

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

@@ -6,21 +6,21 @@
 import kotlin.experimental.ExperimentalTypeInference
 
 fun test_1() {
-    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Nothing")!>sequence {
+    <!DEBUG_INFO_EXPRESSION_TYPE("Inv<kotlin.Int>")!>sequence {
-        <!INAPPLICABLE_CANDIDATE!>yield<!>(<!DEBUG_INFO_EXPRESSION_TYPE("Inv<kotlin.Any?>")!>materialize()<!>)
+        yield(<!DEBUG_INFO_EXPRESSION_TYPE("Inv<kotlin.Int>")!>materialize()<!>)
-        <!INAPPLICABLE_CANDIDATE!>yield<!>(materialize<Int>())
+        yield(materialize<Int>())
     }<!>
 }
 
 fun test_2() {
     sequence {
-        <!INAPPLICABLE_CANDIDATE!>yield<!>(materialize())
+        yield(materialize())
     }
 }
 
 fun test_3() {
     sequence {
-        <!INAPPLICABLE_CANDIDATE!>yield<!>(materialize<Int>())
+        yield(materialize<Int>())
         materialize()
     }
 }

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

@@ -1,34 +0,0 @@
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE -UNUSED_EXPRESSION
-
-import kotlin.experimental.ExperimentalTypeInference
-
-abstract class AnimationVector
-class AnimationVector1D : AnimationVector()
-interface PropKey<T, V : AnimationVector>
-class IntPropKey : PropKey<Int, AnimationVector1D>
-abstract class AnimationBuilder<T>
-abstract class DurationBasedAnimationBuilder<T> : AnimationBuilder<T>()
-class TweenBuilder<T> : DurationBasedAnimationBuilder<T>()
-class TransitionSpec<S> {
-    fun <T> tween(init: TweenBuilder<T>.() -> Unit): DurationBasedAnimationBuilder<T> = TweenBuilder<T>().apply(init)
-    infix fun <T, V : AnimationVector> PropKey<T, V>.using(builder: AnimationBuilder<T>) {}
-}
-class TransitionDefinition<T> {
-    fun transition(fromState: T? = null, toState: T? = null, init: TransitionSpec<T>.() -> Unit) {}
-}
-@OptIn(ExperimentalTypeInference::class)
-fun <T> transitionDefinition(@BuilderInference init: TransitionDefinition<T>.() -> Unit) = TransitionDefinition<T>().apply(init)
-
-fun main() {
-    val intProp = IntPropKey()
-    val defn = transitionDefinition {
-        transition(1, 2) {
-            intProp using tween {
-
-            }
-        }
-    }
-
-    <!DEBUG_INFO_EXPRESSION_TYPE("TransitionDefinition<kotlin.Any?>")!>defn<!>
-}

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

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

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

@@ -1,20 +0,0 @@
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// !DIAGNOSTICS: -UNUSED_PARAMETER
-
-import kotlin.experimental.ExperimentalTypeInference
-
-fun test() {
-    flow {
-        <!INAPPLICABLE_CANDIDATE!>emit<!>(42)
-        kotlin.coroutines.coroutineContext
-    }
-}
-
-@OptIn(ExperimentalTypeInference::class)
-fun <T> flow(@BuilderInference block: suspend FlowCollector<T>.() -> Unit): Flow<T> = TODO()
-
-interface Flow<out T>
-
-interface FlowCollector<in T> {
-    suspend fun emit(value: T)
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !USE_EXPERIMENTAL: kotlin.RequiresOptIn
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 

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

@@ -1,22 +0,0 @@
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
-
-@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()
-
-val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(generate {
-        yield(generate {
-            yield(generate {
-                yield(3)
-            })
-        })
-    })
-}

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

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

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

@@ -18,7 +18,7 @@ suspend fun <S> GenericController<List<S>>.yieldGenerate(g: suspend GenericContr
 
 val test1 = generate {
     // TODO: KT-15185
-    <!INAPPLICABLE_CANDIDATE!>yieldGenerate<!> {
+    yieldGenerate {
-        <!INAPPLICABLE_CANDIDATE!>yield<!>(4)
+        yield(4)
     }
 }

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

@@ -21,7 +21,7 @@ fun <S> Builder<S>.extensionAdd(s: S) {}
 fun <S> Builder<S>.safeExtensionAdd(s: S) {}
 
 val member = build {
-    <!INAPPLICABLE_CANDIDATE!>add<!>(42)
+    add(42)
 }
 
 val memberWithoutAnn = wrongBuild {
@@ -29,9 +29,9 @@ val memberWithoutAnn = wrongBuild {
 }
 
 val extension = build {
-    <!INAPPLICABLE_CANDIDATE!>extensionAdd<!>("foo")
+    extensionAdd("foo")
 }
 
 val safeExtension = build {
-    <!INAPPLICABLE_CANDIDATE!>safeExtensionAdd<!>("foo")
+    safeExtensionAdd("foo")
 }

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

@@ -1,34 +0,0 @@
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
-
-@file:OptIn(ExperimentalTypeInference::class)
-
-import kotlin.experimental.ExperimentalTypeInference
-
-class GenericController<T> {
-    suspend fun yield(t: T) {}
-    suspend fun yieldSet(t: Set<T>) {}
-    suspend fun yieldVararg(vararg t: T) {}
-}
-
-fun <S> generate(@BuilderInference g: suspend GenericController<S>.() -> Unit): S = TODO()
-
-val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(4)
-}
-
-val test2 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yieldSet<!>(setOf(1, 2, 3))
-}
-
-val test3 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yieldVararg<!>(1, 2, 3)
-}
-
-val test4 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yieldVararg<!>(1, 2, "")
-}
-
-
-// Util function
-fun <X> setOf(vararg x: X): Set<X> = TODO()

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

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

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

@@ -16,9 +16,9 @@ fun <S> generate(@BuilderInference g: suspend Controller<S>.() -> Unit): S = TOD
 class A
 
 val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(A)
+    yield(A)
 }
 
 val test2: Int = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(A())
+    yield(A())
 }

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

@@ -1,17 +0,0 @@
-// !USE_EXPERIMENTAL: kotlin.RequiresOptIn
-// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
-
-@file:OptIn(ExperimentalTypeInference::class)
-
-import kotlin.experimental.ExperimentalTypeInference
-
-class GenericController<T>
-
-fun <S> generate(@BuilderInference g: suspend GenericController<S>.() -> Unit): List<S> = TODO()
-
-@BuilderInference
-suspend fun GenericController<List<String>>.test() {}
-
-val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>test<!>()
-}

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

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

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

@@ -21,13 +21,13 @@ suspend fun <S> GenericController<S>.safeExtensionYield(s: S) {}
 fun <S> generate(@BuilderInference g: suspend GenericController<S>.() -> Unit): List<S> = TODO()
 
 val normal = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>(42)
+    yield(42)
 }
 
 val extension = generate {
-    <!INAPPLICABLE_CANDIDATE!>extensionYield<!>("foo")
+    extensionYield("foo")
 }
 
 val safeExtension = generate {
-    <!INAPPLICABLE_CANDIDATE!>safeExtensionYield<!>("foo")
+    safeExtensionYield("foo")
 }

compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/outerYield_1_3.fir.kt

@@ -84,10 +84,10 @@ fun test() {
 
             with(this) {
                 yield("")
-                this@with.yield("")
+                this@with.<!UNRESOLVED_REFERENCE!>yield<!>("")
 
                 yield2("")
-                this@with.yield2("")
+                this@with.<!INAPPLICABLE_CANDIDATE!>yield2<!>("")
             }
         }
     }

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

@@ -14,8 +14,8 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
             FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.FlowCollector<R of <root>.scopedFlow>) returnType:kotlin.Unit [suspend]
               $receiver: VALUE_PARAMETER name:<this> type:<root>.FlowCollector<R of <root>.scopedFlow>
               BLOCK_BODY
-                VAR name:collector type:<root>.FlowCollector<R of <root>.scopedFlow> [val]
+                VAR name:collector type:<root>.FlowCollector<IrErrorType> [val]
-                  GET_VAR '<this>: <root>.FlowCollector<R of <root>.scopedFlow> declared in <root>.scopedFlow.<anonymous>' type=<root>.FlowCollector<R of <root>.scopedFlow> origin=null
+                  GET_VAR '<this>: <root>.FlowCollector<R of <root>.scopedFlow> declared in <root>.scopedFlow.<anonymous>' type=<root>.FlowCollector<IrErrorType> origin=null
                 CALL 'public final fun flowScope <R> (block: kotlin.coroutines.SuspendFunction1<<root>.CoroutineScope, R of <root>.flowScope>): R of <root>.flowScope [suspend] declared in <root>' type=kotlin.Unit origin=null
                   <R>: kotlin.Unit
                   block: FUN_EXPR type=kotlin.coroutines.SuspendFunction1<<root>.CoroutineScope, kotlin.Unit> origin=LAMBDA
@@ -25,7 +25,7 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
                         CALL 'public abstract fun invoke (p1: P1 of kotlin.coroutines.SuspendFunction2, p2: P2 of kotlin.coroutines.SuspendFunction2): R of kotlin.coroutines.SuspendFunction2 [suspend,operator] declared in kotlin.coroutines.SuspendFunction2' type=kotlin.Unit origin=null
                           $this: GET_VAR 'block: kotlin.coroutines.SuspendFunction2<<root>.CoroutineScope, <root>.FlowCollector<R of <root>.scopedFlow>, kotlin.Unit> declared in <root>.scopedFlow' type=kotlin.coroutines.SuspendFunction2<<root>.CoroutineScope, <root>.FlowCollector<R of <root>.scopedFlow>, kotlin.Unit> origin=null
                           p1: GET_VAR '<this>: <root>.CoroutineScope declared in <root>.scopedFlow.<anonymous>.<anonymous>' type=<root>.CoroutineScope origin=null
-                          p2: GET_VAR 'val collector: <root>.FlowCollector<R of <root>.scopedFlow> [val] declared in <root>.scopedFlow.<anonymous>' type=<root>.FlowCollector<R of <root>.scopedFlow> origin=null
+                          p2: GET_VAR 'val collector: <root>.FlowCollector<IrErrorType> [val] declared in <root>.scopedFlow.<anonymous>' type=<root>.FlowCollector<IrErrorType> origin=null
   FUN name:onCompletion visibility:public modality:FINAL <T> ($receiver:<root>.Flow<T of <root>.onCompletion>, action:kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.onCompletion>, kotlin.Throwable?, kotlin.Unit>) returnType:<root>.Flow<T of <root>.onCompletion>
     TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?]
     $receiver: VALUE_PARAMETER name:<this> type:<root>.Flow<T of <root>.onCompletion>
@@ -38,13 +38,13 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
             FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.FlowCollector<T of <root>.onCompletion>) returnType:kotlin.Unit [suspend]
               $receiver: VALUE_PARAMETER name:<this> type:<root>.FlowCollector<T of <root>.onCompletion>
               BLOCK_BODY
-                VAR name:safeCollector type:<root>.SafeCollector<T of <root>.onCompletion> [val]
+                VAR name:safeCollector type:<root>.SafeCollector<IrErrorType> [val]
-                  CONSTRUCTOR_CALL 'public constructor <init> (collector: <root>.FlowCollector<T of <root>.SafeCollector>) [primary] declared in <root>.SafeCollector' type=<root>.SafeCollector<T of <root>.onCompletion> origin=null
+                  CONSTRUCTOR_CALL 'public constructor <init> (collector: <root>.FlowCollector<T of <root>.SafeCollector>) [primary] declared in <root>.SafeCollector' type=<root>.SafeCollector<IrErrorType> origin=null
-                    <class: T>: T of <root>.onCompletion
+                    <class: T>: IrErrorType
-                    collector: GET_VAR '<this>: <root>.FlowCollector<T of <root>.onCompletion> declared in <root>.onCompletion.<anonymous>' type=<root>.FlowCollector<T of <root>.onCompletion> origin=null
+                    collector: GET_VAR '<this>: <root>.FlowCollector<T of <root>.onCompletion> declared in <root>.onCompletion.<anonymous>' type=<root>.FlowCollector<IrErrorType> origin=null
                 CALL 'public final fun invokeSafely <T> (action: kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.invokeSafely>, kotlin.Throwable?, kotlin.Unit>): kotlin.Unit [suspend] declared in <root>' type=kotlin.Unit origin=null
                   <T>: T of <root>.onCompletion
-                  $receiver: GET_VAR 'val safeCollector: <root>.SafeCollector<T of <root>.onCompletion> [val] declared in <root>.onCompletion.<anonymous>' type=<root>.SafeCollector<T of <root>.onCompletion> origin=null
+                  $receiver: GET_VAR 'val safeCollector: <root>.SafeCollector<IrErrorType> [val] declared in <root>.onCompletion.<anonymous>' type=<root>.SafeCollector<IrErrorType> origin=null
                   action: GET_VAR 'action: kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.onCompletion>, kotlin.Throwable?, kotlin.Unit> declared in <root>.onCompletion' type=kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.onCompletion>, kotlin.Throwable?, kotlin.Unit> origin=null
   FUN name:invokeSafely visibility:public modality:FINAL <T> ($receiver:<root>.FlowCollector<T of <root>.invokeSafely>, action:kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.invokeSafely>, kotlin.Throwable?, kotlin.Unit>) returnType:kotlin.Unit [suspend]
     TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?]
@@ -90,8 +90,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.Any> origin=GET_PROPERTY
+                    CALL 'public abstract fun <get-channel> (): <root>.SendChannel<E of <root>.ProducerScope> declared in <root>.ProducerScope' type=<root>.SendChannel<IrErrorType> origin=GET_PROPERTY
-                      $this: GET_VAR '<this>: <root>.ProducerScope<kotlin.Any> declared in <root>.asFairChannel.<anonymous>' type=<root>.ProducerScope<kotlin.Any> origin=null
+                      $this: GET_VAR '<this>: <root>.ProducerScope<kotlin.Any> declared in <root>.asFairChannel.<anonymous>' type=<root>.ProducerScope<IrErrorType> origin=null
                 CALL 'public abstract fun collect (collector: <root>.FlowCollector<T of <root>.Flow>): kotlin.Unit [suspend] declared in <root>.Flow' type=kotlin.Unit origin=null
                   $this: GET_VAR 'flow: <root>.Flow<*> declared in <root>.asFairChannel' type=<root>.Flow<*> origin=null
                   collector: TYPE_OP type=<root>.FlowCollector<kotlin.Any?> origin=SAM_CONVERSION typeOperand=<root>.FlowCollector<kotlin.Any?>
@@ -136,8 +136,8 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
                         BLOCK_BODY
                           RETURN type=kotlin.Nothing from='local final fun <anonymous> (value: kotlin.Any?): kotlin.Unit 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<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<IrErrorType> origin=GET_PROPERTY
-                                $this: GET_VAR '<this>: <root>.ProducerScope<kotlin.Any> declared in <root>.asChannel.<anonymous>' type=<root>.ProducerScope<kotlin.Any> origin=null
+                                $this: GET_VAR '<this>: <root>.ProducerScope<kotlin.Any> declared in <root>.asChannel.<anonymous>' type=<root>.ProducerScope<IrErrorType> origin=null
                               e: BLOCK type=kotlin.Any origin=ELVIS
                                 VAR IR_TEMPORARY_VARIABLE name:tmp_1 type:kotlin.Any? [val]
                                   GET_VAR 'value: kotlin.Any? declared in <root>.asChannel.<anonymous>.<anonymous>' type=kotlin.Any? origin=null