From 633679ac16fb89837c5d69741326ccc0f9637149 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Sat, 14 Apr 2018 12:29:51 +0300
Subject: [PATCH] Preserve comments in "Remove variable" quick fix #KT-23753
 Fixed

---
 .../kotlin/idea/quickfix/RemovePsiElementSimpleFix.kt     | 8 +++++---
 .../variables/unusedVariableWithInitializerAndComment.kt  | 6 ++++++
 .../unusedVariableWithInitializerAndComment.kt.after      | 6 ++++++
 .../kotlin/idea/quickfix/QuickFixTestGenerated.java       | 5 +++++
 4 files changed, 22 insertions(+), 3 deletions(-)
 create mode 100644 idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt
 create mode 100644 idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/RemovePsiElementSimpleFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/RemovePsiElementSimpleFix.kt
index 4b1525ce9eb..accacabef09 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/RemovePsiElementSimpleFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/RemovePsiElementSimpleFix.kt
@@ -20,6 +20,7 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
@@ -63,9 +64,10 @@ open class RemovePsiElementSimpleFix(element: PsiElement, private val text: Stri
                 override fun invoke(project: Project, editor: Editor?, file: KtFile) {
                     val initializer = expression.initializer
                     if (initializer != null && initializer !is KtConstantExpression) {
-                        expression.replace(initializer)
-                    }
-                    else {
+                        val commentSaver = CommentSaver(expression)
+                        val replaced = expression.replace(initializer)
+                        commentSaver.restore(replaced)
+                    } else {
                         expression.delete()
                     }
                 }
diff --git a/idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt b/idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt
new file mode 100644
index 00000000000..edf1cfd11e0
--- /dev/null
+++ b/idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt
@@ -0,0 +1,6 @@
+// "Remove variable 'a'" "true"
+var cnt = 5
+fun getCnt() = cnt++
+fun f() {
+    var <caret>a = getCnt() // comment
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt.after b/idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt.after
new file mode 100644
index 00000000000..ae5523ce62e
--- /dev/null
+++ b/idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt.after
@@ -0,0 +1,6 @@
+// "Remove variable 'a'" "true"
+var cnt = 5
+fun getCnt() = cnt++
+fun f() {
+    getCnt() // comment
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
index e26dd226963..902f3f3aafb 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
@@ -11668,6 +11668,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/variables/unusedVariableWithInitializer.kt");
         }
 
+        @TestMetadata("unusedVariableWithInitializerAndComment.kt")
+        public void testUnusedVariableWithInitializerAndComment() throws Exception {
+            runTest("idea/testData/quickfix/variables/unusedVariableWithInitializerAndComment.kt");
+        }
+
         @TestMetadata("unusedVariableWithNullInitializer.kt")
         public void testUnusedVariableWithNullInitializer() throws Exception {
             runTest("idea/testData/quickfix/variables/unusedVariableWithNullInitializer.kt");