From 08caf0a0112f25cb315ab6e81c31ec26f4c91049 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <Mikhail.Glukhikh@jetbrains.com>
Date: Thu, 26 Jan 2017 16:16:23 +0300
Subject: [PATCH] Object literal to lambda: handle last comment correctly
 #KT-15670 Fixed

---
 .../ObjectLiteralToLambdaIntention.kt          | 18 +++++++++++-------
 .../objectLiteralToLambda/WithComment.kt       |  7 +++++++
 .../objectLiteralToLambda/WithComment.kt.after |  5 +++++
 .../WithCommentAfterExpression.kt              |  5 +++++
 .../WithCommentAfterExpression.kt.after        |  5 +++++
 .../inspectionData/expected.xml                | 18 ++++++++++++++++++
 .../intentions/IntentionTestGenerated.java     | 12 ++++++++++++
 7 files changed, 63 insertions(+), 7 deletions(-)
 create mode 100644 idea/testData/intentions/objectLiteralToLambda/WithComment.kt
 create mode 100644 idea/testData/intentions/objectLiteralToLambda/WithComment.kt.after
 create mode 100644 idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt
 create mode 100644 idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ObjectLiteralToLambdaIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ObjectLiteralToLambdaIntention.kt
index bef56f68b14..d0b95967e42 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ObjectLiteralToLambdaIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ObjectLiteralToLambdaIntention.kt
@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiComment
 import com.intellij.psi.search.LocalSearchScope
 import com.intellij.psi.search.searches.ReferencesSearch
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
@@ -34,10 +35,7 @@ import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.load.java.sam.SingleAbstractMethodUtils
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
-import org.jetbrains.kotlin.psi.psiUtil.contentRange
-import org.jetbrains.kotlin.psi.psiUtil.endOffset
-import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getCalleeExpressionIfAny
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
@@ -85,7 +83,7 @@ class ObjectLiteralToLambdaIntention : SelfTargetingRangeIntention<KtObjectLiter
     override fun applyTo(element: KtObjectLiteralExpression, editor: Editor?) {
         val commentSaver = CommentSaver(element)
 
-        val (baseTypeRef, baseType, singleFunction) = extractData(element)!!
+        val (_, baseType, singleFunction) = extractData(element)!!
 
         val returnSaver = ReturnSaver(singleFunction)
 
@@ -111,13 +109,19 @@ class ObjectLiteralToLambdaIntention : SelfTargetingRangeIntention<KtObjectLiter
                 appendFixedText("->")
             }
 
-            if (singleFunction.hasBlockBody()) {
-                appendChildRange((body as KtBlockExpression).contentRange())
+            val lastCommentOwner = if (singleFunction.hasBlockBody()) {
+                val contentRange = (body as KtBlockExpression).contentRange()
+                appendChildRange(contentRange)
+                contentRange.last
             }
             else {
                 appendExpression(body)
+                body
             }
 
+            if (lastCommentOwner?.anyDescendantOfType<PsiComment> { it.tokenType == KtTokens.EOL_COMMENT } ?: false) {
+                appendFixedText("\n")
+            }
             appendFixedText("}")
         }
 
diff --git a/idea/testData/intentions/objectLiteralToLambda/WithComment.kt b/idea/testData/intentions/objectLiteralToLambda/WithComment.kt
new file mode 100644
index 00000000000..cc12577974b
--- /dev/null
+++ b/idea/testData/intentions/objectLiteralToLambda/WithComment.kt
@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+val r = object: Runnable<caret> {
+    override fun run() {
+        TODO("not implemented") //To change body ...
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/objectLiteralToLambda/WithComment.kt.after b/idea/testData/intentions/objectLiteralToLambda/WithComment.kt.after
new file mode 100644
index 00000000000..5ade78009ef
--- /dev/null
+++ b/idea/testData/intentions/objectLiteralToLambda/WithComment.kt.after
@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+
+val r = Runnable {
+    TODO("not implemented") //To change body ...
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt b/idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt
new file mode 100644
index 00000000000..4c32f6c87e2
--- /dev/null
+++ b/idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt
@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+
+val r = object: Runnable<caret> {
+    override fun run() = TODO("not implemented") //To change body ...
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt.after b/idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt.after
new file mode 100644
index 00000000000..5ade78009ef
--- /dev/null
+++ b/idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt.after
@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+
+val r = Runnable {
+    TODO("not implemented") //To change body ...
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/objectLiteralToLambda/inspectionData/expected.xml b/idea/testData/intentions/objectLiteralToLambda/inspectionData/expected.xml
index 7b4651615d3..581385f43c2 100644
--- a/idea/testData/intentions/objectLiteralToLambda/inspectionData/expected.xml
+++ b/idea/testData/intentions/objectLiteralToLambda/inspectionData/expected.xml
@@ -166,4 +166,22 @@
         <problem_class severity="WARNING" attribute_key="WARNING_ATTRIBUTES">Convert object literal to lambda</problem_class>
         <description>Convert to lambda</description>
     </problem>
+
+    <problem>
+        <file>WithComment.kt</file>
+        <line>3</line>
+        <module>light_idea_test_case</module>
+        <entry_point TYPE="file" FQNAME="temp:///src/WithComment.kt" />
+        <problem_class severity="WARNING" attribute_key="WARNING_ATTRIBUTES">Object literal can be converted to lambda</problem_class>
+        <description>Convert to lambda</description>
+    </problem>
+
+    <problem>
+        <file>WithCommentAfterExpression.kt</file>
+        <line>3</line>
+        <module>light_idea_test_case</module>
+        <entry_point TYPE="file" FQNAME="temp:///src/WithCommentAfterExpression.kt" />
+        <problem_class severity="WARNING" attribute_key="WARNING_ATTRIBUTES">Object literal can be converted to lambda</problem_class>
+        <description>Convert to lambda</description>
+    </problem>
 </problems>
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 5a32d2004e5..c0d40c1c0e4 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -11112,6 +11112,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("WithComment.kt")
+        public void testWithComment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/objectLiteralToLambda/WithComment.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("WithCommentAfterExpression.kt")
+        public void testWithCommentAfterExpression() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/objectLiteralToLambda/WithCommentAfterExpression.kt");
+            doTest(fileName);
+        }
+
     }
 
     @TestMetadata("idea/testData/intentions/operatorToFunction")