From 02d309c9cffc4d7c104764402e4bb7d4b07cc756 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Wed, 23 Mar 2016 14:28:28 +0300
Subject: [PATCH] KT-11481 "Add import" intention is not available for 'is'
 branches in when

 #KT-11481 Fixed
---
 .../idea/intentions/ImportMemberIntention.kt  | 27 +++++++++++--------
 .../intentions/importMember/NestedClass1.kt   | 15 +++++++++++
 .../importMember/NestedClass1.kt.after        | 17 ++++++++++++
 .../intentions/importMember/NestedClass2.kt   | 15 +++++++++++
 .../importMember/NestedClass2.kt.after        | 17 ++++++++++++
 .../intentions/IntentionTestGenerated.java    | 12 +++++++++
 6 files changed, 92 insertions(+), 11 deletions(-)
 create mode 100644 idea/testData/intentions/importMember/NestedClass1.kt
 create mode 100644 idea/testData/intentions/importMember/NestedClass1.kt.after
 create mode 100644 idea/testData/intentions/importMember/NestedClass2.kt
 create mode 100644 idea/testData/intentions/importMember/NestedClass2.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ImportMemberIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ImportMemberIntention.kt
index 229ee7f4cb7..cc2b8ed72cf 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ImportMemberIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ImportMemberIntention.kt
@@ -26,10 +26,13 @@ import org.jetbrains.kotlin.idea.util.ImportInsertHelper
 import org.jetbrains.kotlin.idea.util.ShortenReferences
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
+import org.jetbrains.kotlin.psi.KtUserType
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElement
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
+import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
@@ -38,9 +41,10 @@ class ImportMemberIntention : SelfTargetingOffsetIndependentIntention<KtNameRefe
         "Add import for member"
 ){
     override fun isApplicableTo(element: KtNameReferenceExpression): Boolean {
-        val qualifiedExpression = qualifiedExpression(element) ?: return false
+        val qualifiedElement = element.getQualifiedElement()
+        if (qualifiedElement == element) return false
 
-        val fqName = targetFqName(qualifiedExpression) ?: return false
+        val fqName = targetFqName(qualifiedElement) ?: return false
 
         text = "Add import for '${fqName.asString()}'"
         return true
@@ -59,19 +63,20 @@ class ImportMemberIntention : SelfTargetingOffsetIndependentIntention<KtNameRefe
             val selector = qualifiedExpression.getQualifiedElementSelector() as? KtNameReferenceExpression
             selector?.getReferencedNameAsName() == fqName.shortName() && targetFqName(qualifiedExpression) == fqName
         }
+        val userTypes = file.collectDescendantsOfType<KtUserType> { userType ->
+            val selector = userType.getQualifiedElementSelector() as? KtNameReferenceExpression
+            selector?.getReferencedNameAsName() == fqName.shortName() && targetFqName(userType) == fqName
+        }
 
         //TODO: not deep
-        ShortenReferences.DEFAULT.process(qualifiedExpressions)
+        ShortenReferences.DEFAULT.process(qualifiedExpressions + userTypes)
     }
 
-    private fun qualifiedExpression(element: KtNameReferenceExpression): KtDotQualifiedExpression? {
-        return element.getQualifiedElement() as? KtDotQualifiedExpression
-    }
-
-    private fun targetFqName(qualifiedExpression: KtDotQualifiedExpression): FqName? {
-        val nameExpression = qualifiedExpression.getQualifiedElementSelector() as? KtNameReferenceExpression ?: return null
-        val bindingContext = qualifiedExpression.analyze(BodyResolveMode.PARTIAL)
-        if (bindingContext[BindingContext.QUALIFIER, qualifiedExpression.receiverExpression] == null) return null
+    private fun targetFqName(qualifiedElement: KtElement): FqName? {
+        val nameExpression = qualifiedElement.getQualifiedElementSelector() as? KtNameReferenceExpression ?: return null
+        val receiver = nameExpression.getReceiverExpression() ?: return null
+        val bindingContext = qualifiedElement.analyze(BodyResolveMode.PARTIAL)
+        if (bindingContext[BindingContext.QUALIFIER, receiver] == null) return null
 
         val targets = nameExpression.mainReference.resolveToDescriptors(bindingContext)
         if (targets.isEmpty()) return null
diff --git a/idea/testData/intentions/importMember/NestedClass1.kt b/idea/testData/intentions/importMember/NestedClass1.kt
new file mode 100644
index 00000000000..7ef09cf8aa0
--- /dev/null
+++ b/idea/testData/intentions/importMember/NestedClass1.kt
@@ -0,0 +1,15 @@
+// INTENTION_TEXT: "Add import for 'ppp.Foo.Bar'"
+// WITH_RUNTIME
+package ppp
+
+sealed class Foo {
+    class Bar(val x: Int) : Foo()
+}
+
+fun test() {
+    val foo = Foo.<caret>Bar(5)
+
+    when (foo) {
+        is Foo.Bar -> println(foo.x)
+    }
+}
diff --git a/idea/testData/intentions/importMember/NestedClass1.kt.after b/idea/testData/intentions/importMember/NestedClass1.kt.after
new file mode 100644
index 00000000000..9ee1b0a84d2
--- /dev/null
+++ b/idea/testData/intentions/importMember/NestedClass1.kt.after
@@ -0,0 +1,17 @@
+// INTENTION_TEXT: "Add import for 'ppp.Foo.Bar'"
+// WITH_RUNTIME
+package ppp
+
+import ppp.Foo.Bar
+
+sealed class Foo {
+    class Bar(val x: Int) : Foo()
+}
+
+fun test() {
+    val foo = <caret>Bar(5)
+
+    when (foo) {
+        is Bar -> println(foo.x)
+    }
+}
diff --git a/idea/testData/intentions/importMember/NestedClass2.kt b/idea/testData/intentions/importMember/NestedClass2.kt
new file mode 100644
index 00000000000..05297bbea81
--- /dev/null
+++ b/idea/testData/intentions/importMember/NestedClass2.kt
@@ -0,0 +1,15 @@
+// INTENTION_TEXT: "Add import for 'ppp.Foo.Bar'"
+// WITH_RUNTIME
+package ppp
+
+sealed class Foo {
+    class Bar(val x: Int) : Foo()
+}
+
+fun test() {
+    val foo = Foo.Bar(5)
+
+    when (foo) {
+        is Foo.<caret>Bar -> println(foo.x)
+    }
+}
diff --git a/idea/testData/intentions/importMember/NestedClass2.kt.after b/idea/testData/intentions/importMember/NestedClass2.kt.after
new file mode 100644
index 00000000000..3b7fc6fdc47
--- /dev/null
+++ b/idea/testData/intentions/importMember/NestedClass2.kt.after
@@ -0,0 +1,17 @@
+// INTENTION_TEXT: "Add import for 'ppp.Foo.Bar'"
+// WITH_RUNTIME
+package ppp
+
+import ppp.Foo.Bar
+
+sealed class Foo {
+    class Bar(val x: Int) : Foo()
+}
+
+fun test() {
+    val foo = Bar(5)
+
+    when (foo) {
+        is <caret>Bar -> println(foo.x)
+    }
+}
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index e78af45d318..3af51371e3d 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -5706,6 +5706,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("NestedClass1.kt")
+        public void testNestedClass1() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/importMember/NestedClass1.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("NestedClass2.kt")
+        public void testNestedClass2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/importMember/NestedClass2.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("NoTarget.kt")
         public void testNoTarget() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/importMember/NoTarget.kt");