[NI] Update calls after inference for coroutines

compiler/frontend/src/org/jetbrains/kotlin/resolve/DelegatedPropertyInferenceSession.kt

@@ -15,11 +15,10 @@ import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.inference.model.DelegatedPropertyConstraintPosition
-import org.jetbrains.kotlin.resolve.calls.inference.model.ExpectedTypeConstraintPosition
-import org.jetbrains.kotlin.resolve.calls.inference.model.NewTypeVariable
 import org.jetbrains.kotlin.resolve.calls.model.KotlinCallComponents
 import org.jetbrains.kotlin.resolve.calls.model.KotlinResolutionCandidate
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCallAtom
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedLambdaAtom
 import org.jetbrains.kotlin.resolve.calls.tower.ManyCandidatesResolver
 import org.jetbrains.kotlin.resolve.calls.tower.PSICallResolver
 import org.jetbrains.kotlin.resolve.calls.tower.PSIPartialCallInfo
@@ -83,7 +82,12 @@ class DelegatedPropertyInferenceSession(
         addConstraintForThis(candidateDescriptor, commonSystem)
     }
 
-    override fun inferPostponedVariables(initialStorage: ConstraintStorage): Map<TypeConstructor, UnwrappedType> = emptyMap()
+    override fun inferPostponedVariables(
+        lambda: ResolvedLambdaAtom,
+        initialStorage: ConstraintStorage
+    ): Map<TypeConstructor, UnwrappedType> = emptyMap()
+
+    override fun writeOnlyStubs(): Boolean = false
 }
 
 object InferenceSessionForExistingCandidates : InferenceSession {
@@ -96,5 +100,10 @@ object InferenceSessionForExistingCandidates : InferenceSession {
     override fun addErrorCallInfo(callInfo: ErrorCallInfo) {}
 
     override fun currentConstraintSystem(): ConstraintStorage = ConstraintStorage.Empty
-    override fun inferPostponedVariables(initialStorage: ConstraintStorage): Map<TypeConstructor, UnwrappedType> = emptyMap()
+    override fun inferPostponedVariables(
+        lambda: ResolvedLambdaAtom,
+        initialStorage: ConstraintStorage
+    ): Map<TypeConstructor, UnwrappedType> = emptyMap()
+
+    override fun writeOnlyStubs(): Boolean = false
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/DelegatedPropertyResolver.kt

@@ -569,7 +569,7 @@ class DelegatedPropertyResolver(
             )
         }
 
-        val resolutionCallbacks = psiCallResolver.createResolutionCallbacks(trace, inferenceSession)
+        val resolutionCallbacks = psiCallResolver.createResolutionCallbacks(trace, inferenceSession, context = null)
         inferenceSession.resolveCandidates(resolutionCallbacks)
 
         val resolvedDelegateType = extractResolvedDelegateType(delegateExpression, trace)

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/inference/CoroutineInferenceSession.kt

@@ -8,9 +8,14 @@ package org.jetbrains.kotlin.resolve.calls.inference
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.DeprecationResolver
+import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
+import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.components.CompletedCallInfo
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzer
+import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.inference.components.*
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
@@ -18,10 +23,13 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintSystemImp
 import org.jetbrains.kotlin.resolve.calls.inference.model.NewTypeVariable
 import org.jetbrains.kotlin.resolve.calls.model.KotlinCallComponents
 import org.jetbrains.kotlin.resolve.calls.model.KotlinResolutionCandidate
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedLambdaAtom
 import org.jetbrains.kotlin.resolve.calls.tower.*
 import org.jetbrains.kotlin.types.NonFixedType
 import org.jetbrains.kotlin.types.TypeConstructor
 import org.jetbrains.kotlin.types.UnwrappedType
+import org.jetbrains.kotlin.types.expressions.DoubleColonExpressionResolver
+import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 
 class CoroutineInferenceSession(
     psiCallResolver: PSICallResolver,
@@ -29,9 +37,15 @@ class CoroutineInferenceSession(
     kotlinConstraintSystemCompleter: KotlinConstraintSystemCompleter,
     callComponents: KotlinCallComponents,
     builtIns: KotlinBuiltIns,
+    private val topLevelCallContext: BasicCallResolutionContext,
     private val stubsForPostponedVariables: Map<NewTypeVariable, NonFixedType>,
     private val trace: BindingTrace,
-    private val kotlinToResolvedCallTransformer: KotlinToResolvedCallTransformer
+    private val kotlinToResolvedCallTransformer: KotlinToResolvedCallTransformer,
+    private val expressionTypingServices: ExpressionTypingServices,
+    private val argumentTypeResolver: ArgumentTypeResolver,
+    private val doubleColonExpressionResolver: DoubleColonExpressionResolver,
+    private val deprecationResolver: DeprecationResolver,
+    private val moduleDescriptor: ModuleDescriptor
 ) : ManyCandidatesResolver<CallableDescriptor>(
     psiCallResolver, postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter, callComponents, builtIns
 ) {
@@ -50,17 +64,24 @@ class CoroutineInferenceSession(
             normalCompletedCalls.add(callInfo)
     }
 
+    override fun writeOnlyStubs(): Boolean {
+        return true
+    }
+
     override fun currentConstraintSystem(): ConstraintStorage {
         return ConstraintStorage.Empty
     }
 
-    override fun inferPostponedVariables(initialStorage: ConstraintStorage): Map<TypeConstructor, UnwrappedType> {
+    override fun inferPostponedVariables(
+        lambda: ResolvedLambdaAtom,
+        initialStorage: ConstraintStorage
+    ): Map<TypeConstructor, UnwrappedType> {
         val commonSystem = buildCommonSystem(initialStorage)
 
         val context = commonSystem.asConstraintSystemCompleterContext()
         kotlinConstraintSystemCompleter.completeConstraintSystem(context, builtIns.unitType)
 
-        updateCalls(initialStorage, commonSystem)
+        updateCalls(lambda, commonSystem)
 
         return commonSystem.fixedTypeVariables
     }
@@ -88,7 +109,7 @@ class CoroutineInferenceSession(
         *
         * while substitutor from parameter map non-fixed types to the original type variable
         * */
-        val callSubstitutor = storage.buildResultingSubstitutor()
+        val callSubstitutor = storage.buildResultingSubstitutor() // substitutor only for fixed variables
 
         for (initialConstraint in storage.initialConstraints) {
             val lower = nonFixedToVariablesSubstitutor.safeSubstitute(callSubstitutor.safeSubstitute(initialConstraint.a))
@@ -124,26 +145,22 @@ class CoroutineInferenceSession(
         return commonSystem
     }
 
-    private fun updateCalls(initialStorage: ConstraintStorage, commonSystem: NewConstraintSystemImpl) {
+    private fun updateCalls(lambda: ResolvedLambdaAtom, commonSystem: NewConstraintSystemImpl) {
-        val nonFixedToResult = mutableMapOf<TypeConstructor, UnwrappedType>()
+        val nonFixedToVariablesSubstitutor = createNonFixedTypeToVariableSubstitutor()
-        for (variable in initialStorage.postponedTypeVariables) {
+        val commonSystemSubstitutor = commonSystem.buildCurrentSubstitutor()
-            nonFixedToResult[variable.freshTypeConstructor] = commonSystem.fixedTypeVariables[variable.freshTypeConstructor] ?: continue
-        }
 
-        val commonSystemSubstitutor = NewTypeSubstitutorByConstructorMap(nonFixedToResult)
+        val nonFixedTypesToResult = nonFixedToVariablesSubstitutor.map.mapValues { commonSystemSubstitutor.safeSubstitute(it.value) }
+
+        val nonFixedTypesToResultSubstitutor = ComposedSubstitutor(commonSystemSubstitutor, nonFixedToVariablesSubstitutor)
+
+        val lambdaAtomCompleter = createResolvedAtomCompleter(nonFixedTypesToResultSubstitutor, topLevelCallContext)
+        lambdaAtomCompleter.completeAll(lambda)
 
         for (completedCall in suspendCompletedCalls + normalCompletedCalls) {
             val resultCallAtom = completedCall.callResolutionResult.resultCallAtom
             val call = resultCallAtom.atom.getResolvedPsiKotlinCall<CallableDescriptor>(trace) ?: continue
 
-            val resultingCallSubstitutor = completedCall
+            updateCall(completedCall, nonFixedTypesToResultSubstitutor, nonFixedTypesToResult)
-                .callResolutionResult
-                .constraintSystem
-                .fixedTypeVariables
-                .entries
-                .associate { it.key to commonSystemSubstitutor.safeSubstitute(it.value) }
-
-            call.setResultingSubstitutor(NewTypeSubstitutorByConstructorMap(resultingCallSubstitutor))
 
             val resultingDescriptor = call.resultingDescriptor
             kotlinToResolvedCallTransformer.reportCallDiagnostic(
@@ -151,4 +168,44 @@ class CoroutineInferenceSession(
             )
         }
     }
+
+    private fun updateCall(
+        completedCall: PSICompletedCallInfo,
+        nonFixedTypesToResultSubstitutor: NewTypeSubstitutor,
+        nonFixedTypesToResult: Map<TypeConstructor, UnwrappedType>
+    ) {
+        val resultingCallSubstitutor = completedCall.callResolutionResult.constraintSystem.fixedTypeVariables.entries
+            .associate { it.key to nonFixedTypesToResultSubstitutor.safeSubstitute(it.value) }
+
+        val resultingSubstitutor = NewTypeSubstitutorByConstructorMap(resultingCallSubstitutor + nonFixedTypesToResult)
+
+        val atomCompleter = createResolvedAtomCompleter(resultingSubstitutor, completedCall.context)
+        val resultCallAtom = completedCall.callResolutionResult.resultCallAtom
+
+        for (subResolvedAtom in resultCallAtom.subResolvedAtoms) {
+            atomCompleter.completeAll(subResolvedAtom)
+        }
+        atomCompleter.completeResolvedCall(resultCallAtom, completedCall.callResolutionResult.diagnostics)
+
+        val callTrace = completedCall.context.trace
+        if (callTrace is TemporaryBindingTrace) {
+            callTrace.commit()
+        }
+    }
+
+    private fun createResolvedAtomCompleter(
+        resultSubstitutor: NewTypeSubstitutor,
+        context: BasicCallResolutionContext
+    ): ResolvedAtomCompleter {
+        return ResolvedAtomCompleter(
+            resultSubstitutor, context.trace, context, kotlinToResolvedCallTransformer, expressionTypingServices,
+            argumentTypeResolver, doubleColonExpressionResolver, deprecationResolver, moduleDescriptor, context.dataFlowValueFactory
+        )
+    }
+}
+
+class ComposedSubstitutor(val left: NewTypeSubstitutor, val right: NewTypeSubstitutor) : NewTypeSubstitutor {
+    override fun substituteNotNullTypeWithConstructor(constructor: TypeConstructor): UnwrappedType? {
+        return left.safeSubstitute(right.substituteNotNullTypeWithConstructor(constructor) ?: return null)
+    }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt

@@ -11,19 +11,18 @@ import org.jetbrains.kotlin.builtins.KotlinBuiltIns.isUnderKotlinPackage
 import org.jetbrains.kotlin.builtins.createFunctionType
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtPsiUtil
 import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
-import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.*
-import org.jetbrains.kotlin.resolve.BindingTrace
-import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
-import org.jetbrains.kotlin.resolve.TypeResolver
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.components.KotlinResolutionCallbacks
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzer
+import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
 import org.jetbrains.kotlin.resolve.calls.inference.CoroutineInferenceSession
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
@@ -38,6 +37,7 @@ import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValueWithSmartCastInfo
 import org.jetbrains.kotlin.types.*
+import org.jetbrains.kotlin.types.expressions.DoubleColonExpressionResolver
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.KotlinTypeInfo
 import org.jetbrains.kotlin.types.typeUtil.isUnit
@@ -65,7 +65,11 @@ class KotlinResolutionCallbacksImpl(
     val psiCallResolver: PSICallResolver,
     val postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
     val kotlinConstraintSystemCompleter: KotlinConstraintSystemCompleter,
-    val callComponents: KotlinCallComponents
+    val callComponents: KotlinCallComponents,
+    val doubleColonExpressionResolver: DoubleColonExpressionResolver,
+    val deprecationResolver: DeprecationResolver,
+    val moduleDescriptor: ModuleDescriptor,
+    val topLevelCallContext: BasicCallResolutionContext?
 ) : KotlinResolutionCallbacks {
     class LambdaInfo(val expectedType: UnwrappedType, val contextDependency: ContextDependency) {
         val returnStatements = ArrayList<Pair<KtReturnExpression, LambdaContextInfo?>>()
@@ -129,13 +133,18 @@ class KotlinResolutionCallbacksImpl(
             typeApproximator.approximateToSubType(expectedType, TypeApproximatorConfiguration.LocalDeclaration) ?: expectedType
 
         val coroutineSession =
-            if (stubsForPostponedVariables.isNotEmpty())
+            if (stubsForPostponedVariables.isNotEmpty()) {
+                require(topLevelCallContext != null) { "Top level call context should not be null to analyze coroutine-lambda" }
+
                 CoroutineInferenceSession(
                     psiCallResolver, postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter,
-                    callComponents, builtIns, stubsForPostponedVariables, trace, kotlinToResolvedCallTransformer
+                    callComponents, builtIns, topLevelCallContext, stubsForPostponedVariables, trace,
+                    kotlinToResolvedCallTransformer, expressionTypingServices, argumentTypeResolver,
+                    doubleColonExpressionResolver, deprecationResolver, moduleDescriptor
                 )
-            else
+            } else {
                 null
+            }
 
         val actualContext = outerCallContext
             .replaceBindingTrace(trace)

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.resolve.calls.context.CallPosition
 import org.jetbrains.kotlin.resolve.calls.inference.buildResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.FreshVariableNewTypeSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
+import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutorByConstructorMap
 import org.jetbrains.kotlin.resolve.calls.inference.substitute
 import org.jetbrains.kotlin.resolve.calls.inference.substituteAndApproximateCapturedTypes
 import org.jetbrains.kotlin.resolve.calls.model.*
@@ -96,9 +97,18 @@ class KotlinToResolvedCallTransformer(
             is CompletedCallResolutionResult, is ErrorCallResolutionResult -> {
                 val candidate = (baseResolvedCall as SingleCallResolutionResult).resultCallAtom
 
-                context.inferenceSession.addCompletedCallInfo(
+                if (baseResolvedCall is CompletedCallResolutionResult) {
-                    PSICompletedCallInfo(baseResolvedCall as CompletedCallResolutionResult, context, tracingStrategy)
+                    context.inferenceSession.addCompletedCallInfo(PSICompletedCallInfo(baseResolvedCall, context, tracingStrategy))
-                )
+                }
+
+                if (context.inferenceSession.writeOnlyStubs()) {
+                    return createStubResolvedCallAndWriteItToTrace(
+                        candidate,
+                        context.trace,
+                        baseResolvedCall.diagnostics,
+                        completedCall = true
+                    )
+                }
 
                 val resultSubstitutor = baseResolvedCall.constraintSystem.buildResultingSubstitutor()
                 val ktPrimitiveCompleter = ResolvedAtomCompleter(
@@ -121,9 +131,11 @@ class KotlinToResolvedCallTransformer(
     fun <D : CallableDescriptor> createStubResolvedCallAndWriteItToTrace(
         candidate: ResolvedCallAtom,
         trace: BindingTrace,
-        diagnostics: Collection<KotlinCallDiagnostic>
+        diagnostics: Collection<KotlinCallDiagnostic>,
+        completedCall: Boolean = false
     ): ResolvedCall<D> {
-        val result = transformToResolvedCall<D>(candidate, trace, null, diagnostics)
+        val substitutor = if (completedCall) NewTypeSubstitutorByConstructorMap(emptyMap()) else null
+        val result = transformToResolvedCall<D>(candidate, trace, substitutor, diagnostics)
         val psiKotlinCall = candidate.atom.psiKotlinCall
         val tracing = psiKotlinCall.safeAs<PSIKotlinCallForInvoke>()?.baseCall?.tracingStrategy ?: psiKotlinCall.tracingStrategy
 
@@ -499,6 +511,7 @@ class NewResolvedCallImpl<D : CallableDescriptor>(
     private lateinit var typeArguments: List<UnwrappedType>
 
     private var extensionReceiver = resolvedCallAtom.extensionReceiverArgument?.receiver?.receiverValue
+    private var dispatchReceiver = resolvedCallAtom.dispatchReceiverArgument?.receiver?.receiverValue
     private var smartCastDispatchReceiverType: KotlinType? = null
 
     override val kotlinCall: KotlinCall get() = resolvedCallAtom.atom
@@ -511,7 +524,7 @@ class NewResolvedCallImpl<D : CallableDescriptor>(
     override fun getCandidateDescriptor(): D = resolvedCallAtom.candidateDescriptor as D
     override fun getResultingDescriptor(): D = resultingDescriptor
     override fun getExtensionReceiver(): ReceiverValue? = extensionReceiver
-    override fun getDispatchReceiver(): ReceiverValue? = resolvedCallAtom.dispatchReceiverArgument?.receiver?.receiverValue
+    override fun getDispatchReceiver(): ReceiverValue? = dispatchReceiver
     override fun getExplicitReceiverKind(): ExplicitReceiverKind = resolvedCallAtom.explicitReceiverKind
 
     override fun getTypeArguments(): Map<TypeParameterDescriptor, KotlinType> {
@@ -538,6 +551,16 @@ class NewResolvedCallImpl<D : CallableDescriptor>(
         diagnostics = completedDiagnostics
     }
 
+    private fun updateExtensionReceiverType(newType: KotlinType) {
+        if (extensionReceiver?.type == newType) return
+        extensionReceiver = extensionReceiver?.replaceType(newType)
+    }
+
+    private fun updateDispatchReceiverType(newType: KotlinType) {
+        if (dispatchReceiver?.type == newType) return
+        dispatchReceiver = dispatchReceiver?.replaceType(newType)
+    }
+
     fun setResultingSubstitutor(substitutor: NewTypeSubstitutor?) {
         //clear cached values
         argumentToParameterMap = null
@@ -545,6 +568,16 @@ class NewResolvedCallImpl<D : CallableDescriptor>(
         if (substitutor != null) {
             // todo: add asset that we do not complete call many times
             isCompleted = true
+
+            dispatchReceiver?.type?.let {
+                val newType = substitutor.safeSubstitute(it.unwrap())
+                updateDispatchReceiverType(newType)
+            }
+
+            extensionReceiver?.type?.let {
+                val newType = substitutor.safeSubstitute(it.unwrap())
+                updateExtensionReceiverType(newType)
+            }
         }
 
         resultingDescriptor = run {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/NewCallArguments.kt

@@ -59,6 +59,7 @@ val KotlinCallArgument.psiExpression: KtExpression?
         return when (this) {
             is ReceiverExpressionKotlinCallArgument -> receiver.receiverValue.safeAs<ExpressionReceiver>()?.expression
             is QualifierReceiverKotlinCallArgument -> receiver.safeAs<Qualifier>()?.expression
+            is EmptyLabeledReturn -> null // questionable, maybe unit?
             else -> psiCallArgument.valueArgument.getArgumentExpression()
         }
     }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/PSICallResolver.kt

@@ -65,7 +65,9 @@ class PSICallResolver(
     private val constantExpressionEvaluator: ConstantExpressionEvaluator,
     private val dataFlowValueFactory: DataFlowValueFactory,
     private val postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
-    private val kotlinConstraintSystemCompleter: KotlinConstraintSystemCompleter
+    private val kotlinConstraintSystemCompleter: KotlinConstraintSystemCompleter,
+    private val deprecationResolver: DeprecationResolver,
+    private val moduleDescriptor: ModuleDescriptor
 ) {
     private val givenCandidatesName = Name.special("<given candidates>")
 
@@ -158,14 +160,15 @@ class PSICallResolver(
     }
 
     private fun createResolutionCallbacks(context: BasicCallResolutionContext) =
-        createResolutionCallbacks(context.trace, context.inferenceSession)
+        createResolutionCallbacks(context.trace, context.inferenceSession, context)
 
-    fun createResolutionCallbacks(trace: BindingTrace, inferenceSession: InferenceSession) =
+    fun createResolutionCallbacks(trace: BindingTrace, inferenceSession: InferenceSession, context: BasicCallResolutionContext?) =
         KotlinResolutionCallbacksImpl(
             trace, expressionTypingServices, typeApproximator,
             argumentTypeResolver, languageVersionSettings, kotlinToResolvedCallTransformer,
             dataFlowValueFactory, inferenceSession, constantExpressionEvaluator, typeResolver,
-            this, postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter, callComponents
+            this, postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter, callComponents,
+            doubleColonExpressionResolver, deprecationResolver, moduleDescriptor, context
         )
 
     private fun calculateExpectedType(context: BasicCallResolutionContext): UnwrappedType? {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt

@@ -5,12 +5,12 @@
 
 package org.jetbrains.kotlin.resolve.calls.tower
 
-import org.jetbrains.kotlin.builtins.replaceReturnType
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl
+import org.jetbrains.kotlin.descriptors.impl.ReceiverParameterDescriptorImpl
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -88,7 +88,7 @@ class ResolvedAtomCompleter(
     private fun completeLambda(lambda: ResolvedLambdaAtom) {
         val returnType = resultSubstitutor.safeSubstitute(lambda.returnType)
 
-        updateTraceForLambdaReturnType(lambda, trace, returnType)
+        updateTraceForLambda(lambda, trace, returnType)
 
         for (lambdaResult in lambda.resultArguments) {
             val resultValueArgument = lambdaResult as? PSIKotlinCallArgument ?: continue
@@ -102,7 +102,7 @@ class ResolvedAtomCompleter(
         }
     }
 
-    private fun updateTraceForLambdaReturnType(lambda: ResolvedLambdaAtom, trace: BindingTrace, returnType: UnwrappedType) {
+    private fun updateTraceForLambda(lambda: ResolvedLambdaAtom, trace: BindingTrace, returnType: UnwrappedType) {
         val psiCallArgument = lambda.atom.psiCallArgument
 
         val ktArgumentExpression: KtExpression
@@ -124,7 +124,24 @@ class ResolvedAtomCompleter(
         functionDescriptor.setReturnType(returnType)
 
         val existingLambdaType = trace.getType(ktArgumentExpression) ?: throw AssertionError("No type for resolved lambda argument")
-        trace.recordType(ktArgumentExpression, existingLambdaType.replaceReturnType(returnType))
+        trace.recordType(ktArgumentExpression, resultSubstitutor.substituteKeepAnnotations(existingLambdaType.unwrap()))
+
+        // Mainly this is needed for builder-like inference, when we have type `SomeType<K, V>.() -> Unit` and now we want to update those K, V
+        val extensionReceiverParameter = functionDescriptor.extensionReceiverParameter
+        if (extensionReceiverParameter != null) {
+            require(extensionReceiverParameter is ReceiverParameterDescriptorImpl) {
+                "Extension receiver for anonymous function ($extensionReceiverParameter) should be ReceiverParameterDescriptorImpl"
+            }
+
+            val valueType = extensionReceiverParameter.value.type.unwrap()
+            val newValueType = resultSubstitutor.substituteKeepAnnotations(valueType)
+
+            val newReceiverValue = extensionReceiverParameter.value.replaceType(newValueType)
+
+            functionDescriptor.setExtensionReceiverParameter(
+                ReceiverParameterDescriptorImpl(extensionReceiverParameter.containingDeclaration, newReceiverValue)
+            )
+        }
     }
 
     private fun completeCallableReference(

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

@@ -7,10 +7,7 @@ package org.jetbrains.kotlin.resolve.calls.components
 
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.inference.model.NewTypeVariable
-import org.jetbrains.kotlin.resolve.calls.model.CallResolutionResult
+import org.jetbrains.kotlin.resolve.calls.model.*
-import org.jetbrains.kotlin.resolve.calls.model.CompletedCallResolutionResult
-import org.jetbrains.kotlin.resolve.calls.model.KotlinResolutionCandidate
-import org.jetbrains.kotlin.resolve.calls.model.PartialCallResolutionResult
 import org.jetbrains.kotlin.types.TypeConstructor
 import org.jetbrains.kotlin.types.UnwrappedType
 
@@ -22,7 +19,12 @@ interface InferenceSession {
             override fun addErrorCallInfo(callInfo: ErrorCallInfo) {}
             override fun addCompletedCallInfo(callInfo: CompletedCallInfo) {}
             override fun currentConstraintSystem(): ConstraintStorage = ConstraintStorage.Empty
-            override fun inferPostponedVariables(initialStorage: ConstraintStorage): Map<TypeConstructor, UnwrappedType> = emptyMap()
+            override fun inferPostponedVariables(
+                lambda: ResolvedLambdaAtom,
+                initialStorage: ConstraintStorage
+            ): Map<TypeConstructor, UnwrappedType> = emptyMap()
+
+            override fun writeOnlyStubs(): Boolean = false
         }
     }
 
@@ -31,7 +33,8 @@ interface InferenceSession {
     fun addCompletedCallInfo(callInfo: CompletedCallInfo)
     fun addErrorCallInfo(callInfo: ErrorCallInfo)
     fun currentConstraintSystem(): ConstraintStorage
-    fun inferPostponedVariables(initialStorage: ConstraintStorage): Map<TypeConstructor, UnwrappedType>
+    fun inferPostponedVariables(lambda: ResolvedLambdaAtom, initialStorage: ConstraintStorage): Map<TypeConstructor, UnwrappedType>
+    fun writeOnlyStubs(): Boolean
 }
 
 interface PartialCallInfo {

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

@@ -101,10 +101,12 @@ class PostponedArgumentsAnalyzer(
             c.getBuilder().addSubtypeConstraint(lambda.returnType.let(::substitute), unitType, LambdaArgumentConstraintPosition(lambda))
         }
 
+        lambda.setAnalyzedResults(returnArguments, subResolvedKtPrimitives)
+
         if (inferenceSession != null) {
             val storageSnapshot = c.getBuilder().currentStorage()
 
-            val postponedVariables = inferenceSession.inferPostponedVariables(storageSnapshot)
+            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot)
 
             for ((constructor, resultType) in postponedVariables) {
                 val variableWithConstraints = storageSnapshot.notFixedTypeVariables[constructor] ?: continue
@@ -114,7 +116,5 @@ class PostponedArgumentsAnalyzer(
                 c.getBuilder().addEqualityConstraint(variable.defaultType, resultType, CoroutinePosition())
             }
         }
-
-        lambda.setAnalyzedResults(returnArguments, subResolvedKtPrimitives)
     }
 }

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

@@ -195,6 +195,7 @@ internal object CreateFreshVariablesSubstitutor : ResolutionPart() {
 
 internal object PostponedVariablesInitializerResolutionPart : ResolutionPart() {
     override fun KotlinResolutionCandidate.process(workIndex: Int) {
+        // smartset!
         val typesForCoroutineCall = resolvedCall.argumentToCandidateParameter
             .filter { (argument, parameter) -> callComponents.statelessCallbacks.isCoroutineCall(argument, parameter) }
             .mapNotNull { it.value.type.getReceiverTypeFromFunctionType() }

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

@@ -1,5 +1,6 @@
 // !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_ANONYMOUS_PARAMETER -UNUSED_VARIABLE
 // !WITH_NEW_INFERENCE
+// NI_EXPECTED_FILE
 
 class Controller<T> {
     suspend fun yield(t: T) {}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/applyInsideCoroutine.ni.txt

@@ -0,0 +1,34 @@
+package
+
+public val test1: kotlin.Int
+public val test2: A
+public val test3: A
+public fun </*0*/ S> generate(/*0*/ g: suspend Controller<S>.() -> kotlin.Unit): S
+
+public interface A {
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public object B : A {
+    private constructor B()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public object C : A {
+    private constructor C()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Controller</*0*/ T> {
+    public constructor Controller</*0*/ T>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    public final suspend fun yield(/*0*/ t: T): kotlin.Unit
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/applyInsideCoroutine.txt

@@ -1,8 +1,25 @@
 package
 
-public val test1: kotlin.Int
+public val test1: [ERROR : Type for generate {
-public val test2: A
+    apply {
-public val test3: A
+        yield(4)
+    }
+}]
+public val test2: [ERROR : Type for generate {
+    yield(B)
+    apply {
+        yield(C)
+    }
+}]
+public val test3: [ERROR : Type for generate {
+    this.let {
+        yield(B)
+    }
+
+    apply {
+        yield(C)
+    }
+}]
 public fun </*0*/ S> generate(/*0*/ g: suspend Controller<S>.() -> kotlin.Unit): S
 
 public interface A {

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

@@ -1,4 +1,6 @@
 // !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
+// !WITH_NEW_INFERENCE
+// NI_EXPECTED_FILE
 
 class GenericController<T> {
     suspend fun yield(t: T) {}
@@ -14,17 +16,17 @@ val test1 = generate {
     yield(3)
 }
 
-val test2 = <!TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
+val test2 = <!OI;TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
     yield(3)
     notYield(3)
 }
 
-val test3 = <!TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
+val test3 = <!OI;TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
     yield(3)
     yieldBarReturnType(3)
 }
 
-val test4 = <!TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
+val test4 = <!OI;TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
     yield(3)
     barReturnType()
 }

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/incorrectCalls.ni.txt

@@ -0,0 +1,18 @@
+package
+
+public val test1: kotlin.collections.List<kotlin.Int>
+public val test2: kotlin.collections.List<kotlin.Int>
+public val test3: kotlin.collections.List<kotlin.Int>
+public val test4: kotlin.collections.List<kotlin.Int>
+public fun </*0*/ S> generate(/*0*/ g: suspend GenericController<S>.() -> kotlin.Unit): kotlin.collections.List<S>
+
+public final class GenericController</*0*/ T> {
+    public constructor GenericController</*0*/ T>()
+    public final fun barReturnType(): T
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final fun notYield(/*0*/ t: T): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    public final suspend fun yield(/*0*/ t: T): kotlin.Unit
+    public final suspend fun yieldBarReturnType(/*0*/ t: T): T
+}

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

@@ -1,3 +1,4 @@
+// !WITH_NEW_INFERENCE
 // SKIP_TXT
 class StateMachine<Q> internal constructor() {
     fun getInputStub(): Q = null <!UNCHECKED_CAST!>as Q<!>
@@ -12,7 +13,7 @@ class Problem<F>(){
 
     fun createStateMachine(): StateMachine<F> = stateMachine {
         val letter = getInputStub()
-        if (letter is Any)
+        if (<!NI;USELESS_IS_CHECK!>letter is Any<!>)
             println("yes")
     }
 }

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

@@ -1,4 +1,6 @@
 // !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE
+// !WITH_NEW_INFERENCE
+// NI_EXPECTED_FILE
 
 class GenericController<T> {
     suspend fun yield(t: T) {}
@@ -10,7 +12,7 @@ suspend fun <S> GenericController<List<S>>.yieldGenerate(g: suspend GenericContr
 
 val test1 = generate {
     // TODO: KT-15185
-    <!TYPE_MISMATCH, TYPE_INFERENCE_PARAMETER_CONSTRAINT_ERROR!>yieldGenerate<!> {
+    <!OI;TYPE_INFERENCE_PARAMETER_CONSTRAINT_ERROR, OI;TYPE_MISMATCH!>yieldGenerate<!> {
         yield(4)
     }
 }

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/recursiveGenerators2.ni.txt

@@ -0,0 +1,13 @@
+package
+
+public val test1: kotlin.collections.List<kotlin.collections.List<kotlin.Int>>
+public fun </*0*/ S> generate(/*0*/ g: suspend GenericController<S>.() -> kotlin.Unit): kotlin.collections.List<S>
+public suspend fun </*0*/ S> GenericController<kotlin.collections.List<S>>.yieldGenerate(/*0*/ g: suspend GenericController<S>.() -> kotlin.Unit): kotlin.Unit
+
+public final class GenericController</*0*/ T> {
+    public constructor GenericController</*0*/ T>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    public final suspend fun yield(/*0*/ t: T): kotlin.Unit
+}

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

@@ -1,5 +1,6 @@
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // !WITH_NEW_INFERENCE
+// NI_EXPECTED_FILE
 
 class Controller<T> {
     suspend fun yield(t: T) {}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/suspendCallsWithErrors.ni.txt

@@ -0,0 +1,20 @@
+package
+
+public val test1: [ERROR : Error type for ParseError-argument VALUE_ARGUMENT]
+public val test2: kotlin.Int
+public fun </*0*/ S> generate(/*0*/ g: suspend Controller<S>.() -> kotlin.Unit): S
+
+public final class A {
+    public constructor A()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Controller</*0*/ T> {
+    public constructor Controller</*0*/ T>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    public final suspend fun yield(/*0*/ t: T): kotlin.Unit
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/suspendCallsWithErrors.txt

@@ -1,6 +1,8 @@
 package
 
-public val test1: [ERROR : Error type for ParseError-argument VALUE_ARGUMENT]
+public val test1: [ERROR : Type for generate {
+    yield(A)
+}]
 public val test2: kotlin.Int
 public fun </*0*/ S> generate(/*0*/ g: suspend Controller<S>.() -> kotlin.Unit): S
 

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

@@ -1,5 +1,6 @@
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // !WITH_NEW_INFERENCE
+// NI_EXPECTED_FILE
 
 class Controller<T : Number> {
     suspend fun yield(t: T) {}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/suspendCallsWrongUpperBound.ni.txt

@@ -0,0 +1,12 @@
+package
+
+public val test: kotlin.String
+public fun </*0*/ S : kotlin.Number> generate(/*0*/ g: suspend Controller<S>.() -> kotlin.Unit): S
+
+public final class Controller</*0*/ T : kotlin.Number> {
+    public constructor Controller</*0*/ T : kotlin.Number>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    public final suspend fun yield(/*0*/ t: T): kotlin.Unit
+}

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

@@ -1,6 +1,8 @@
 package
 
-public val test: kotlin.String
+public val test: [ERROR : Type for generate {
+    yield("foo")
+}]
 public fun </*0*/ S : kotlin.Number> generate(/*0*/ g: suspend Controller<S>.() -> kotlin.Unit): S
 
 public final class Controller</*0*/ T : kotlin.Number> {

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

@@ -1,4 +1,6 @@
 // !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_ANONYMOUS_PARAMETER -UNUSED_VARIABLE
+// !WITH_NEW_INFERENCE
+// NI_EXPECTED_FILE
 
 class GenericController<T> {
     suspend fun yield(t: T) {}
@@ -6,7 +8,7 @@ class GenericController<T> {
 
 fun <S> generate(g: suspend GenericController<S>.(S) -> Unit): S = TODO()
 
-val test1 = <!TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
+val test1 = <!OI;TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>generate<!> {
     yield(4)
 }
 

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/withUninferredParameter.ni.txt

@@ -0,0 +1,14 @@
+package
+
+public val test1: kotlin.Int
+public val test2: kotlin.Int
+public val test3: kotlin.Int
+public fun </*0*/ S> generate(/*0*/ g: suspend GenericController<S>.(S) -> kotlin.Unit): S
+
+public final class GenericController</*0*/ T> {
+    public constructor GenericController</*0*/ T>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    public final suspend fun yield(/*0*/ t: T): kotlin.Unit
+}

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/FunctionDescriptorImpl.java

@@ -172,6 +172,10 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
         this.hasSynthesizedParameterNames = hasSynthesizedParameterNames;
     }
 
+    public void setExtensionReceiverParameter(@NotNull ReceiverParameterDescriptor extensionReceiverParameter) {
+        this.extensionReceiverParameter = extensionReceiverParameter;
+    }
+
     @Nullable
     @Override
     public ReceiverParameterDescriptor getExtensionReceiverParameter() {

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/ReceiverParameterDescriptorImpl.java

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.descriptors.impl;
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor;
+import org.jetbrains.kotlin.descriptors.ReceiverParameterDescriptor;
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue;
 
 public class ReceiverParameterDescriptorImpl extends AbstractReceiverParameterDescriptor {

core/descriptors/src/org/jetbrains/kotlin/resolve/inlineClassesUtils.kt

@@ -12,6 +12,8 @@ import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import java.lang.reflect.ParameterizedType
+import kotlin.reflect.KFunction1
 
 fun ClassDescriptor.underlyingRepresentation(): ValueParameterDescriptor? {
     if (!isInline) return null