Correct detection of parameter index or receiver from invoke call

idea/src/org/jetbrains/kotlin/idea/slicer/OutflowSlicer.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.slicer
 import com.intellij.codeInsight.highlighting.ReadWriteAccessDetector.Access
 import com.intellij.psi.PsiMethod
 import com.intellij.usageView.UsageInfo
+import org.jetbrains.kotlin.builtins.isExtensionFunctionType
 import org.jetbrains.kotlin.cfg.pseudocode.PseudoValue
 import org.jetbrains.kotlin.cfg.pseudocode.instructions.Instruction
 import org.jetbrains.kotlin.cfg.pseudocode.instructions.eval.*
@@ -180,13 +181,16 @@ class OutflowSlicer(
                                     when (val createdAt = receiverPseudoValue.createdAt) {
                                         is ReadValueInstruction -> {
                                             val accessedDescriptor = createdAt.target.accessedDescriptor ?: return@processPseudocodeUsages
-                                            val accessedDeclaration = accessedDescriptor.originalSource.getPsi() ?: return@processPseudocodeUsages
+                                            if (accessedDescriptor is ValueParameterDescriptor) {
-                                            when (accessedDescriptor) {
+                                                val accessedDeclaration = accessedDescriptor.originalSource.getPsi() ?: return@processPseudocodeUsages
-                                                is ValueParameterDescriptor -> {
+                                                val isExtension = accessedDescriptor.type.isExtensionFunctionType
-                                                    //TODO: argument index is not always correct - first argument can be receiver
+                                                val shift = if (isExtension) 1 else 0
-                                                    val newMode = mode.withBehaviour(LambdaArgumentInflowBehaviour(accessedDescriptor.index))
+                                                val argumentIndex = parameterDescriptor.index - shift
-                                                    accessedDeclaration.passToProcessor(newMode)
+                                                val newMode = if (argumentIndex >= 0)
-                                                }
+                                                    mode.withBehaviour(LambdaArgumentInflowBehaviour(argumentIndex))
+                                                else
+                                                    mode.withBehaviour(LambdaReceiverInflowBehaviour)
+                                                accessedDeclaration.passToProcessor(newMode)
                                             }
                                         }
                                     }

idea/testData/slicer/outflow/invokeExtensionLambda.kt

@@ -0,0 +1,21 @@
+// FLOW: OUT
+
+fun String.foo(<caret>p: String) {
+    val v1 = f(p, { this })
+
+    val v2 = g("a", "b", p, { p1, p2 -> p2 })
+
+    val v3 = inlineF(p, { this })
+}
+
+fun f(receiver: String, lambda: String.() -> String): String {
+    return lambda.invoke(receiver)
+}
+
+fun g(a: String, b: String, c: String, lambda: String.(String, String) -> String): String {
+    return lambda.invoke(a, b, c)
+}
+
+inline fun inlineF(receiver: String, lambda: String.() -> String): String {
+    return lambda.invoke(receiver)
+}

idea/testData/slicer/outflow/invokeExtensionLambda.results.txt

@@ -0,0 +1,43 @@
+3 fun String.foo(<bold>p: String</bold>) {
+4 	val v1 = f(<bold>p</bold>, { this })
+11 		fun f(<bold>receiver: String</bold>, lambda: String.() -> String): String {
+12 			return lambda.invoke(<bold>receiver</bold>)
+11 				[LAMBDA RECEIVER IN] fun f(receiver: String, <bold>lambda: String.() -> String</bold>): String {
+4 					[LAMBDA RECEIVER IN] val v1 = f(p, <bold>{ this }</bold>)
+4 						val v1 = f(p, { <bold>this</bold> })
+4 							val v1 = f(p, <bold>{ this }</bold>)
+4 								[LAMBDA CALLS] val v1 = f(p, <bold>{ this }</bold>)
+11 									[LAMBDA CALLS] fun f(receiver: String, <bold>lambda: String.() -> String</bold>): String {
+12 										[LAMBDA CALLS] return <bold>lambda</bold>.invoke(receiver)
+12 											return lambda.<bold>invoke(receiver)</bold>
+11 												fun <bold>f(receiver: String, lambda: String.() -> String): String {</bold>
+4 													val v1 = <bold>f(p, { this })</bold>
+4 														val <bold>v1 = f(p, { this })</bold>
+6 	val v2 = g("a", "b", <bold>p</bold>, { p1, p2 -> p2 })
+15 		fun g(a: String, b: String, <bold>c: String</bold>, lambda: String.(String, String) -> String): String {
+16 			return lambda.invoke(a, b, <bold>c</bold>)
+15 				[LAMBDA ARGUMENT IN] fun g(a: String, b: String, c: String, <bold>lambda: String.(String, String) -> String</bold>): String {
+6 					[LAMBDA ARGUMENT IN] val v2 = g("a", "b", p, <bold>{ p1, p2 -> p2 }</bold>)
+6 						val v2 = g("a", "b", p, { p1, <bold>p2</bold> -> p2 })
+6 							val v2 = g("a", "b", p, { p1, p2 -> <bold>p2</bold> })
+6 								val v2 = g("a", "b", p, <bold>{ p1, p2 -> p2 }</bold>)
+6 									[LAMBDA CALLS] val v2 = g("a", "b", p, <bold>{ p1, p2 -> p2 }</bold>)
+15 										[LAMBDA CALLS] fun g(a: String, b: String, c: String, <bold>lambda: String.(String, String) -> String</bold>): String {
+16 											[LAMBDA CALLS] return <bold>lambda</bold>.invoke(a, b, c)
+16 												return lambda.<bold>invoke(a, b, c)</bold>
+15 													fun <bold>g(a: String, b: String, c: String, lambda: String.(String, String) -> String): String {</bold>
+6 														val v2 = <bold>g("a", "b", p, { p1, p2 -> p2 })</bold>
+6 															val <bold>v2 = g("a", "b", p, { p1, p2 -> p2 })</bold>
+8 	val v3 = inlineF(<bold>p</bold>, { this })
+19 		(INLINE CALL inlineF) inline fun inlineF(<bold>receiver: String</bold>, lambda: String.() -> String): String {
+20 			(INLINE CALL inlineF) return lambda.invoke(<bold>receiver</bold>)
+19 				(INLINE CALL inlineF) [LAMBDA RECEIVER IN] inline fun inlineF(receiver: String, <bold>lambda: String.() -> String</bold>): String {
+8 					[LAMBDA RECEIVER IN] val v3 = inlineF(p, <bold>{ this }</bold>)
+8 						val v3 = inlineF(p, { <bold>this</bold> })
+8 							val v3 = inlineF(p, <bold>{ this }</bold>)
+8 								[LAMBDA CALLS] val v3 = inlineF(p, <bold>{ this }</bold>)
+19 									(INLINE CALL inlineF) [LAMBDA CALLS] inline fun inlineF(receiver: String, <bold>lambda: String.() -> String</bold>): String {
+20 										(INLINE CALL inlineF) [LAMBDA CALLS] return <bold>lambda</bold>.invoke(receiver)
+20 											(INLINE CALL inlineF) return lambda.<bold>invoke(receiver)</bold>
+8 												val v3 = <bold>inlineF(p, { this })</bold>
+8 													val <bold>v3 = inlineF(p, { this })</bold>

idea/testData/slicer/outflow/invokeLambdaSecondParam.kt

@@ -0,0 +1,9 @@
+// FLOW: OUT
+
+fun String.foo(<caret>p: String) {
+    val v = f({ p1, p2 -> p2 }, p)
+}
+
+fun f(lambda: (String, String) -> String, receiver: String): String {
+    return lambda("a", receiver)
+}

idea/testData/slicer/outflow/invokeLambdaSecondParam.results.txt

@@ -0,0 +1,15 @@
+3 fun String.foo(<bold>p: String</bold>) {
+4 	val v = f({ p1, p2 -> p2 }, <bold>p</bold>)
+7 		fun f(lambda: (String, String) -> String, <bold>receiver: String</bold>): String {
+8 			return lambda("a", <bold>receiver</bold>)
+7 				[LAMBDA ARGUMENT IN] fun f(<bold>lambda: (String, String) -> String</bold>, receiver: String): String {
+4 					[LAMBDA ARGUMENT IN] val v = f(<bold>{ p1, p2 -> p2 }</bold>, p)
+4 						val v = f({ p1, <bold>p2</bold> -> p2 }, p)
+4 							val v = f({ p1, p2 -> <bold>p2</bold> }, p)
+4 								val v = f(<bold>{ p1, p2 -> p2 }</bold>, p)
+4 									[LAMBDA CALLS] val v = f(<bold>{ p1, p2 -> p2 }</bold>, p)
+7 										[LAMBDA CALLS] fun f(<bold>lambda: (String, String) -> String</bold>, receiver: String): String {
+8 											return <bold>lambda("a", receiver)</bold>
+7 												fun <bold>f(lambda: (String, String) -> String, receiver: String): String {</bold>
+4 													val v = <bold>f({ p1, p2 -> p2 }, p)</bold>
+4 														val <bold>v = f({ p1, p2 -> p2 }, p)</bold>

idea/tests/org/jetbrains/kotlin/idea/slicer/SlicerTreeTestGenerated.java

@@ -598,6 +598,16 @@ public class SlicerTreeTestGenerated extends AbstractSlicerTreeTest {
             runTest("idea/testData/slicer/outflow/indexingDereferences.kt");
         }
 
+        @TestMetadata("invokeExtensionLambda.kt")
+        public void testInvokeExtensionLambda() throws Exception {
+            runTest("idea/testData/slicer/outflow/invokeExtensionLambda.kt");
+        }
+
+        @TestMetadata("invokeLambdaSecondParam.kt")
+        public void testInvokeLambdaSecondParam() throws Exception {
+            runTest("idea/testData/slicer/outflow/invokeLambdaSecondParam.kt");
+        }
+
         @TestMetadata("jvmFieldMemberPropertyJavaUsages.kt")
         public void testJvmFieldMemberPropertyJavaUsages() throws Exception {
             runTest("idea/testData/slicer/outflow/jvmFieldMemberPropertyJavaUsages.kt");