[NI] Update calls after inference for coroutines

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.CompletedCallResolutionResult
-import org.jetbrains.kotlin.resolve.calls.model.KotlinResolutionCandidate
-import org.jetbrains.kotlin.resolve.calls.model.PartialCallResolutionResult
+import org.jetbrains.kotlin.resolve.calls.model.*
 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() }