From a8d9c761c8d524c30dab4b4746a49952a673d47b Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Mon, 19 Oct 2015 15:09:37 +0300
Subject: [PATCH] Correct detection of non-qualified references in ReplaceWith
 analyzer

---
 .../kotlin/idea/references/referenceUtil.kt       | 13 +++++++++++++
 .../KotlinCopyPasteReferenceProcessor.kt          | 15 +++------------
 .../kotlin/idea/imports/KotlinImportOptimizer.kt  |  9 ++-------
 .../replaceWith/ReplaceWithAnnotationAnalyzer.kt  |  7 ++++---
 .../toMethodFromCompanionObject.after.kt          |  4 ++--
 5 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt
index 19fa169d541..540fe98cb47 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt
@@ -18,9 +18,12 @@ package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.*
 import org.jetbrains.kotlin.asJava.unwrapped
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.imports.canBeReferencedViaImport
 import org.jetbrains.kotlin.idea.intentions.OperatorToFunctionIntention
 import org.jetbrains.kotlin.idea.kdoc.KDocReference
+import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
 import org.jetbrains.kotlin.lexer.JetTokens
@@ -30,6 +33,7 @@ import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.isReallySuccess
+import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import org.jetbrains.kotlin.utils.addToStdlib.constant
@@ -182,3 +186,12 @@ public fun JetExpression.readWriteAccess(useResolveForReadWrite: Boolean): Refer
     else
         ReferenceAccess.READ
 }
+
+public fun JetReference.canBeResolvedViaImport(target: DeclarationDescriptor): Boolean {
+    if (!target.canBeReferencedViaImport()) return false
+    if (target.isExtension) return true // assume that any type of reference can use imports when resolved to extension
+    val referenceExpression = this.element as? JetNameReferenceExpression ?: return false
+    if (CallTypeAndReceiver.detect(referenceExpression).receiver != null) return false
+    if (element.parent is JetThisExpression || element.parent is JetSuperExpression) return false // TODO: it's a bad design of PSI tree, we should change it
+    return true
+}
\ No newline at end of file
diff --git a/idea/src/org/jetbrains/kotlin/idea/codeInsight/KotlinCopyPasteReferenceProcessor.kt b/idea/src/org/jetbrains/kotlin/idea/codeInsight/KotlinCopyPasteReferenceProcessor.kt
index a26ac078add..91d0c503fe1 100644
--- a/idea/src/org/jetbrains/kotlin/idea/codeInsight/KotlinCopyPasteReferenceProcessor.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/codeInsight/KotlinCopyPasteReferenceProcessor.kt
@@ -42,11 +42,7 @@ import org.jetbrains.kotlin.idea.conversion.copy.end
 import org.jetbrains.kotlin.idea.conversion.copy.range
 import org.jetbrains.kotlin.idea.conversion.copy.start
 import org.jetbrains.kotlin.idea.imports.importableFqName
-import org.jetbrains.kotlin.idea.references.JetMultiReference
-import org.jetbrains.kotlin.idea.references.JetReference
-import org.jetbrains.kotlin.idea.references.JetSimpleNameReference
-import org.jetbrains.kotlin.idea.references.mainReference
-import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
+import org.jetbrains.kotlin.idea.references.*
 import org.jetbrains.kotlin.idea.util.ImportInsertHelper
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
@@ -56,7 +52,6 @@ import org.jetbrains.kotlin.psi.psiUtil.elementsInRange
 import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.resolve.DescriptorUtils
-import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.scopes.JetScope
 import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -154,13 +149,9 @@ public class KotlinCopyPasteReferenceProcessor() : CopyPastePostProcessor<Kotlin
                 val declaration = declarations.singleOrNull()
                 if (declaration != null && declaration.isInCopiedArea(file, startOffsets, endOffsets)) continue
 
-                if (!descriptor.isExtension) {
-                    if (element !is JetNameReferenceExpression) continue
-                    if (CallTypeAndReceiver.detect(element).receiver != null) continue
-                    if (element.parent is JetThisExpression || element.parent is JetSuperExpression) continue // TODO: it's a bad design of PSI tree, we should change it
-                }
+                if (!reference.canBeResolvedViaImport(descriptor)) continue
 
-                val fqName = descriptor.importableFqName ?: continue
+                val fqName = descriptor.importableFqName!!
 
                 val kind = KotlinReferenceData.Kind.fromDescriptor(descriptor) ?: continue
                 add(KotlinReferenceData(element.range.start - startOffset, element.range.end - startOffset, fqName.asString(), kind))
diff --git a/idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt b/idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt
index 8bc37c24b47..4866915100b 100644
--- a/idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt
@@ -28,7 +28,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.core.formatter.JetCodeStyleSettings
 import org.jetbrains.kotlin.idea.core.getResolutionScope
 import org.jetbrains.kotlin.idea.references.JetReference
-import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
+import org.jetbrains.kotlin.idea.references.canBeResolvedViaImport
 import org.jetbrains.kotlin.idea.util.ImportInsertHelper
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
@@ -39,7 +39,6 @@ import org.jetbrains.kotlin.renderer.render
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.ImportPath
 import org.jetbrains.kotlin.resolve.descriptorUtil.getImportableDescriptor
-import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.scopes.JetScope
 import org.jetbrains.kotlin.resolve.scopes.utils.asJetScope
 import org.jetbrains.kotlin.resolve.scopes.utils.getImplicitReceiversHierarchy
@@ -103,11 +102,7 @@ public class KotlinImportOptimizer() : ImportOptimizer {
                     if (target is PackageViewDescriptor && parentFqName == FqName.ROOT) continue // no need to import top-level packages
                     if (target !is PackageViewDescriptor && parentFqName == currentPackageName) continue
 
-                    if (!target.isExtension) { // for non-extension targets, count only non-qualified simple name usages
-                        if (element !is JetNameReferenceExpression) continue
-                        if (CallTypeAndReceiver.detect(element).receiver != null) continue
-                        if (element.parent is JetThisExpression || element.parent is JetSuperExpression) continue // TODO: it's a bad design of PSI tree, we should change it
-                    }
+                    if (!reference.canBeResolvedViaImport(target)) continue
 
                     val importableDescriptor = target.getImportableDescriptor()
 
diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/replaceWith/ReplaceWithAnnotationAnalyzer.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/replaceWith/ReplaceWithAnnotationAnalyzer.kt
index 0c67b36e324..8fbe00737ff 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/replaceWith/ReplaceWithAnnotationAnalyzer.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/replaceWith/ReplaceWithAnnotationAnalyzer.kt
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.intentions.InsertExplicitTypeArgumentsIntention
 import org.jetbrains.kotlin.idea.references.JetSimpleNameReference
+import org.jetbrains.kotlin.idea.references.canBeResolvedViaImport
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.name.FqName
@@ -50,7 +51,6 @@ import org.jetbrains.kotlin.storage.LockBasedStorageManager
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor
-import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.*
 
 data class ReplaceWith(val pattern: String, val imports: List<String>)
@@ -121,8 +121,9 @@ object ReplaceWithAnnotationAnalyzer {
         expression.forEachDescendantOfType<JetSimpleNameExpression> { expression ->
             val target = bindingContext[BindingContext.REFERENCE_TARGET, expression] ?: return@forEachDescendantOfType
 
-            if (target.isExtension || expression.getReceiverExpression() == null) {
-                importFqNames.addIfNotNull(target.importableFqName)
+            //TODO: other types of references ('[]' etc)
+            if (expression.mainReference.canBeResolvedViaImport(target)) {
+                importFqNames.add(target.importableFqName!!)
             }
 
             if (expression.getReceiverExpression() == null) {
diff --git a/idea/testData/quickfix/deprecatedSymbolUsage/toMethodFromCompanionObject.after.kt b/idea/testData/quickfix/deprecatedSymbolUsage/toMethodFromCompanionObject.after.kt
index 944907d6d00..be7e6d377b6 100644
--- a/idea/testData/quickfix/deprecatedSymbolUsage/toMethodFromCompanionObject.after.kt
+++ b/idea/testData/quickfix/deprecatedSymbolUsage/toMethodFromCompanionObject.after.kt
@@ -1,7 +1,7 @@
-import dependency.C
+import dependency.C.Companion.newFun
 
 // "Replace with 'newFun(this)'" "true"
 
 fun foo(c: dependency.C) {
-    C.<caret>newFun(c)
+    newFun(c)
 }
\ No newline at end of file