[NI] Fix checking for the same instance inside for suspend functions

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/coroutineCallChecker.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.psi.KtThisExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.tower.NewResolvedCallImpl
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.inline.InlineUtil
 import org.jetbrains.kotlin.resolve.scopes.HierarchicalScope
@@ -156,9 +157,19 @@ private fun checkRestrictsSuspension(
     if (!enclosingSuspendReceiverValue.isRestrictsSuspensionReceiver()) return
 
     // member of suspend receiver
-    if (enclosingSuspendReceiverValue sameInstance resolvedCall.dispatchReceiver) return
+    val (dispatchReceiver, extensionReceiver) = if (context.languageVersionSettings.supportsFeature(LanguageFeature.NewInference)) {
+        require(resolvedCall is NewResolvedCallImpl<*>) {
+            "Resolved call with enabled new inference should be instance of NewResolvedCallImpl"
+        }
 
-    if (enclosingSuspendReceiverValue sameInstance resolvedCall.extensionReceiver &&
+        resolvedCall.originalDispatchReceiver() to resolvedCall.originalExtensionReceiver()
+    } else {
+        resolvedCall.dispatchReceiver to resolvedCall.extensionReceiver
+    }
+
+    if (enclosingSuspendReceiverValue sameInstance dispatchReceiver) return
+
+    if (enclosingSuspendReceiverValue sameInstance extensionReceiver &&
         resolvedCall.candidateDescriptor.extensionReceiverParameter!!.value.isRestrictsSuspensionReceiver()) return
 
     context.trace.report(Errors.ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL.on(reportOn))

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

@@ -153,9 +153,6 @@ class CoroutineInferenceSession(
 
         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
@@ -167,6 +164,9 @@ class CoroutineInferenceSession(
                 completedCall.context, trace, resultCallAtom, resultingDescriptor, commonSystem.diagnostics
             )
         }
+
+        val lambdaAtomCompleter = createResolvedAtomCompleter(nonFixedTypesToResultSubstitutor, topLevelCallContext)
+        lambdaAtomCompleter.completeAll(lambda)
     }
 
     private fun updateCall(

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

@@ -534,6 +534,9 @@ class NewResolvedCallImpl<D : CallableDescriptor>(
 
     override fun getSmartCastDispatchReceiverType(): KotlinType? = smartCastDispatchReceiverType
 
+    internal fun originalExtensionReceiver(): ReceiverValue? = resolvedCallAtom.extensionReceiverArgument?.receiver?.receiverValue
+    internal fun originalDispatchReceiver(): ReceiverValue? = resolvedCallAtom.dispatchReceiverArgument?.receiver?.receiverValue
+
     fun updateExtensionReceiverWithSmartCastIfNeeded(smartCastExtensionReceiverType: KotlinType) {
         if (extensionReceiver is ImplicitClassReceiver) {
             extensionReceiver = CastImplicitClassReceiver(