From 3ea0222d43615016282e04c4bea514e65065fec1 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Wed, 29 Aug 2018 13:10:36 +0900
Subject: [PATCH] "Add annotation target" quickfix: add SOURCE retention when
 target is EXPRESSION #KT-26306 Fixed

---
 .../idea/quickfix/AddAnnotationTargetFix.kt     | 17 +++++++++++++----
 .../quickfix/addAnnotationTarget/expression2.kt |  7 +++++++
 .../addAnnotationTarget/expression2.kt.after    |  9 +++++++++
 .../quickfix/addAnnotationTarget/expression3.kt |  8 ++++++++
 .../addAnnotationTarget/expression3.kt.after    |  9 +++++++++
 .../idea/quickfix/QuickFixTestGenerated.java    | 10 ++++++++++
 6 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 idea/testData/quickfix/addAnnotationTarget/expression2.kt
 create mode 100644 idea/testData/quickfix/addAnnotationTarget/expression2.kt.after
 create mode 100644 idea/testData/quickfix/addAnnotationTarget/expression3.kt
 create mode 100644 idea/testData/quickfix/addAnnotationTarget/expression3.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/AddAnnotationTargetFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/AddAnnotationTargetFix.kt
index 34f2a90f27b..4fc6abe7005 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/AddAnnotationTargetFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/AddAnnotationTargetFix.kt
@@ -119,9 +119,9 @@ private fun PsiAnnotation.getActualTargetList(): List<KotlinTarget> {
 
 private fun KtAnnotationEntry.getActualTargetList(): List<KotlinTarget> {
     val annotatedElement = getStrictParentOfType<KtModifierList>()?.owner as? KtElement
-            ?: getStrictParentOfType<KtAnnotatedExpression>()?.baseExpression
-            ?: getStrictParentOfType<KtFile>()
-            ?: return emptyList()
+        ?: getStrictParentOfType<KtAnnotatedExpression>()?.baseExpression
+        ?: getStrictParentOfType<KtFile>()
+        ?: return emptyList()
 
     val targetList = AnnotationChecker.getActualTargetList(annotatedElement, null, BindingTraceContext().bindingContext)
 
@@ -147,8 +147,17 @@ private fun KtAnnotationEntry.getActualTargetList(): List<KotlinTarget> {
 }
 
 private fun KtClass.addAnnotationTargets(annotationTargets: List<KotlinTarget>, psiFactory: KtPsiFactory) {
-    val targetAnnotationName = KotlinBuiltIns.FQ_NAMES.target.shortName().asString()
+    val retentionAnnotationName = KotlinBuiltIns.FQ_NAMES.retention.shortName().asString()
+    if (annotationTargets.any { it == KotlinTarget.EXPRESSION }
+        && annotationEntries.none { it.typeReference?.text == retentionAnnotationName }) {
+        addAnnotationEntry(
+            psiFactory.createAnnotationEntry(
+                "@$retentionAnnotationName(${KotlinBuiltIns.FQ_NAMES.annotationRetention.shortName()}.${AnnotationRetention.SOURCE.name})"
+            )
+        )
+    }
 
+    val targetAnnotationName = KotlinBuiltIns.FQ_NAMES.target.shortName().asString()
     val targetAnnotationEntry = annotationEntries.find { it.typeReference?.text == targetAnnotationName } ?: run {
         val text = "@$targetAnnotationName${annotationTargets.toArgumentListString()}"
         addAnnotationEntry(psiFactory.createAnnotationEntry(text))
diff --git a/idea/testData/quickfix/addAnnotationTarget/expression2.kt b/idea/testData/quickfix/addAnnotationTarget/expression2.kt
new file mode 100644
index 00000000000..2e359226df1
--- /dev/null
+++ b/idea/testData/quickfix/addAnnotationTarget/expression2.kt
@@ -0,0 +1,7 @@
+// "Add annotation target" "true"
+annotation class Foo
+
+fun test() {
+    var v = 0
+    <caret>@Foo v++
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/addAnnotationTarget/expression2.kt.after b/idea/testData/quickfix/addAnnotationTarget/expression2.kt.after
new file mode 100644
index 00000000000..8060c88cff4
--- /dev/null
+++ b/idea/testData/quickfix/addAnnotationTarget/expression2.kt.after
@@ -0,0 +1,9 @@
+// "Add annotation target" "true"
+@Target(AnnotationTarget.EXPRESSION)
+@Retention(AnnotationRetention.SOURCE)
+annotation class Foo
+
+fun test() {
+    var v = 0
+    @Foo v++
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/addAnnotationTarget/expression3.kt b/idea/testData/quickfix/addAnnotationTarget/expression3.kt
new file mode 100644
index 00000000000..92c947c1e10
--- /dev/null
+++ b/idea/testData/quickfix/addAnnotationTarget/expression3.kt
@@ -0,0 +1,8 @@
+// "Add annotation target" "true"
+@Retention
+annotation class Foo
+
+fun test() {
+    var v = 0
+    <caret>@Foo v++
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/addAnnotationTarget/expression3.kt.after b/idea/testData/quickfix/addAnnotationTarget/expression3.kt.after
new file mode 100644
index 00000000000..3f709ace784
--- /dev/null
+++ b/idea/testData/quickfix/addAnnotationTarget/expression3.kt.after
@@ -0,0 +1,9 @@
+// "Add annotation target" "true"
+@Target(AnnotationTarget.EXPRESSION)
+@Retention
+annotation class Foo
+
+fun test() {
+    var v = 0
+    @Foo v++
+}
\ 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 da91aa6b372..4d4d758b6d8 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
@@ -239,6 +239,16 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/addAnnotationTarget/expression.kt");
         }
 
+        @TestMetadata("expression2.kt")
+        public void testExpression2() throws Exception {
+            runTest("idea/testData/quickfix/addAnnotationTarget/expression2.kt");
+        }
+
+        @TestMetadata("expression3.kt")
+        public void testExpression3() throws Exception {
+            runTest("idea/testData/quickfix/addAnnotationTarget/expression3.kt");
+        }
+
         @TestMetadata("file.kt")
         public void testFile() throws Exception {
             runTest("idea/testData/quickfix/addAnnotationTarget/file.kt");