From 9302beb5fa1eb4219795a3049d2fc9caacc176aa Mon Sep 17 00:00:00 2001
From: peter <peter@jetbrains.com>
Date: Thu, 15 Sep 2011 13:42:59 +0200
Subject: [PATCH] preserve annotations on groovy import statements while
 opimizing (IDEA-74323)

---
 .../lang/editor/GroovyImportOptimizer.java    | 23 +++++++++++++++++++
 .../OptimizeImportsTest.groovy                |  2 ++
 .../PreserveImportAnnotations.groovy          |  6 +++++
 .../PreserveImportAnnotations_after.groovy    |  5 ++++
 4 files changed, 36 insertions(+)
 create mode 100644 plugins/groovy/testdata/optimizeImports/PreserveImportAnnotations.groovy
 create mode 100644 plugins/groovy/testdata/optimizeImports/PreserveImportAnnotations_after.groovy

diff --git a/plugins/groovy/src/org/jetbrains/plugins/groovy/lang/editor/GroovyImportOptimizer.java b/plugins/groovy/src/org/jetbrains/plugins/groovy/lang/editor/GroovyImportOptimizer.java
index ae9c32ed6690..6c0729035cf6 100644
--- a/plugins/groovy/src/org/jetbrains/plugins/groovy/lang/editor/GroovyImportOptimizer.java
+++ b/plugins/groovy/src/org/jetbrains/plugins/groovy/lang/editor/GroovyImportOptimizer.java
@@ -30,6 +30,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.groovy.lang.psi.*;
 import org.jetbrains.plugins.groovy.lang.psi.api.GroovyResolveResult;
+import org.jetbrains.plugins.groovy.lang.psi.api.auxiliary.modifiers.GrModifierList;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrReferenceExpression;
 import org.jetbrains.plugins.groovy.lang.psi.api.toplevel.imports.GrImportStatement;
 import org.jetbrains.plugins.groovy.lang.psi.api.types.GrCodeReferenceElement;
@@ -97,12 +98,20 @@ public class GroovyImportOptimizer implements ImportOptimizer {
         return;
       }
 
+      Map<String, String> annotations = new HashMap<String, String>();
+      
       // Getting aliased imports
       GroovyPsiElementFactory factory = GroovyPsiElementFactory.getInstance(myFile.getProject());
       ArrayList<GrImportStatement> aliased = new ArrayList<GrImportStatement>();
       for (GrImportStatement oldImport : oldImports) {
         if (oldImport.isAliasedImport() && usedImports.contains(oldImport)) {
           aliased.add(factory.createImportStatementFromText(oldImport.getText()));
+        } else {
+          String importReference = getImportReferenceText(oldImport);
+          GrModifierList annotationList = oldImport.getAnnotationList();
+          if (importReference != null && annotationList != null) {
+            annotations.put(importReference, annotationList.getText());
+          }
         }
       }
 
@@ -116,6 +125,12 @@ public class GroovyImportOptimizer implements ImportOptimizer {
         myFile.addImport(aliasedImport);
       }
       for (GrImportStatement newImport : newImports) {
+        String imported = getImportReferenceText(newImport);
+        String annos = annotations.get(imported);
+        if (imported != null && StringUtil.isNotEmpty(annos)) {
+          newImport = factory.createImportStatementFromText(annos + " " + newImport.getText());
+        }
+        
         myFile.addImport(newImport);
       }
 
@@ -126,6 +141,14 @@ public class GroovyImportOptimizer implements ImportOptimizer {
       }
     }
 
+    @Nullable String getImportReferenceText(GrImportStatement statement) {
+      GrCodeReferenceElement importReference = statement.getImportReference();
+      if (importReference != null) {
+        return statement.getText().substring(importReference.getStartOffsetInParent());
+      }
+      return null;
+    }
+
     public List<GrImportStatement> findUnusedImports(final Set<String> importedClasses,
                                                      final Set<String> staticallyImportedMembers,
                                                      final Set<GrImportStatement> usedImports,
diff --git a/plugins/groovy/test/org/jetbrains/plugins/groovy/refactoring/optimizeImports/OptimizeImportsTest.groovy b/plugins/groovy/test/org/jetbrains/plugins/groovy/refactoring/optimizeImports/OptimizeImportsTest.groovy
index dffc2e9911e2..1fcda5dfba14 100644
--- a/plugins/groovy/test/org/jetbrains/plugins/groovy/refactoring/optimizeImports/OptimizeImportsTest.groovy
+++ b/plugins/groovy/test/org/jetbrains/plugins/groovy/refactoring/optimizeImports/OptimizeImportsTest.groovy
@@ -103,6 +103,8 @@ public class OptimizeImportsTest extends LightCodeInsightFixtureTestCase {
     doTest();
   }
 
+  public void testPreserveImportAnnotations() throws Throwable { doTest(); }
+
   public void testUtilListMasked() throws Throwable {
     myFixture.addClass("package java.awt; public class List {}");
     doTest();
diff --git a/plugins/groovy/testdata/optimizeImports/PreserveImportAnnotations.groovy b/plugins/groovy/testdata/optimizeImports/PreserveImportAnnotations.groovy
new file mode 100644
index 000000000000..9a438b12e91d
--- /dev/null
+++ b/plugins/groovy/testdata/optimizeImports/PreserveImportAnnotations.groovy
@@ -0,0 +1,6 @@
+@Anno import javax.swing.*
+import java.sql.Blob
+
+def frame = new JFrame()
+def table = new JTable()
+def dialog = new JDialog()
\ No newline at end of file
diff --git a/plugins/groovy/testdata/optimizeImports/PreserveImportAnnotations_after.groovy b/plugins/groovy/testdata/optimizeImports/PreserveImportAnnotations_after.groovy
new file mode 100644
index 000000000000..2008cb926869
--- /dev/null
+++ b/plugins/groovy/testdata/optimizeImports/PreserveImportAnnotations_after.groovy
@@ -0,0 +1,5 @@
+@Anno import javax.swing.*
+
+def frame = new JFrame()
+def table = new JTable()
+def dialog = new JDialog()
\ No newline at end of file