[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.csBuilder
+import org.jetbrains.kotlin.fir.resolve.inference.*
 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
-            ?.safeAs<FirNamedReferenceWithCandidate>()
-            ?.candidate
+            ?.second
             ?.system
             ?.currentStorage()
             ?: ConstraintStorage.Empty
 
     private lateinit var resultingConstraintSystem: NewConstraintSystem
 
-    override fun shouldRunCompletion(candidate: Candidate): Boolean {
-        return false
-    }
-
-    override fun <T> addCompetedCall(call: T) where T : FirResolvable, T : FirStatement {
+    override fun <T> addCompetedCall(call: T, candidate: Candidate) 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.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.*
 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, null)
+            return ReturnArgumentsAnalysisResult(returnArguments, builderInferenceSession)
         }
     }
 

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 writeOnlyStubs(): Boolean
+    abstract fun <T> writeOnlyStubs(call: T): Boolean where T : FirResolvable, T : FirStatement
     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) {
-                    when (postponedAtom) {
-                        is ResolvedLambdaAtom -> postponedAtom.typeVariableForLambdaReturnType
+                for (lambdaAtom in candidate.postponedAtoms) {
+                    if (lambdaAtom is ResolvedLambdaAtom) {
+                        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.ConeFlexibleType
+import org.jetbrains.kotlin.fir.types.*
 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 }
+}