Refactor JavaEnumValueAnnotationArgument and implementations

Only require implementations to be able to compute the enum class name
and the corresponding entry name. Only this should be necessary to
correctly resolve the argument; the resolvers will find the
corresponding class descriptor if necessary

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/structure/impl/annotationArgumentsImpl.kt

@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.load.java.structure.impl
 
 import com.intellij.psi.*
 import org.jetbrains.kotlin.load.java.structure.*
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
 abstract class JavaAnnotationArgumentImpl(
@@ -64,14 +66,17 @@ class JavaEnumValueAnnotationArgumentImpl(
         private val psiReference: PsiReferenceExpression,
         name: Name?
 ) : JavaAnnotationArgumentImpl(name), JavaEnumValueAnnotationArgument {
-    override fun resolve(): JavaField? {
-        val element = psiReference.resolve()
-        return when (element) {
-            null -> null
-            is PsiEnumConstant -> JavaFieldImpl(element)
-            else -> throw IllegalStateException("Reference argument should be an enum value, but was $element: ${element.text}")
+    override val enumClassId: ClassId?
+        get() {
+            val element = psiReference.resolve()
+            if (element is PsiEnumConstant) {
+                return JavaFieldImpl(element).containingClass.classId
+            }
+
+            val fqName = (psiReference.qualifier as? PsiReferenceExpression)?.qualifiedName ?: return null
+            // TODO: find a way to construct a correct name (with nested classes) for unresolved enums
+            return ClassId.topLevel(FqName(fqName))
         }
-    }
 
     override val entryName: Name?
         get() = psiReference.referenceName?.let(Name::identifier)

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/structure/impl/classFiles/Annotations.kt

@@ -127,7 +127,7 @@ class BinaryJavaAnnotation private constructor(
     }
 
     override val classId: ClassId?
-        get() = classifierResolutionResult.classifier.safeAs<JavaClass>()?.classId()
+        get() = classifierResolutionResult.classifier.safeAs<JavaClass>()?.classId
                 ?: ClassId.topLevel(FqName(classifierResolutionResult.qualifiedName))
 
     override fun resolve() = classifierResolutionResult.classifier as? JavaClass
@@ -152,7 +152,7 @@ class BinaryJavaAnnotationVisitor(
     }
 
     override fun visitEnum(name: String?, desc: String, value: String) {
-         addArgument(PlainJavaEnumValueAnnotationArgument(name, desc, value, context))
+        addArgument(PlainJavaEnumValueAnnotationArgument(name, context.mapDescToClassId(desc), value))
     }
 
     override fun visit(name: String?, value: Any?) {
@@ -216,23 +216,8 @@ class PlainJavaAnnotationAsAnnotationArgument(
 
 class PlainJavaEnumValueAnnotationArgument(
         name: String?,
-        private val desc: String,
-        entryName: String,
-        private val context: ClassifierResolutionContext
+        override val enumClassId: ClassId,
+        entryName: String
 ) : PlainJavaAnnotationArgument(name), JavaEnumValueAnnotationArgument {
     override val entryName = Name.identifier(entryName)
-
-    override fun resolve(): JavaField? {
-        val javaClass = context.resolveByInternalName(Type.getType(desc).internalName).classifier as? JavaClass ?: return null
-        return javaClass.fields.singleOrNull { it.name == entryName }
-    }
-}
-
-private fun JavaClass.classId(): ClassId? {
-    val fqName = fqName ?: return null
-    if (outerClass == null) return ClassId.topLevel(fqName)
-
-    val outerClassId = outerClass!!.classId() ?: return null
-
-    return ClassId(outerClassId.packageFqName, outerClassId.relativeClassName.child(name), false)
 }

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/structure/impl/classFiles/ClassifierResolutionContext.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.load.java.structure.JavaTypeParameter
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.org.objectweb.asm.Type
 
 typealias ClassIdToJavaClass = (ClassId) -> JavaClass?
 
@@ -117,4 +118,6 @@ class ClassifierResolutionContext private constructor(
     private fun createClassIdForTopLevel(internalName: String) = ClassId.topLevel(FqName(internalName.replace('/', '.')))
 
     internal fun resolveByInternalName(c: String) = resolveClass(mapInternalNameToClassId(c))
+
+    internal fun mapDescToClassId(desc: String): ClassId = mapInternalNameToClassId(Type.getType(desc).internalName)
 }

compiler/javac-wrapper/src/org/jetbrains/kotlin/javac/wrappers/symbols/symbolBasedAnnotationArguments.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.javac.wrappers.symbols
 import com.sun.tools.javac.code.Symbol
 import org.jetbrains.kotlin.javac.JavacWrapper
 import org.jetbrains.kotlin.load.java.structure.*
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import javax.lang.model.element.AnnotationMirror
 import javax.lang.model.element.AnnotationValue
@@ -63,16 +64,17 @@ class SymbolBasedReferenceAnnotationArgument(
         name: Name,
         javac: JavacWrapper
 ) : SymbolBasedAnnotationArgument(name, javac), JavaEnumValueAnnotationArgument {
-    override val entryName: Name?
-        get() = name
-
-    override fun resolve(): SymbolBasedField {
-        val containingClass = (element.enclosingElement as Symbol.ClassSymbol).let {
-            SymbolBasedClass(it, javac, null, it.classfile)
-        }
-        return SymbolBasedField(element, containingClass, javac)
+    // TODO: do not create extra objects here
+    private val javaField: JavaField? by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        val containingClass = element.enclosingElement as Symbol.ClassSymbol
+        SymbolBasedField(element, SymbolBasedClass(containingClass, javac, null, containingClass.classfile), javac)
     }
 
+    override val enumClassId: ClassId?
+        get() = javaField?.containingClass?.classId
+
+    override val entryName: Name?
+        get() = javaField?.name
 }
 
 class SymbolBasedClassObjectAnnotationArgument(

compiler/javac-wrapper/src/org/jetbrains/kotlin/javac/wrappers/trees/TreeBasedAnnotation.kt

@@ -50,21 +50,26 @@ class TreeBasedLiteralAnnotationArgument(name: Name,
                                          override val value: Any?,
                                          javac: JavacWrapper) : TreeBasedAnnotationArgument(name, javac), JavaLiteralAnnotationArgument
 
-class TreeBasedReferenceAnnotationArgument(name: Name,
-                                           private val compilationUnit: CompilationUnitTree,
-                                           private val field: JCTree.JCFieldAccess,
-                                           javac: JavacWrapper,
-                                           private val onElement: JavaElement) : TreeBasedAnnotationArgument(name, javac), JavaEnumValueAnnotationArgument {
-    override val entryName: Name?
-        get() = name
+class TreeBasedReferenceAnnotationArgument(
+        name: Name,
+        private val compilationUnit: CompilationUnitTree,
+        private val field: JCTree.JCFieldAccess,
+        javac: JavacWrapper,
+        private val onElement: JavaElement
+) : TreeBasedAnnotationArgument(name, javac), JavaEnumValueAnnotationArgument {
+    // TODO: do not run resolve here
+    private val javaField: JavaField? by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        val javaClass = javac.resolve(field.selected, compilationUnit, onElement) as? JavaClass
+        val fieldName = Name.identifier(field.name.toString())
 
-    override fun resolve(): JavaField? {
-        val javaClass = javac.resolve(field.selected, compilationUnit, onElement) as? JavaClass ?: return null
-        val fieldName = field.name.toString().let { Name.identifier(it) }
-
-        return javaClass.fields.find { it.name == fieldName }
+        javaClass?.fields?.find { it.name == fieldName }
     }
 
+    override val enumClassId: ClassId?
+        get() = javaField?.containingClass?.classId
+
+    override val entryName: Name?
+        get() = javaField?.name
 }
 
 class TreeBasedArrayAnnotationArgument(val args: List<JavaAnnotationArgument>,

compiler/javac-wrapper/src/org/jetbrains/kotlin/javac/wrappers/trees/utils.kt

@@ -20,7 +20,6 @@ import com.sun.tools.javac.tree.JCTree
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.javac.wrappers.symbols.SymbolBasedArrayAnnotationArgument
-import org.jetbrains.kotlin.javac.wrappers.symbols.SymbolBasedField
 import org.jetbrains.kotlin.javac.wrappers.symbols.SymbolBasedReferenceAnnotationArgument
 import org.jetbrains.kotlin.load.java.JavaVisibilities
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotation
@@ -81,18 +80,16 @@ fun Collection<JavaAnnotation>.filterTypeAnnotations(): Collection<JavaAnnotatio
             }
             is TreeBasedArrayAnnotationArgument -> {
                 elementTypeArg.args.find {
-                    val field = (it as? TreeBasedReferenceAnnotationArgument)?.resolve() as? SymbolBasedField
-                    field?.element?.simpleName?.contentEquals("TYPE_USE") ?: false
+                    (it as? TreeBasedReferenceAnnotationArgument)?.entryName?.asString() == "TYPE_USE"
                 }?.let { filteredAnnotations.add(annotation) }
             }
             is TreeBasedReferenceAnnotationArgument -> {
-                (elementTypeArg.resolve() as? SymbolBasedField)?.let { field ->
-                    field.element.simpleName.takeIf { it.contentEquals("TYPE_USE") }
-                            ?.let { filteredAnnotations.add(annotation) }
+                if (elementTypeArg.entryName?.asString() == "TYPE_USE") {
+                    filteredAnnotations.add(annotation)
                 }
             }
         }
     }
 
     return filteredAnnotations
-}
+}