From 8f80851b9ad1249678ffcb72a9722b3b2f438260 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Fri, 24 Aug 2018 19:42:55 +0900
Subject: [PATCH] if then to elvis: support calling extension function in class
 #KT-26343 Fixed

---
 .../idea/intentions/branchedTransformations/IfThenUtils.kt | 7 ++++++-
 .../branched/ifThenToElvis/extensionFunctionInClass.kt     | 5 +++++
 .../ifThenToElvis/extensionFunctionInClass.kt.after        | 5 +++++
 .../kotlin/idea/intentions/IntentionTestGenerated.java     | 5 +++++
 4 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/IfThenUtils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/IfThenUtils.kt
index ac3290ce80e..c08d718cb43 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/IfThenUtils.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/IfThenUtils.kt
@@ -43,6 +43,7 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.callUtil.getType
+import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getExplicitReceiverValue
 import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getImplicitReceiverValue
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
@@ -229,7 +230,11 @@ data class IfThenToSelectData(
         }
     }
 
-    internal fun getImplicitReceiver(): ImplicitReceiver? = baseClause.getResolvedCall(context)?.getImplicitReceiverValue()
+    internal fun getImplicitReceiver(): ImplicitReceiver? {
+        val resolvedCall = baseClause.getResolvedCall(context) ?: return null
+        if (resolvedCall.getExplicitReceiverValue() != null) return null
+        return resolvedCall.getImplicitReceiverValue()
+    }
 
     internal fun hasImplicitReceiver(): Boolean = getImplicitReceiver() != null
 
diff --git a/idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt b/idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt
new file mode 100644
index 00000000000..4ce7869cbe0
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt
@@ -0,0 +1,5 @@
+class Foo {
+    fun Foo.bar(): Int = 1
+}
+
+fun Foo.test(foo: Foo?): Int = <caret>if (foo == null) { 0 } else { foo.bar() }
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt.after b/idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt.after
new file mode 100644
index 00000000000..e3842c7b503
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt.after
@@ -0,0 +1,5 @@
+class Foo {
+    fun Foo.bar(): Int = 1
+}
+
+fun Foo.test(foo: Foo?): Int = foo?.bar() ?: 0
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 209bc416121..b78b6e80de5 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -2276,6 +2276,11 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 runTest("idea/testData/intentions/branched/ifThenToElvis/emptyThenBlock.kt");
             }
 
+            @TestMetadata("extensionFunctionInClass.kt")
+            public void testExtensionFunctionInClass() throws Exception {
+                runTest("idea/testData/intentions/branched/ifThenToElvis/extensionFunctionInClass.kt");
+            }
+
             @TestMetadata("ifAndElseBothInBlocks.kt")
             public void testIfAndElseBothInBlocks() throws Exception {
                 runTest("idea/testData/intentions/branched/ifThenToElvis/ifAndElseBothInBlocks.kt");