From bfd20e18908818dd8500d7abf2053f266ff84383 Mon Sep 17 00:00:00 2001
From: Anastasiya Shadrina <a.shadrina5@mail.ru>
Date: Thu, 10 Jun 2021 02:26:48 +0700
Subject: [PATCH] [FE] Move context receivers checks to the separate resolution
 part

---
 .../calls/components/ResolutionParts.kt       | 92 +++++++++++--------
 .../contextReceivers/typeParameterized.kt     |  2 +-
 2 files changed, 55 insertions(+), 39 deletions(-)

diff --git a/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt b/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt
index 735dfc691d6..ace70a429c6 100644
--- a/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt
+++ b/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt
@@ -697,47 +697,10 @@ internal object CheckReceivers : ResolutionPart() {
                     shouldCheckImplicitInvoke = false, // reproduce old inference behaviour
                 )
             }
-            else -> {
-                resolvedCall.contextReceiversArguments = searchForAdditionalReceivers() ?: return
-                for (i in resolvedCall.contextReceiversArguments.indices) {
-                    checkReceiver(
-                        resolvedCall.contextReceiversArguments[i],
-                        candidateDescriptor.contextReceiverParameters[i],
-                        shouldCheckImplicitInvoke = false
-                    )
-                }
-            }
         }
     }
 
-    override fun ResolutionCandidate.workCount() = 3
-
-    private fun ResolutionCandidate.searchForAdditionalReceivers(): List<SimpleKotlinCallArgument>? {
-        val result = mutableListOf<ReceiverValueWithSmartCastInfo>()
-        val candidateReceivers = scopeTower.lexicalScope.parentsWithSelf
-            .flatMap { if (it is LexicalScope) scopeTower.getImplicitReceivers(it) else emptyList() }
-
-        fun KotlinType.prepared(): KotlinType = if (containsTypeParameter()) replaceArgumentsWithStarProjections() else this
-
-        for (receiver in candidateDescriptor.contextReceiverParameters) {
-            val expectedReceiverType = receiver.type
-            val expectedReceiverTypeClosestBound =
-                if (expectedReceiverType.isTypeParameter()) expectedReceiverType.supertypes().first()
-                else expectedReceiverType
-            val selectedCandidate = candidateReceivers.firstOrNull { candidateReceiver ->
-                val candidateReceiverType = candidateReceiver.receiverValue.type
-                org.jetbrains.kotlin.types.checker.NewKotlinTypeChecker.Default.isSubtypeOf(
-                    candidateReceiverType.prepared(),
-                    expectedReceiverTypeClosestBound.prepared()
-                )
-            } ?: run {
-                this.diagnosticsFromResolutionParts.add(NoContextReceiver(receiver))
-                return null
-            }
-            result.add(selectedCandidate)
-        }
-        return result.map { ReceiverExpressionKotlinCallArgument(it) }
-    }
+    override fun ResolutionCandidate.workCount() = 2
 }
 
 internal object CheckArgumentsInParenthesis : ResolutionPart() {
@@ -841,3 +804,56 @@ internal object ErrorDescriptorResolutionPart : ResolutionPart() {
         }
     }
 }
+
+internal object CheckContextReceiversResolutionPart : ResolutionPart() {
+    private fun ResolutionCandidate.checkReceiver(
+        receiverArgument: SimpleKotlinCallArgument?,
+        receiverParameter: ReceiverParameterDescriptor?
+    ) {
+        if ((receiverArgument == null) != (receiverParameter == null)) {
+            error("Inconsistency receiver state for call $kotlinCall and candidate descriptor: $candidateDescriptor")
+        }
+        if (receiverArgument == null || receiverParameter == null) return
+        val receiverInfo = ReceiverInfo(
+            isReceiver = true,
+            shouldReportUnsafeCall = true,
+            reportUnsafeCallAsUnsafeImplicitInvoke = false
+        )
+
+        resolveKotlinArgument(receiverArgument, receiverParameter, receiverInfo)
+    }
+
+    private fun ResolutionCandidate.searchForAdditionalReceivers(): List<SimpleKotlinCallArgument>? {
+        val result = mutableListOf<ReceiverValueWithSmartCastInfo>()
+        val candidateReceivers = scopeTower.lexicalScope.parentsWithSelf
+            .flatMap { if (it is LexicalScope) scopeTower.getImplicitReceivers(it) else emptyList() }
+
+        fun KotlinType.prepared(): KotlinType = if (containsTypeParameter()) replaceArgumentsWithStarProjections() else this
+
+        for (receiver in candidateDescriptor.contextReceiverParameters) {
+            val expectedReceiverType = receiver.type
+            val expectedReceiverTypeClosestBound =
+                if (expectedReceiverType.isTypeParameter()) expectedReceiverType.supertypes().first()
+                else expectedReceiverType
+            val selectedCandidate = candidateReceivers.firstOrNull { candidateReceiver ->
+                val candidateReceiverType = candidateReceiver.receiverValue.type
+                NewKotlinTypeChecker.Default.isSubtypeOf(candidateReceiverType.prepared(), expectedReceiverTypeClosestBound.prepared())
+            } ?: run {
+                this.diagnosticsFromResolutionParts.add(NoContextReceiver(receiver))
+                return null
+            }
+            result.add(selectedCandidate)
+        }
+        return result.map { ReceiverExpressionKotlinCallArgument(it) }
+    }
+
+    override fun ResolutionCandidate.process(workIndex: Int) {
+        resolvedCall.contextReceiversArguments = searchForAdditionalReceivers() ?: return
+        for (i in resolvedCall.contextReceiversArguments.indices) {
+            checkReceiver(
+                resolvedCall.contextReceiversArguments[i],
+                candidateDescriptor.contextReceiverParameters[i]
+            )
+        }
+    }
+}
\ No newline at end of file
diff --git a/compiler/testData/diagnostics/tests/extensions/contextReceivers/typeParameterized.kt b/compiler/testData/diagnostics/tests/extensions/contextReceivers/typeParameterized.kt
index 430bbbf4e44..e3a5545c9cd 100644
--- a/compiler/testData/diagnostics/tests/extensions/contextReceivers/typeParameterized.kt
+++ b/compiler/testData/diagnostics/tests/extensions/contextReceivers/typeParameterized.kt
@@ -7,5 +7,5 @@ context(T)
 fun <T> T.f(t: B<T>) {}
 
 fun Int.main(a: A, b: B<String>) {
-    a.f(<!TYPE_MISMATCH, TYPE_MISMATCH!>b<!>)
+    a.<!UNRESOLVED_REFERENCE_WRONG_RECEIVER!>f<!>(b)
 }
\ No newline at end of file