Extract function for adding annotation to KtModifierListOwner

idea/idea-analysis/src/org/jetbrains/kotlin/idea/quickfix/KotlinSuppressIntentionAction.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.util.PsiPrecedences
+import org.jetbrains.kotlin.idea.util.addAnnotation
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.replaceFileAnnotationList
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -50,7 +51,13 @@ class KotlinSuppressIntentionAction private constructor(
         val id = "\"$suppressKey\""
         when (suppressAt) {
             is KtModifierListOwner ->
-                suppressAtModifierListOwner(suppressAt, id)
+                suppressAt.addAnnotation(KotlinBuiltIns.FQ_NAMES.suppress.toSafe(),
+                                         "$id",
+                                         whiteSpaceText = if (kind.newLineNeeded) "\n" else " ",
+                                         addToExistingAnnotation = { entry ->
+                                             addArgumentToSuppressAnnotation(entry, id)
+                                             true
+                                         })
 
             is KtAnnotatedExpression ->
                 suppressAtAnnotatedExpression(CaretBox(suppressAt, editor), id)
@@ -87,32 +94,6 @@ class KotlinSuppressIntentionAction private constructor(
         addArgumentToSuppressAnnotation(suppressAnnotation, id)
     }
 
-    private fun suppressAtModifierListOwner(suppressAt: KtModifierListOwner, id: String) {
-        val modifierList = suppressAt.modifierList
-        val psiFactory = KtPsiFactory(suppressAt)
-        if (modifierList == null) {
-            // create a modifier list from scratch
-            val newModifierList = psiFactory.createModifierList(suppressAnnotationText(id))
-            val replaced = KtPsiUtil.replaceModifierList(suppressAt, newModifierList)
-            val whiteSpace = psiFactory.createWhiteSpace(kind)
-            suppressAt.addAfter(whiteSpace, replaced)
-        }
-        else {
-            val entry = findSuppressAnnotation(suppressAt)
-            if (entry == null) {
-                // no [suppress] annotation
-                val newAnnotation = psiFactory.createAnnotationEntry(suppressAnnotationText(id))
-                val addedAnnotation = modifierList.addBefore(newAnnotation, modifierList.firstChild)
-                val whiteSpace = psiFactory.createWhiteSpace(kind)
-                modifierList.addAfter(whiteSpace, addedAnnotation)
-            }
-            else {
-                // already annotated with [suppress]
-                addArgumentToSuppressAnnotation(entry, id)
-            }
-        }
-    }
-
     private fun suppressAtAnnotatedExpression(suppressAt: CaretBox<KtAnnotatedExpression>, id: String) {
         val entry = findSuppressAnnotation(suppressAt.expression)
         if (entry != null) {
@@ -200,4 +181,4 @@ private class CaretBox<out E: KtExpression>(
         if (editor == null) return
         editor.caretModel.moveToOffset(copy.textOffset + offsetInExpression)
     }
-}
+}

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/modifierListModifactor.kt

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.util
+
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.js.descriptorUtils.getJetTypeFqName
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
+
+
+fun KtModifierListOwner.addAnnotation(
+        annotationFqName: FqName,
+        annotationInnerText: String? = null,
+        whiteSpaceText: String = "\n",
+        addToExistingAnnotation: ((KtAnnotationEntry) -> Boolean)? = null): Boolean {
+    val annotationText = when (annotationInnerText) {
+        null -> "@${annotationFqName.asString()}"
+        else -> "@${annotationFqName.asString()}($annotationInnerText)"
+    }
+
+    val psiFactory = KtPsiFactory(this)
+    val modifierList = modifierList
+
+    if (modifierList == null) {
+        // create a modifier list from scratch
+        val newModifierList = psiFactory.createModifierList(annotationText)
+        val replaced = KtPsiUtil.replaceModifierList(this, newModifierList)!!
+        val whiteSpace = psiFactory.createWhiteSpace(whiteSpaceText)
+        addAfter(whiteSpace, replaced)
+
+        ShortenReferences.DEFAULT.process(replaced)
+
+        return true
+    }
+
+    val entry = findAnnotation(annotationFqName)
+    if (entry == null) {
+        // no annotation
+        val newAnnotation = psiFactory.createAnnotationEntry(annotationText)
+        val addedAnnotation = modifierList.addBefore(newAnnotation, modifierList.firstChild) as KtElement
+        val whiteSpace = psiFactory.createWhiteSpace(whiteSpaceText)
+        modifierList.addAfter(whiteSpace, addedAnnotation)
+
+        ShortenReferences.DEFAULT.process(addedAnnotation)
+
+        return true
+    }
+
+    if (addToExistingAnnotation != null) {
+        return addToExistingAnnotation(entry)
+    }
+
+    return false
+}
+
+fun KtAnnotated.findAnnotation(annotationFqName: FqName): KtAnnotationEntry? {
+    val context = analyze()
+
+    for (entry in annotationEntries) {
+        val annotationDescriptor = context.get(BindingContext.ANNOTATION, entry)
+        if (annotationDescriptor != null) {
+            val fqName = annotationDescriptor.type.getJetTypeFqName(false)
+            if (fqName == annotationFqName.asString()) {
+                return entry
+            }
+        }
+    }
+    return null
+}
+