From 18c181641f69bcb872c188082c348fe5180d551b Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <nikolay.krasko@jetbrains.com>
Date: Tue, 25 Sep 2018 13:24:33 +0300
Subject: [PATCH] Return updated element after reference shortening

Fresh new expression obtained after qualified expression binding is
invalidated after shortening.
---
 .../idea/references/KtSimpleNameReference.kt   | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/KtSimpleNameReference.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/KtSimpleNameReference.kt
index fd7d61cc34e..098d179e069 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/KtSimpleNameReference.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/KtSimpleNameReference.kt
@@ -178,16 +178,18 @@ class KtSimpleNameReference(expression: KtSimpleNameExpression) : KtSimpleRefere
 
         if (shorteningMode == ShorteningMode.NO_SHORTENING) return newExpression
 
-        val needToShorten =
-                PsiTreeUtil.getParentOfType(expression, KtImportDirective::class.java, KtPackageDirective::class.java) == null
-        if (needToShorten) {
-            when (shorteningMode) {
-                ShorteningMode.FORCED_SHORTENING -> ShortenReferences.DEFAULT.process(newQualifiedElement)
-                else -> newQualifiedElement.addToShorteningWaitSet()
-            }
+        val needToShorten = PsiTreeUtil.getParentOfType(expression, KtImportDirective::class.java, KtPackageDirective::class.java) == null
+        if (!needToShorten) {
+            return newExpression
         }
 
-        return newExpression
+        return when (shorteningMode) {
+            ShorteningMode.FORCED_SHORTENING -> ShortenReferences.DEFAULT.process(newQualifiedElement)
+            else -> {
+                newQualifiedElement.addToShorteningWaitSet()
+                newExpression
+            }
+        }
     }
 
     /**