Fixed important bug in reference shortening

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/ShortenReferences.kt

@@ -154,6 +154,7 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
             elements: List<JetElement>,
             elementFilter: (PsiElement) -> FilterResult
     ) {
+        //TODO: that's not correct since we have options!
         val elementsToUse = dropNestedElements(elements)
 
         val importInserter = ImportInserter(file)
@@ -169,7 +170,7 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
                     ShortenQualifiedExpressionsVisitor(file, elementFilter, failedToImportDescriptors)
             )
             val descriptorsToImport = visitors.flatMap { analyzeReferences(elementsToUse, it) }.toSet()
-            visitors.forEach { elementsToUse.removeAll(it.shortenElements()) }
+            visitors.forEach { it.shortenElements(elementsToUse) }
 
             var anyChange = false
             for (descriptor in descriptorsToImport) {
@@ -232,7 +233,7 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
 
         protected abstract fun qualifier(element: T): JetElement
 
-        protected abstract fun shortenElement(element: T)
+        protected abstract fun shortenElement(element: T): JetElement
 
         override fun visitElement(element: PsiElement) {
             if (elementFilter(element) != FilterResult.SKIP) {
@@ -240,12 +241,15 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
             }
         }
 
-        public fun shortenElements(): Collection<T> {
+        public fun shortenElements(elementSetToUpdate: MutableSet<JetElement>) {
             for (element in elementsToShorten) {
                 if (!element.isValid()) continue
-                shortenElement(element)
+                val newElement = shortenElement(element)
+                if (element in elementSetToUpdate && newElement != element) {
+                    elementSetToUpdate.remove(element)
+                    elementSetToUpdate.add(newElement)
+                }
             }
-            return elementsToShorten
         }
 
         public fun getDescriptorsToImport(): Set<DeclarationDescriptor> = descriptorsToImport
@@ -288,8 +292,9 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
 
         override fun qualifier(element: JetUserType) = element.getQualifier()!!
 
-        override fun shortenElement(element: JetUserType) {
+        override fun shortenElement(element: JetUserType): JetElement {
             element.deleteQualifier()
+            return element
         }
     }
 
@@ -365,8 +370,8 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
 
         override fun qualifier(element: JetQualifiedExpression) = element.getReceiverExpression()
 
-        override fun shortenElement(element: JetQualifiedExpression) {
-            element.replace(element.getSelectorExpression()!!)
+        override fun shortenElement(element: JetQualifiedExpression): JetElement {
+            return element.replace(element.getSelectorExpression()!!) as JetElement
         }
     }
 
@@ -400,8 +405,8 @@ public class ShortenReferences(val options: (JetElement) -> Options = { Options.
         override fun qualifier(element: JetThisExpression): JetElement =
                 throw AssertionError("Qualifier requested: ${JetPsiUtil.getElementTextWithContext(element)}")
 
-        override fun shortenElement(element: JetThisExpression) {
-            element.replace(simpleThis)
+        override fun shortenElement(element: JetThisExpression): JetElement {
+            return element.replace(simpleThis) as JetElement
         }
     }