From 480982b69013cbf6d737ab6547eac2fc56e972ac Mon Sep 17 00:00:00 2001
From: Alexey Sedunov <alexey.sedunov@jetbrains.com>
Date: Fri, 21 Apr 2017 15:27:34 +0300
Subject: [PATCH] Shorten References: Wrap elements to shorten in smart
 pointers

 #KT-17526 Fixed
---
 .../jetbrains/kotlin/idea/core/ShortenReferences.kt  |  9 ++++++---
 .../shortCompanionRefInsideShortenedCall.kt          | 12 ++++++++++++
 .../shortCompanionRefInsideShortenedCall.kt.after    | 12 ++++++++++++
 .../kotlin/shortenRefs/ShortenRefsTestGenerated.java |  6 ++++++
 4 files changed, 36 insertions(+), 3 deletions(-)
 create mode 100644 idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt
 create mode 100644 idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt.after

diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt
index eee2997aafa..d7f2052b24f 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt
@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.core
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
+import com.intellij.psi.SmartPsiElementPointer
 import com.intellij.psi.impl.source.PostprocessReformattingAspect
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.descriptors.*
@@ -33,6 +34,7 @@ import org.jetbrains.kotlin.idea.util.ShadowedDeclarationsFilter
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForReceiver
 import org.jetbrains.kotlin.psi.psiUtil.parents
@@ -260,7 +262,7 @@ class ShortenReferences(val options: (KtElement) -> Options = { Options.DEFAULT
             protected val failedToImportDescriptors: Set<DeclarationDescriptor>
     ) {
         protected val resolutionFacade = file.getResolutionFacade()
-        private val elementsToShorten = ArrayList<TElement>()
+        private val elementsToShorten = ArrayList<SmartPsiElementPointer<TElement>>()
         private val descriptorsToImport = LinkedHashSet<DeclarationDescriptor>()
 
         abstract val collectElementsVisitor: CollectElementsVisitor<TElement>
@@ -277,7 +279,7 @@ class ShortenReferences(val options: (KtElement) -> Options = { Options.DEFAULT
                 val toBeShortened: Boolean
                 when (result) {
                     is AnalyzeQualifiedElementResult.ShortenNow -> {
-                        elementsToShorten.add(element)
+                        elementsToShorten.add(element.createSmartPointer())
                         toBeShortened = true
                     }
 
@@ -324,7 +326,8 @@ class ShortenReferences(val options: (KtElement) -> Options = { Options.DEFAULT
         protected abstract fun shortenElement(element: TElement): KtElement
 
         fun shortenElements(elementSetToUpdate: MutableSet<KtElement>) {
-            for (element in elementsToShorten) {
+            for (elementPointer in elementsToShorten) {
+                val element = elementPointer.element ?: continue
                 if (!element.isValid) continue
 
                 var newElement: KtElement? = null
diff --git a/idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt b/idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt
new file mode 100644
index 00000000000..0f758ca27c2
--- /dev/null
+++ b/idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt
@@ -0,0 +1,12 @@
+// REMOVE_COMPANION_REF
+package test
+
+class CompanionExtraPlain { companion object { val cmpVal = 1 } }
+
+fun wrap(p: Int) = p + 1
+
+class Dome {
+    fun refer() {
+        <selection>test.wrap(test.CompanionExtraPlain.Companion.cmpVal)</selection>
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt.after b/idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt.after
new file mode 100644
index 00000000000..326f3b9f694
--- /dev/null
+++ b/idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt.after
@@ -0,0 +1,12 @@
+// REMOVE_COMPANION_REF
+package test
+
+class CompanionExtraPlain { companion object { val cmpVal = 1 } }
+
+fun wrap(p: Int) = p + 1
+
+class Dome {
+    fun refer() {
+        wrap(CompanionExtraPlain.cmpVal)
+    }
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java b/idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java
index 312e7e19555..dbc09ab33b7 100644
--- a/idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java
@@ -138,6 +138,12 @@ public class ShortenRefsTestGenerated extends AbstractShortenRefsTest {
         doTest(fileName);
     }
 
+    @TestMetadata("shortCompanionRefInsideShortenedCall.kt")
+    public void testShortCompanionRefInsideShortenedCall() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/shortenRefs/shortCompanionRefInsideShortenedCall.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("UnresolvedOverload.kt")
     public void testUnresolvedOverload() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/shortenRefs/UnresolvedOverload.kt");