Handling implicit receivers in InflowSlicer

JetBrains · Apr 17, 2020 · b4b58d1 · b4b58d1
1 parent 90188ef
commit b4b58d1
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 15 deletions.
diff --git 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
@@ -7,7 +7,14 @@ import kotlin.jvm.JvmName
 fun <caret>Any.extensionFun() {
 }
 
-fun foo() {
+fun String.foo() {
     "".extensionFun()
+
     1.extensionFun()
-}
+
+    extensionFun()
+}
+
+fun main() {
+    "A".foo()
+}
diff --git 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
@@ -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
@@ -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()