From 98be3a370a75d2ebcd8029d6640ba7d1e0067ba0 Mon Sep 17 00:00:00 2001
From: Max Medvedev <maxim.medvedev@jetbrains.com>
Date: Mon, 9 Dec 2013 18:38:20 +0400
Subject: [PATCH] IDEA-117584 more stubs for annotations. Optimization for
 argument searching

---
 .../annotation/GrAnnotationImpl.java          | 42 ++++++++++++-------
 1 file changed, 27 insertions(+), 15 deletions(-)

diff --git a/plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/impl/auxiliary/annotation/GrAnnotationImpl.java b/plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/impl/auxiliary/annotation/GrAnnotationImpl.java
index f82b8a08cd81..c672ca128c8c 100644
--- a/plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/impl/auxiliary/annotation/GrAnnotationImpl.java
+++ b/plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/impl/auxiliary/annotation/GrAnnotationImpl.java
@@ -23,6 +23,7 @@ import com.intellij.psi.*;
 import com.intellij.psi.impl.PsiImplUtil;
 import com.intellij.psi.impl.light.LightClassReference;
 import com.intellij.psi.meta.PsiMetaData;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.PairFunction;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
@@ -82,10 +83,6 @@ public class GrAnnotationImpl extends GrStubElementBase<GrAnnotationStub> implem
 
   @NotNull
   public GrAnnotationArgumentList getParameterList() {
-    final GrAnnotationStub stub = getStub();
-    if (stub != null) {
-      return stub.getPsiElement().getParameterList();
-    }
     return findNotNullChildByClass(GrAnnotationArgumentList.class);
   }
 
@@ -105,29 +102,44 @@ public class GrAnnotationImpl extends GrStubElementBase<GrAnnotationStub> implem
 
   @Nullable
   public PsiJavaCodeReferenceElement getNameReferenceElement() {
-    final GrAnnotationStub stub = getStub();
-    if (stub != null) {
-      return stub.getPsiElement().getNameReferenceElement();
-    }
+    final GroovyResolveResult resolveResult = resolveWithStub();
 
-    final GroovyResolveResult resolveResult = getClassReference().advancedResolve();
     final PsiElement resolved = resolveResult.getElement();
+    if (!(resolved instanceof PsiClass)) return null;
 
-    if (resolved instanceof PsiClass) {
-      return new LightClassReference(getManager(), getClassReference().getText(), (PsiClass)resolved, resolveResult.getSubstitutor());
-    }
-    else {
-      return null;
-    }
+    return new LightClassReference(getManager(), getClassReference().getText(), (PsiClass)resolved, resolveResult.getSubstitutor());
+  }
+
+  @NotNull
+  private GroovyResolveResult resolveWithStub() {
+    final GrAnnotationStub stub = getStub();
+    final GrCodeReferenceElement reference = stub != null ? stub.getPsiElement().getClassReference() : getClassReference();
+    return reference.advancedResolve();
   }
 
   @Nullable
   public PsiAnnotationMemberValue findAttributeValue(@Nullable String attributeName) {
+    final GrAnnotationStub stub = getStub();
+    if (stub != null) {
+      final GrAnnotation stubbedPsi = stub.getPsiElement();
+      final PsiAnnotationMemberValue value = PsiImplUtil.findAttributeValue(stubbedPsi, attributeName);
+      if (value == null || !PsiTreeUtil.isAncestor(stubbedPsi, value, true)) {         // if value is a default value we can use it
+        return value;
+      }
+    }
     return PsiImplUtil.findAttributeValue(this, attributeName);
   }
 
   @Nullable
   public PsiAnnotationMemberValue findDeclaredAttributeValue(@NonNls final String attributeName) {
+    final GrAnnotationStub stub = getStub();
+    if (stub != null) {
+      final GrAnnotation stubbedPsi = stub.getPsiElement();
+      final PsiAnnotationMemberValue value = PsiImplUtil.findDeclaredAttributeValue(stubbedPsi, attributeName);
+      if (value == null) {
+        return null;
+      }
+    }
     return PsiImplUtil.findDeclaredAttributeValue(this, attributeName);
   }