Extract common logic for annotation entries collection

Before this change annotation entries starting with '@' within expressions hadn't been resolved

Note that order of annotation entries in result may change according to order of children in PSI (see changed testData)

compiler/frontend/src/org/jetbrains/kotlin/psi/JetAnnotatedExpression.java

@@ -20,12 +20,13 @@ import com.intellij.lang.ASTNode;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.JetNodeTypes;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilPackage;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
-public class JetAnnotatedExpression extends JetExpressionImpl implements JetAnnotated {
+public class JetAnnotatedExpression extends JetExpressionImpl implements JetAnnotated, JetAnnotationsContainer {
     public JetAnnotatedExpression(@NotNull ASTNode node) {
         super(node);
     }
@@ -49,11 +50,6 @@ public class JetAnnotatedExpression extends JetExpressionImpl implements JetAnno
     @Override
     @NotNull
     public List<JetAnnotationEntry> getAnnotationEntries() {
-        List<JetAnnotationEntry> answer = null;
+        return PsiUtilPackage.collectAnnotationEntriesFromStubOrPsi(this);
-        for (JetAnnotation annotation : getAnnotations()) {
-            if (answer == null) answer = new ArrayList<JetAnnotationEntry>();
-            answer.addAll(annotation.getEntries());
-        }
-        return answer != null ? answer : Collections.<JetAnnotationEntry>emptyList();
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/psi/JetAnnotationsContainer.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2015 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.psi;
+
+public interface JetAnnotationsContainer extends JetElement {
+}

compiler/frontend/src/org/jetbrains/kotlin/psi/JetFileAnnotationList.java

@@ -17,15 +17,14 @@
 package org.jetbrains.kotlin.psi;
 
 import com.intellij.lang.ASTNode;
-import kotlin.Function1;
-import kotlin.KotlinPackage;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilPackage;
 import org.jetbrains.kotlin.psi.stubs.KotlinPlaceHolderStub;
 import org.jetbrains.kotlin.psi.stubs.elements.JetStubElementTypes;
 
 import java.util.List;
 
-public class JetFileAnnotationList extends JetElementImplStub<KotlinPlaceHolderStub<JetFileAnnotationList>> {
+public class JetFileAnnotationList extends JetElementImplStub<KotlinPlaceHolderStub<JetFileAnnotationList>> implements JetAnnotationsContainer {
 
     public JetFileAnnotationList(@NotNull ASTNode node) {
         super(node);
@@ -47,11 +46,6 @@ public class JetFileAnnotationList extends JetElementImplStub<KotlinPlaceHolderS
 
     @NotNull
     public List<JetAnnotationEntry> getAnnotationEntries() {
-        return KotlinPackage.flatMap(getAnnotations(), new Function1<JetAnnotation, List<JetAnnotationEntry>>() {
+        return PsiUtilPackage.collectAnnotationEntriesFromStubOrPsi(this);
-            @Override
-            public List<JetAnnotationEntry> invoke(JetAnnotation annotation) {
-                return annotation.getEntries();
-            }
-        });
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/psi/JetModifierList.java

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.psi;
 
-import com.google.common.collect.Lists;
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.stubs.IStubElementType;
@@ -24,14 +23,13 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.lexer.JetModifierKeywordToken;
 import org.jetbrains.kotlin.lexer.JetToken;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilPackage;
 import org.jetbrains.kotlin.psi.stubs.KotlinModifierListStub;
 import org.jetbrains.kotlin.psi.stubs.elements.JetStubElementTypes;
 
-import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 
-public abstract class JetModifierList extends JetElementImplStub<KotlinModifierListStub> {
+public abstract class JetModifierList extends JetElementImplStub<KotlinModifierListStub> implements JetAnnotationsContainer {
 
     public JetModifierList(@NotNull KotlinModifierListStub stub, @NotNull IStubElementType nodeType) {
         super(stub, nodeType);
@@ -53,13 +51,7 @@ public abstract class JetModifierList extends JetElementImplStub<KotlinModifierL
 
     @NotNull
     public List<JetAnnotationEntry> getAnnotationEntries() {
-        List<JetAnnotationEntry> entries = getStubOrPsiChildrenAsList(JetStubElementTypes.ANNOTATION_ENTRY);
+        return PsiUtilPackage.collectAnnotationEntriesFromStubOrPsi(this);
-        List<JetAnnotationEntry> answer = entries.isEmpty() ? null : Lists.newArrayList(entries);
-        for (JetAnnotation annotation : getAnnotations()) {
-            if (answer == null) answer = new ArrayList<JetAnnotationEntry>();
-            answer.addAll(annotation.getEntries());
-        }
-        return answer != null ? answer : Collections.<JetAnnotationEntry>emptyList();
     }
 
     public boolean hasModifier(@NotNull JetModifierKeywordToken token) {

compiler/frontend/src/org/jetbrains/kotlin/psi/JetTypeReference.java

@@ -20,6 +20,7 @@ import com.intellij.lang.ASTNode;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.lexer.JetTokens;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilPackage;
 import org.jetbrains.kotlin.psi.stubs.KotlinPlaceHolderStub;
 import org.jetbrains.kotlin.psi.stubs.elements.JetStubElementTypes;
 
@@ -33,7 +34,7 @@ import static org.jetbrains.kotlin.psi.stubs.elements.JetStubElementTypes.ANNOTA
  * Type reference element.
  * Underlying token is {@link org.jetbrains.kotlin.JetNodeTypes#TYPE_REFERENCE}
  */
-public class JetTypeReference extends JetElementImplStub<KotlinPlaceHolderStub<JetTypeReference>> {
+public class JetTypeReference extends JetElementImplStub<KotlinPlaceHolderStub<JetTypeReference>> implements JetAnnotationsContainer {
 
     public JetTypeReference(@NotNull ASTNode node) {
         super(node);
@@ -48,23 +49,13 @@ public class JetTypeReference extends JetElementImplStub<KotlinPlaceHolderStub<J
         return visitor.visitTypeReference(this, data);
     }
 
-    @NotNull
-    public List<JetAnnotation> getAttributeAnnotations() {
-        return getStubOrPsiChildrenAsList(ANNOTATION);
-    }
-
     @Nullable
     public JetTypeElement getTypeElement() {
         return JetStubbedPsiUtil.getStubOrPsiChild(this, JetStubElementTypes.TYPE_ELEMENT_TYPES, JetTypeElement.ARRAY_FACTORY);
     }
 
     public List<JetAnnotationEntry> getAnnotations() {
-        List<JetAnnotationEntry> answer = null;
+        return PsiUtilPackage.collectAnnotationEntriesFromStubOrPsi(this);
-        for (JetAnnotation annotation : getAttributeAnnotations()) {
-            if (answer == null) answer = new ArrayList<JetAnnotationEntry>();
-            answer.addAll(annotation.getEntries());
-        }
-        return answer != null ? answer : Collections.<JetAnnotationEntry>emptyList();
     }
 
     public boolean hasParentheses() {

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/jetPsiUtil.kt

@@ -41,6 +41,8 @@ import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.resolve.calls.CallTransformer.CallForImplicitInvoke
 import com.intellij.openapi.util.TextRange
+import com.intellij.psi.stubs.StubElement
+import org.jetbrains.kotlin.JetNodeTypes
 import org.jetbrains.kotlin.name.FqName
 
 public fun JetCallElement.getCallNameExpression(): JetSimpleNameExpression? {
@@ -542,3 +544,29 @@ public fun PsiFile.getFqNameByDirectory(): FqName {
 }
 
 public fun JetFile.packageMatchesDirectory(): Boolean = getPackageFqName() == getFqNameByDirectory()
+
+public fun JetAnnotationsContainer.collectAnnotationEntriesFromStubOrPsi(): List<JetAnnotationEntry> =
+    when (this) {
+        is StubBasedPsiElementBase<*> -> getStub()?.collectAnnotationEntriesFromStubElement() ?: collectAnnotationEntriesFromPsi()
+        else -> collectAnnotationEntriesFromPsi()
+    }
+
+private fun StubElement<*>.collectAnnotationEntriesFromStubElement() =
+    getChildrenStubs().flatMap {
+        child ->
+        when (child.getStubType()) {
+            JetNodeTypes.ANNOTATION_ENTRY -> listOf(child.getPsi() as JetAnnotationEntry)
+            JetNodeTypes.ANNOTATION -> (child.getPsi() as JetAnnotation).getEntries()
+            else -> emptyList<JetAnnotationEntry>()
+        }
+    }
+
+private fun JetAnnotationsContainer.collectAnnotationEntriesFromPsi() =
+    getChildren().flatMap {
+        child ->
+        when (child) {
+            is JetAnnotationEntry -> listOf(child)
+            is JetAnnotation -> child.getEntries()
+            else -> emptyList<JetAnnotationEntry>()
+        }
+    }

compiler/testData/diagnostics/tests/annotations/atAnnotationResolve.kt

@@ -0,0 +1,14 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+annotation class Ann(val x: Int = 6)
+
+@Ann(1) [Ann(2)] @Ann(3) @private class A [Ann] () {
+    @Ann(x = 5) fun foo() {
+        1 + @Ann(1) 1 * @Ann(<!TYPE_MISMATCH!>""<!>) 6
+
+        @Ann fun local() {}
+    }
+
+    @Ann val x = 1
+
+    fun bar(x: @Ann(1) [Ann(2)] @Ann(3) Int) {}
+}

compiler/testData/diagnostics/tests/annotations/atAnnotationResolve.txt

@@ -0,0 +1,19 @@
+package
+
+Ann(x = IntegerValueType(1): IntegerValueType(1)) Ann(x = IntegerValueType(2): IntegerValueType(2)) Ann(x = IntegerValueType(3): IntegerValueType(3)) private final class A {
+    Ann() public constructor A()
+    Ann() internal final val x: kotlin.Int = 1
+    internal final fun bar(/*0*/ x: [Ann(x = IntegerValueType(1): IntegerValueType(1)) Ann(x = IntegerValueType(2): IntegerValueType(2)) Ann(x = IntegerValueType(3): IntegerValueType(3))] kotlin.Int): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    Ann(x = IntegerValueType(5): IntegerValueType(5)) internal final fun foo(): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+internal final annotation class Ann : kotlin.Annotation {
+    public constructor Ann(/*0*/ x: kotlin.Int = ...)
+    internal final val x: kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/testData/loadJava/compiledKotlinWithStdlib/platformNames/functionName.txt

@@ -3,7 +3,7 @@ package test
 test.A(s = "2": kotlin.String) internal var v: kotlin.Int
     kotlin.platform.platformName(name = "vget": kotlin.String) test.A(s = "3": kotlin.String) internal fun <get-v>(): kotlin.Int
     kotlin.platform.platformName(name = "vset": kotlin.String) test.A(s = "4": kotlin.String) internal fun <set-v>(/*0*/ <set-?>: kotlin.Int): kotlin.Unit
-test.A(s = "1": kotlin.String) kotlin.platform.platformName(name = "bar": kotlin.String) internal fun foo(): kotlin.String
+kotlin.platform.platformName(name = "bar": kotlin.String) test.A(s = "1": kotlin.String) internal fun foo(): kotlin.String
 
 internal final annotation class A : kotlin.Annotation {
     /*primary*/ public constructor A(/*0*/ s: kotlin.String)

compiler/tests/org/jetbrains/kotlin/checkers/JetDiagnosticsTestGenerated.java

@@ -639,6 +639,12 @@ public class JetDiagnosticsTestGenerated extends AbstractJetDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("atAnnotationResolve.kt")
+            public void testAtAnnotationResolve() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/annotations/atAnnotationResolve.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("BasicAnnotations.kt")
             public void testBasicAnnotations() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/annotations/BasicAnnotations.kt");