From 8ce9b2d061b719c603637be8878994b85d31aaed Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Thu, 20 Aug 2020 17:15:00 +0900
Subject: [PATCH] KT-34572 Convert to block body action improperly works with
 suppress annotations (#2969)

* Convert to block body: place @Suppress annotation before return expression

#KT-34572 Fixed

* Convert to block body: place annotations before return expression
---
 .../intentions/ConvertToBlockBodyIntention.kt  | 18 ++++++++++++++----
 .../convertToBlockBody/annotatedExpr.kt.after  |  3 ++-
 .../convertToBlockBody/annotatedExpr2.kt       | 12 ++++++++++++
 .../convertToBlockBody/annotatedExpr2.kt.after | 13 +++++++++++++
 .../intentions/IntentionTestGenerated.java     |  5 +++++
 5 files changed, 46 insertions(+), 5 deletions(-)
 create mode 100644 idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt
 create mode 100644 idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt
index f7471303643..0ef3bd6526f 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt
@@ -58,10 +58,20 @@ class ConvertToBlockBodyIntention : SelfTargetingIntention<KtDeclarationWithBody
                 val factory = KtPsiFactory(declaration)
                 if (bodyType != null && bodyType.isUnit() && body is KtNameReferenceExpression) return factory.createEmptyBody()
                 val unitWhenAsResult = (bodyType == null || bodyType.isUnit()) && body.resultingWhens().isNotEmpty()
-                val needReturn = returnsValue &&
-                        (bodyType == null || (!bodyType.isUnit() && !bodyType.isNothing()))
-                val statement = if (needReturn || unitWhenAsResult) factory.createExpressionByPattern("return $0", body) else body
-                return factory.createSingleStatementBlock(statement)
+                val needReturn = returnsValue && (bodyType == null || (!bodyType.isUnit() && !bodyType.isNothing()))
+                return if (needReturn || unitWhenAsResult) {
+                    val annotatedExpr = body as? KtAnnotatedExpression
+                    val returnedExpr = annotatedExpr?.baseExpression ?: body
+                    val block = factory.createSingleStatementBlock(factory.createExpressionByPattern("return $0", returnedExpr))
+                    val statement = block.firstStatement
+                    annotatedExpr?.annotationEntries?.forEach {
+                        block.addBefore(it, statement)
+                        block.addBefore(factory.createNewLine(), statement)
+                    }
+                    block
+                } else {
+                    factory.createSingleStatementBlock(body)
+                }
             }
 
             val newBody = when (declaration) {
diff --git a/idea/testData/intentions/convertToBlockBody/annotatedExpr.kt.after b/idea/testData/intentions/convertToBlockBody/annotatedExpr.kt.after
index b867c306ec4..21a89e88dbb 100644
--- a/idea/testData/intentions/convertToBlockBody/annotatedExpr.kt.after
+++ b/idea/testData/intentions/convertToBlockBody/annotatedExpr.kt.after
@@ -3,5 +3,6 @@
 annotation class ann
 
 fun foo(): Int {
-    return (@ann 1)
+    @ann
+    return 1
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt b/idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt
new file mode 100644
index 00000000000..c111f8685c6
--- /dev/null
+++ b/idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt
@@ -0,0 +1,12 @@
+@Target(AnnotationTarget.EXPRESSION)
+@Retention(AnnotationRetention.SOURCE)
+annotation class ann
+
+@Target(AnnotationTarget.EXPRESSION)
+@Retention(AnnotationRetention.SOURCE)
+annotation class ann2
+
+fun test() =<caret>
+    @ann
+    @ann2
+    1
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt.after b/idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt.after
new file mode 100644
index 00000000000..0d85b74fdd7
--- /dev/null
+++ b/idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt.after
@@ -0,0 +1,13 @@
+@Target(AnnotationTarget.EXPRESSION)
+@Retention(AnnotationRetention.SOURCE)
+annotation class ann
+
+@Target(AnnotationTarget.EXPRESSION)
+@Retention(AnnotationRetention.SOURCE)
+annotation class ann2
+
+fun test(): Int {
+    @ann
+    @ann2
+    return 1
+}
\ 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 41ef4523197..45891da9f04 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -6887,6 +6887,11 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             runTest("idea/testData/intentions/convertToBlockBody/annotatedExpr.kt");
         }
 
+        @TestMetadata("annotatedExpr2.kt")
+        public void testAnnotatedExpr2() throws Exception {
+            runTest("idea/testData/intentions/convertToBlockBody/annotatedExpr2.kt");
+        }
+
         @TestMetadata("explicitlyNonUnitFun.kt")
         public void testExplicitlyNonUnitFun() throws Exception {
             runTest("idea/testData/intentions/convertToBlockBody/explicitlyNonUnitFun.kt");