From bf7c325f437f163f1f0130258f0c79630e7ebc8c Mon Sep 17 00:00:00 2001
From: Dmitry Gridin <dmitry.gridin@jetbrains.com>
Date: Tue, 1 Oct 2019 16:40:39 +0700
Subject: [PATCH] Completion: should prefer ShortenReferences by elements
 instead of text range

---
 .../kotlin/idea/completion/CompletionUtils.kt | 20 +++++++++++++++++--
 .../handlers/KotlinCallableInsertHandler.kt   |  3 ++-
 .../handlers/KotlinClassifierInsertHandler.kt |  4 ++--
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionUtils.kt b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionUtils.kt
index 0ee78fb6e19..b2248ae6745 100644
--- a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionUtils.kt
+++ b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionUtils.kt
@@ -27,7 +27,9 @@ import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.util.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.renderer.render
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
@@ -36,6 +38,7 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.isError
 import org.jetbrains.kotlin.types.typeUtil.TypeNullability
 import org.jetbrains.kotlin.types.typeUtil.nullability
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import java.util.*
 
 tailrec fun <T : Any> LookupElement.putUserDataDeep(key: Key<T>, value: T?) {
@@ -333,9 +336,22 @@ private open class BaseTypeLookupElement(type: KotlinType, baseLookupElement: Lo
     }
 }
 
-fun shortenReferences(context: InsertionContext, startOffset: Int, endOffset: Int) {
+fun shortenReferences(
+    context: InsertionContext,
+    startOffset: Int,
+    endOffset: Int,
+    shortenReferences: ShortenReferences = ShortenReferences.DEFAULT
+) {
     PsiDocumentManager.getInstance(context.project).commitAllDocuments()
-    ShortenReferences.DEFAULT.process(context.file as KtFile, startOffset, endOffset)
+    val file = context.file as KtFile
+    val element = file.findElementAt(startOffset)?.parentsWithSelf?.find {
+        it.startOffset == startOffset && it.endOffset == endOffset
+    }?.safeAs<KtElement>()
+
+    if (element != null)
+        shortenReferences.process(element)
+    else
+        shortenReferences.process(file, startOffset, endOffset)
 }
 
 infix fun <T> ElementPattern<T>.and(rhs: ElementPattern<T>) = StandardPatterns.and(this, rhs)
diff --git a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/KotlinCallableInsertHandler.kt b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/KotlinCallableInsertHandler.kt
index ad390b00532..f2afb44b437 100644
--- a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/KotlinCallableInsertHandler.kt
+++ b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/KotlinCallableInsertHandler.kt
@@ -10,6 +10,7 @@ import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.psi.PsiDocumentManager
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.idea.completion.isArtificialImportAliasedDescriptor
+import org.jetbrains.kotlin.idea.completion.shortenReferences
 import org.jetbrains.kotlin.idea.core.ShortenReferences
 import org.jetbrains.kotlin.idea.core.completion.DeclarationLookupObject
 import org.jetbrains.kotlin.idea.core.withRootPrefixIfNeeded
@@ -54,7 +55,7 @@ abstract class KotlinCallableInsertHandler(val callType: CallType<*>) : BaseDecl
 
                 psiDocumentManager.commitAllDocuments()
 
-                SHORTEN_REFERENCES.process(file, context.startOffset, context.tailOffset - 1)
+                shortenReferences(context, context.startOffset, context.tailOffset - 1, SHORTEN_REFERENCES)
 
                 psiDocumentManager.doPostponedOperationsAndUnblockDocument(context.document)
 
diff --git a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/KotlinClassifierInsertHandler.kt b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/KotlinClassifierInsertHandler.kt
index fd8ecaef2ce..74b41e76528 100644
--- a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/KotlinClassifierInsertHandler.kt
+++ b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/KotlinClassifierInsertHandler.kt
@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.allowResolveInDispatchThread
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.completion.isAfterDot
 import org.jetbrains.kotlin.idea.completion.isArtificialImportAliasedDescriptor
-import org.jetbrains.kotlin.idea.core.ShortenReferences
+import org.jetbrains.kotlin.idea.completion.shortenReferences
 import org.jetbrains.kotlin.idea.core.canAddRootPrefix
 import org.jetbrains.kotlin.idea.core.completion.DeclarationLookupObject
 import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
@@ -83,7 +83,7 @@ object KotlinClassifierInsertHandler : BaseDeclarationInsertHandler() {
                 val rangeMarker = document.createRangeMarker(classNameStart, classNameEnd)
                 val wholeRangeMarker = document.createRangeMarker(startOffset, classNameEnd + tempSuffix.length)
 
-                ShortenReferences.DEFAULT.process(file, classNameStart, classNameEnd)
+                shortenReferences(context, classNameStart, classNameEnd)
                 psiDocumentManager.doPostponedOperationsAndUnblockDocument(document)
 
                 if (rangeMarker.isValid && wholeRangeMarker.isValid) {