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)) {
+    if (enclosingSuspendReceiverValue sameInstance resolvedCall.dispatchReceiver?.original) return
-        require(resolvedCall is NewResolvedCallImpl<*>) {
-            "Resolved call with enabled new inference should be instance of NewResolvedCallImpl"
-        }
 
-        resolvedCall.originalDispatchReceiver() to resolvedCall.originalExtensionReceiver()
+    if (enclosingSuspendReceiverValue sameInstance resolvedCall.extensionReceiver?.original &&
-    } 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/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
+            override val expression: KtExpression, type: KotlinType, original: ReceiverValue?
-        ) : AbstractReceiverValue(type), ExpressionReceiver {
+        ) : AbstractReceiverValue(type, original), ExpressionReceiver {
-            override fun replaceType(newType: KotlinType) = ExpressionReceiverImpl(expression, newType)
+            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
+            type: KotlinType,
-        ) : ExpressionReceiverImpl(expression, type), ThisClassReceiver {
+            original: ReceiverValue?
-            override fun replaceType(newType: KotlinType) = ThisExpressionClassReceiver(classDescriptor, expression, newType)
+        ) : 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
+            type: KotlinType,
-        ) : ExpressionReceiverImpl(expression, type), SuperCallReceiverValue {
+            original: ReceiverValue?
-            override fun replaceType(newType: KotlinType) = SuperExpressionReceiver(thisType, expression, newType)
+        ) : 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
 }

core/descriptors/src/org/jetbrains/kotlin/descriptors/Visibilities.java

@@ -402,6 +402,12 @@ public class Visibilities {
         public ReceiverValue replaceType(@NotNull KotlinType newType) {
             throw new IllegalStateException("This method should not be called");
         }
+
+        @NotNull
+        @Override
+        public ReceiverValue getOriginal() {
+            return this;
+        }
     };
 
     /**
@@ -420,6 +426,12 @@ public class Visibilities {
         public ReceiverValue replaceType(@NotNull KotlinType newType) {
             throw new IllegalStateException("This method should not be called");
         }
+
+        @NotNull
+        @Override
+        public ReceiverValue getOriginal() {
+            return this;
+        }
     };
 
     // This constant is not intended to use somewhere else from
@@ -436,6 +448,12 @@ public class Visibilities {
         public ReceiverValue replaceType(@NotNull KotlinType newType) {
             throw new IllegalStateException("This method should not be called");
         }
+
+        @NotNull
+        @Override
+        public ReceiverValue getOriginal() {
+            return this;
+        }
     };
 
     public static boolean isPrivate(@NotNull Visibility visibility) {

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

@@ -28,7 +28,7 @@ public class LazyClassReceiverParameterDescriptor extends AbstractReceiverParame
 
     public LazyClassReceiverParameterDescriptor(@NotNull ClassDescriptor descriptor) {
         this.descriptor = descriptor;
-        this.receiverValue = new ImplicitClassReceiver(descriptor);
+        this.receiverValue = new ImplicitClassReceiver(descriptor, null);
     }
 
     @NotNull

core/descriptors/src/org/jetbrains/kotlin/resolve/DescriptorFactory.java

@@ -159,6 +159,6 @@ public class DescriptorFactory {
     ) {
         return receiverParameterType == null
                ? null
-               : new ReceiverParameterDescriptorImpl(owner, new ExtensionReceiver(owner, receiverParameterType));
+               : new ReceiverParameterDescriptorImpl(owner, new ExtensionReceiver(owner, receiverParameterType, null));
     }
 }

core/descriptors/src/org/jetbrains/kotlin/resolve/scopes/receivers/AbstractReceiverValue.java

@@ -17,13 +17,16 @@
 package org.jetbrains.kotlin.resolve.scopes.receivers;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.types.KotlinType;
 
 public abstract class AbstractReceiverValue implements ReceiverValue {
     protected final KotlinType receiverType;
+    private final ReceiverValue original;
 
-    public AbstractReceiverValue(@NotNull KotlinType receiverType) {
+    public AbstractReceiverValue(@NotNull KotlinType receiverType, @Nullable ReceiverValue original) {
         this.receiverType = receiverType;
+        this.original = original != null ? original : this;
     }
 
     @Override
@@ -31,4 +34,10 @@ public abstract class AbstractReceiverValue implements ReceiverValue {
     public KotlinType getType() {
         return receiverType;
     }
+
+    @NotNull
+    @Override
+    public ReceiverValue getOriginal() {
+        return original;
+    }
 }

core/descriptors/src/org/jetbrains/kotlin/resolve/scopes/receivers/ExtensionReceiver.java

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.resolve.scopes.receivers;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.CallableDescriptor;
 import org.jetbrains.kotlin.types.KotlinType;
 
@@ -24,8 +25,12 @@ public class ExtensionReceiver extends AbstractReceiverValue implements Implicit
 
     private final CallableDescriptor descriptor;
 
-    public ExtensionReceiver(@NotNull CallableDescriptor callableDescriptor, @NotNull KotlinType receiverType) {
+    public ExtensionReceiver(
-        super(receiverType);
+            @NotNull CallableDescriptor callableDescriptor,
+            @NotNull KotlinType receiverType,
+            @Nullable ReceiverValue original
+    ) {
+        super(receiverType, original);
         this.descriptor = callableDescriptor;
     }
 
@@ -38,7 +43,7 @@ public class ExtensionReceiver extends AbstractReceiverValue implements Implicit
     @NotNull
     @Override
     public ReceiverValue replaceType(@NotNull KotlinType newType) {
-        return new ExtensionReceiver(descriptor, newType);
+        return new ExtensionReceiver(descriptor, newType, getOriginal());
     }
 
     @Override

core/descriptors/src/org/jetbrains/kotlin/resolve/scopes/receivers/ImplicitClassReceiver.kt

@@ -30,7 +30,12 @@ interface ThisClassReceiver : ReceiverValue {
 /**
  * Same but implicit only
  */
-open class ImplicitClassReceiver(final override val classDescriptor: ClassDescriptor) : ThisClassReceiver, ImplicitReceiver {
+open class ImplicitClassReceiver(
+    final override val classDescriptor: ClassDescriptor,
+    original: ImplicitClassReceiver? = null
+) : ThisClassReceiver, ImplicitReceiver {
+
+    private val original = original ?: this
 
     override fun getType() = classDescriptor.defaultType
 
@@ -44,4 +49,6 @@ open class ImplicitClassReceiver(final override val classDescriptor: ClassDescri
 
     override fun replaceType(newType: KotlinType) =
             throw UnsupportedOperationException("Replace type should not be called for this receiver")
+
+    override fun getOriginal() = original
 }

core/descriptors/src/org/jetbrains/kotlin/resolve/scopes/receivers/ReceiverValue.java

@@ -26,4 +26,7 @@ public interface ReceiverValue extends Receiver {
 
     @NotNull
     ReceiverValue replaceType(@NotNull KotlinType newType);
+
+    @NotNull
+    ReceiverValue getOriginal();
 }

core/descriptors/src/org/jetbrains/kotlin/resolve/scopes/receivers/TransientReceiver.java

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.resolve.scopes.receivers;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.types.KotlinType;
 
 /**
@@ -25,7 +26,11 @@ import org.jetbrains.kotlin.types.KotlinType;
  */
 public class TransientReceiver extends AbstractReceiverValue {
     public TransientReceiver(@NotNull KotlinType type) {
-        super(type);
+        this(type, null);
+    }
+
+    private TransientReceiver(@NotNull KotlinType type, @Nullable ReceiverValue original) {
+        super(type, original);
     }
 
     @Override
@@ -36,6 +41,6 @@ public class TransientReceiver extends AbstractReceiverValue {
     @NotNull
     @Override
     public ReceiverValue replaceType(@NotNull KotlinType newType) {
-        return new TransientReceiver(newType);
+        return new TransientReceiver(newType, getOriginal());
     }
 }