Introduce original for ReceiverValue and use it inside coroutine checker

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

@@ -157,19 +157,9 @@ private fun checkRestrictsSuspension(
     if (!enclosingSuspendReceiverValue.isRestrictsSuspensionReceiver()) return
 
     // member of suspend receiver
-    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.dispatchReceiver?.original) return
 
-        resolvedCall.originalDispatchReceiver() to resolvedCall.originalExtensionReceiver()
-    } else {
-        resolvedCall.dispatchReceiver to resolvedCall.extensionReceiver
-    }
-
-    if (enclosingSuspendReceiverValue sameInstance dispatchReceiver) return
-
-    if (enclosingSuspendReceiverValue sameInstance extensionReceiver &&
+    if (enclosingSuspendReceiverValue sameInstance resolvedCall.extensionReceiver?.original &&
         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/tower/KotlinToResolvedCallTransformer.kt

@@ -534,9 +534,6 @@ 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(

compiler/frontend/src/org/jetbrains/kotlin/resolve/scopes/receivers/ExpressionReceiver.kt

@@ -26,9 +26,9 @@ interface ExpressionReceiver : ReceiverValue {
 
     companion object {
         private open class ExpressionReceiverImpl(
-            override val expression: KtExpression, type: KotlinType
-        ) : AbstractReceiverValue(type), ExpressionReceiver {
-            override fun replaceType(newType: KotlinType) = ExpressionReceiverImpl(expression, newType)
+            override val expression: KtExpression, type: KotlinType, original: ReceiverValue?
+        ) : AbstractReceiverValue(type, original), ExpressionReceiver {
+            override fun replaceType(newType: KotlinType) = ExpressionReceiverImpl(expression, newType, original)
 
             override fun toString() = "$type {$expression: ${expression.text}}"
         }
@@ -36,17 +36,19 @@ interface ExpressionReceiver : ReceiverValue {
         private class ThisExpressionClassReceiver(
             override val classDescriptor: ClassDescriptor,
             expression: KtExpression,
-            type: KotlinType
-        ) : ExpressionReceiverImpl(expression, type), ThisClassReceiver {
-            override fun replaceType(newType: KotlinType) = ThisExpressionClassReceiver(classDescriptor, expression, newType)
+            type: KotlinType,
+            original: ReceiverValue?
+        ) : ExpressionReceiverImpl(expression, type, original), ThisClassReceiver {
+            override fun replaceType(newType: KotlinType) = ThisExpressionClassReceiver(classDescriptor, expression, newType, original)
         }
 
         private class SuperExpressionReceiver(
             override val thisType: KotlinType,
             expression: KtExpression,
-            type: KotlinType
-        ) : ExpressionReceiverImpl(expression, type), SuperCallReceiverValue {
-            override fun replaceType(newType: KotlinType) = SuperExpressionReceiver(thisType, expression, newType)
+            type: KotlinType,
+            original: ReceiverValue?
+        ) : ExpressionReceiverImpl(expression, type, original), SuperCallReceiverValue {
+            override fun replaceType(newType: KotlinType) = SuperExpressionReceiver(thisType, expression, newType, original)
         }
 
         fun create(
@@ -64,17 +66,17 @@ interface ExpressionReceiver : ReceiverValue {
             if (referenceExpression != null) {
                 val descriptor = bindingContext.get(BindingContext.REFERENCE_TARGET, referenceExpression)
                 if (descriptor is ClassDescriptor) {
-                    return ThisExpressionClassReceiver(descriptor.original, expression, type)
+                    return ThisExpressionClassReceiver(descriptor.original, expression, type, original = null)
                 }
             } else if (expression is KtSuperExpression) {
                 // if there is no THIS_TYPE_FOR_SUPER_EXPRESSION in binding context, we fall through into more restrictive option
                 // i.e. just return common ExpressionReceiverImpl
                 bindingContext[BindingContext.THIS_TYPE_FOR_SUPER_EXPRESSION, expression]?.let { thisType ->
-                    return SuperExpressionReceiver(thisType, expression, type)
+                    return SuperExpressionReceiver(thisType, expression, type, original = null)
                 }
             }
 
-            return ExpressionReceiverImpl(expression, type)
+            return ExpressionReceiverImpl(expression, type, original = null)
         }
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/scopes/receivers/Qualifier.kt

@@ -123,11 +123,19 @@ class TypeAliasQualifier(
     }
 }
 
-class ClassValueReceiver(val classQualifier: ClassifierQualifier, private val type: KotlinType) : ExpressionReceiver {
+class ClassValueReceiver @JvmOverloads constructor(
+    val classQualifier: ClassifierQualifier,
+    private val type: KotlinType,
+    original: ClassValueReceiver? = null
+) : ExpressionReceiver {
+    private val original = original ?: this
+
     override fun getType() = type
 
     override val expression: KtExpression
         get() = classQualifier.expression
 
-    override fun replaceType(newType: KotlinType) = ClassValueReceiver(classQualifier, newType)
+    override fun replaceType(newType: KotlinType) = ClassValueReceiver(classQualifier, newType, original)
+
+    override fun getOriginal() = original
 }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/smartcasts/smartCastUtil.kt

@@ -22,9 +22,12 @@ import org.jetbrains.kotlin.types.KotlinType
 fun getReceiverValueWithSmartCast(
     receiverArgument: ReceiverValue?,
     smartCastType: KotlinType?
-) = smartCastType?.let(::SmartCastReceiverValue) ?: receiverArgument
+) = smartCastType?.let { type -> SmartCastReceiverValue(type, original = null) } ?: receiverArgument
+
+private class SmartCastReceiverValue(private val type: KotlinType, original: SmartCastReceiverValue?) : ReceiverValue {
+    private val original = original ?: this
 
-private class SmartCastReceiverValue(private val type: KotlinType) : ReceiverValue {
     override fun getType() = type
-    override fun replaceType(newType: KotlinType) = SmartCastReceiverValue(newType)
+    override fun replaceType(newType: KotlinType) = SmartCastReceiverValue(newType, original)
+    override fun getOriginal() = original
 }