From b4b58d1d99bbecd0e3b8df2516359713e18ad204 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Fri, 10 Apr 2020 09:18:24 +0300
Subject: [PATCH] Handling implicit receivers in InflowSlicer

---
 .../kotlin/idea/slicer/InflowSlicer.kt        | 24 ++++++++++++-------
 .../slicer/inflow/onFunctionReceiverType.kt   | 11 +++++++--
 .../onFunctionReceiverType.leafGroups.txt     |  9 +++++--
 .../onFunctionReceiverType.nullnessGroups.txt |  8 +++++--
 .../inflow/onFunctionReceiverType.results.txt |  4 +++-
 5 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/slicer/InflowSlicer.kt b/idea/src/org/jetbrains/kotlin/idea/slicer/InflowSlicer.kt
index f5c36cbbfc3..5ccbc8aa4fb 100644
--- a/idea/src/org/jetbrains/kotlin/idea/slicer/InflowSlicer.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/slicer/InflowSlicer.kt
@@ -45,6 +45,7 @@ import org.jetbrains.kotlin.resolve.calls.model.DefaultValueArgument
 import org.jetbrains.kotlin.resolve.calls.model.ExpressionValueArgument
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
+import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
@@ -164,24 +165,31 @@ class InflowSlicer(
     }
 
     private fun processExtensionReceiver(declaration: KtCallableDeclaration, includeOverriders: Boolean) {
-        fun extractReceiverExpression(refElement: PsiElement): PsiElement? {
-            return when (refElement) {
+        fun processReceiver(refElement: PsiElement) {
+            when (refElement) {
                 is KtExpression -> {
-                    val resolvedCall = refElement.resolveToCall() ?: return null
-                    //TODO: implicit receiver
-                    val expressionReceiver = resolvedCall.extensionReceiver as? ExpressionReceiver ?: return null
-                    return expressionReceiver.expression
+                    val resolvedCall = refElement.resolveToCall() ?: return
+                    when (val receiver = resolvedCall.extensionReceiver) {
+                        is ExpressionReceiver -> {
+                            receiver.expression.passToProcessorAsValue()
+                        }
+
+                        is ImplicitReceiver -> {
+                            val callableDeclaration = (receiver.declarationDescriptor as? CallableDescriptor)?.originalSource?.getPsi()
+                            (callableDeclaration as? KtCallableDeclaration)?.receiverTypeReference?.passToProcessor()
+                        }
+                    }
                 }
                 
                 else -> {
-                    (refElement.parent as? PsiCall)?.argumentList?.expressions?.getOrNull(0)
+                    (refElement.parent as? PsiCall)?.argumentList?.expressions?.getOrNull(0)?.passToProcessorAsValue()
                 }
             }
         }
 
         fun processCall(usageInfo: UsageInfo) {
             usageInfo.element?.let {
-                extractReceiverExpression(it)?.passToProcessorAsValue()
+                processReceiver(it)
             }
         }
 
diff --git a/idea/testData/slicer/inflow/onFunctionReceiverType.kt b/idea/testData/slicer/inflow/onFunctionReceiverType.kt
index 51d86fa4b00..17e71e4ba81 100644
--- a/idea/testData/slicer/inflow/onFunctionReceiverType.kt
+++ b/idea/testData/slicer/inflow/onFunctionReceiverType.kt
@@ -7,7 +7,14 @@ import kotlin.jvm.JvmName
 fun <caret>Any.extensionFun() {
 }
 
-fun foo() {
+fun String.foo() {
     "".extensionFun()
+
     1.extensionFun()
-}
\ No newline at end of file
+
+    extensionFun()
+}
+
+fun main() {
+    "A".foo()
+}
diff --git a/idea/testData/slicer/inflow/onFunctionReceiverType.leafGroups.txt b/idea/testData/slicer/inflow/onFunctionReceiverType.leafGroups.txt
index 7fbeaa3cfb4..adb5234e9da 100644
--- a/idea/testData/slicer/inflow/onFunctionReceiverType.leafGroups.txt
+++ b/idea/testData/slicer/inflow/onFunctionReceiverType.leafGroups.txt
@@ -2,11 +2,16 @@
 7 	fun <bold>Any</bold>.extensionFun() {
 11 		<bold>""</bold>.extensionFun()
 
+19 <bold>"A"</bold>.foo()
+7 	fun <bold>Any</bold>.extensionFun() {
+10 		fun <bold>String</bold>.foo() {
+19 			<bold>"A"</bold>.foo()
+
 3 KotlinUtil.extensionFun(<bold>"x"</bold>);
 7 	fun <bold>Any</bold>.extensionFun() {
 3 		KotlinUtil.extensionFun(<bold>"x"</bold>);
 
-12 <bold>1</bold>.extensionFun()
+13 <bold>1</bold>.extensionFun()
 7 	fun <bold>Any</bold>.extensionFun() {
-12 		<bold>1</bold>.extensionFun()
+13 		<bold>1</bold>.extensionFun()
 
diff --git a/idea/testData/slicer/inflow/onFunctionReceiverType.nullnessGroups.txt b/idea/testData/slicer/inflow/onFunctionReceiverType.nullnessGroups.txt
index a6c683edd00..ce34565c438 100644
--- a/idea/testData/slicer/inflow/onFunctionReceiverType.nullnessGroups.txt
+++ b/idea/testData/slicer/inflow/onFunctionReceiverType.nullnessGroups.txt
@@ -2,9 +2,13 @@
 11 	<bold>""</bold>.extensionFun()
 7 		fun <bold>Any</bold>.extensionFun() {
 11 			<bold>""</bold>.extensionFun()
-12 	<bold>1</bold>.extensionFun()
+13 	<bold>1</bold>.extensionFun()
 7 		fun <bold>Any</bold>.extensionFun() {
-12 			<bold>1</bold>.extensionFun()
+13 			<bold>1</bold>.extensionFun()
+19 	<bold>"A"</bold>.foo()
+7 		fun <bold>Any</bold>.extensionFun() {
+10 			fun <bold>String</bold>.foo() {
+19 				<bold>"A"</bold>.foo()
 
 [Other Values]
 3 	KotlinUtil.extensionFun(<bold>"x"</bold>);
diff --git a/idea/testData/slicer/inflow/onFunctionReceiverType.results.txt b/idea/testData/slicer/inflow/onFunctionReceiverType.results.txt
index d2614d5f30a..41a5b251139 100644
--- a/idea/testData/slicer/inflow/onFunctionReceiverType.results.txt
+++ b/idea/testData/slicer/inflow/onFunctionReceiverType.results.txt
@@ -1,4 +1,6 @@
 7 fun <bold>Any</bold>.extensionFun() {
 3 	KotlinUtil.extensionFun(<bold>"x"</bold>);
+10 	fun <bold>String</bold>.foo() {
+19 		<bold>"A"</bold>.foo()
 11 	<bold>""</bold>.extensionFun()
-12 	<bold>1</bold>.extensionFun()
+13 	<bold>1</bold>.extensionFun()