Light Classes: Constant expression evaluator for light annotation arguments

compiler/light-classes/src/org/jetbrains/kotlin/asJava/KtLightAnnotation.kt

@@ -17,17 +17,73 @@
 package org.jetbrains.kotlin.asJava
 
 import com.intellij.openapi.util.TextRange
-import com.intellij.psi.PsiAnnotation
-import com.intellij.psi.PsiAnnotationOwner
-import com.intellij.psi.PsiElement
+import com.intellij.psi.*
 import com.intellij.util.IncorrectOperationException
+import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.psi.KtAnnotationEntry
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtParameter
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.DefaultValueArgument
+import org.jetbrains.kotlin.resolve.calls.model.ExpressionValueArgument
+import org.jetbrains.kotlin.resolve.calls.model.VarargValueArgument
+import org.jetbrains.kotlin.resolve.source.getPsi
 
 class KtLightAnnotation(
         private val delegate: PsiAnnotation,
         private val originalElement: KtAnnotationEntry,
         private val owner: PsiAnnotationOwner
 ) : PsiAnnotation by delegate, KtLightElement<KtAnnotationEntry, PsiAnnotation> {
+    inner class LightExpressionValue(private val delegate: PsiExpression) : PsiAnnotationMemberValue, PsiExpression by delegate {
+        val originalExpression: KtExpression? by lazy {
+            val nameAndValue = getStrictParentOfType<PsiNameValuePair>() ?: return@lazy null
+            val annotationEntry = this@KtLightAnnotation.getOrigin()
+            val context = LightClassGenerationSupport.getInstance(project).analyze(annotationEntry)
+            val resolvedCall = annotationEntry.getResolvedCall(context) ?: return@lazy null
+            val annotationConstructor = resolvedCall.resultingDescriptor
+            val parameterName = nameAndValue.name ?: "value"
+            val parameter = annotationConstructor.valueParameters.singleOrNull { it.name.asString() == parameterName }
+                            ?: return@lazy null
+            val resolvedArgument = resolvedCall.valueArguments[parameter] ?: return@lazy null
+            when (resolvedArgument) {
+                is DefaultValueArgument -> {
+                    (parameter.source.getPsi() as? KtParameter)?.defaultValue
+                }
+
+                is ExpressionValueArgument -> {
+                    resolvedArgument.valueArgument?.getArgumentExpression()
+                }
+
+                is VarargValueArgument -> {
+                    val arrayInitializer = parent as? PsiArrayInitializerMemberValue ?: return@lazy null
+                    val exprIndex = arrayInitializer.initializers.indexOf(delegate as PsiAnnotationMemberValue)
+                    if (exprIndex < 0) return@lazy null
+                    resolvedArgument.arguments[exprIndex].getArgumentExpression()
+                }
+
+                else -> null
+            }
+        }
+
+        override fun getLanguage() = KotlinLanguage.INSTANCE
+    }
+
+    inner class LightArrayInitializerValue(private val delegate: PsiArrayInitializerMemberValue) : PsiArrayInitializerMemberValue by delegate {
+        private val _initializers by lazy { delegate.initializers.map { wrapAnnotationValue(it) }.toTypedArray() }
+
+        override fun getInitializers() = _initializers
+        override fun getLanguage() = KotlinLanguage.INSTANCE
+    }
+
+    private fun wrapAnnotationValue(value: PsiAnnotationMemberValue): PsiAnnotationMemberValue {
+        return when (value) {
+            is PsiExpression -> LightExpressionValue(value)
+            is PsiArrayInitializerMemberValue -> LightArrayInitializerValue(value)
+            else -> value
+        }
+    }
+
     override fun getDelegate() = delegate
     override fun getOrigin() = originalElement
 
@@ -36,6 +92,9 @@ class KtLightAnnotation(
 
     override fun getOwner() = owner
 
+    override fun findAttributeValue(name: String?) = delegate.findAttributeValue(name)?.let { wrapAnnotationValue(it) }
+    override fun findDeclaredAttributeValue(name: String?) = delegate.findDeclaredAttributeValue(name)?.let { wrapAnnotationValue(it) }
+
     override fun getText() = originalElement.text ?: ""
     override fun getTextRange() = originalElement.textRange ?: TextRange.EMPTY_RANGE
 

compiler/light-classes/src/org/jetbrains/kotlin/asJava/KtLightParameter.java

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.asJava;
 
 import com.intellij.lang.Language;
+import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.search.SearchScope;

compiler/light-classes/src/org/jetbrains/kotlin/asJava/LightClassGenerationSupport.java

@@ -27,7 +27,9 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.psi.KtClassOrObject;
 import org.jetbrains.kotlin.psi.KtDeclaration;
+import org.jetbrains.kotlin.psi.KtElement;
 import org.jetbrains.kotlin.psi.KtFile;
+import org.jetbrains.kotlin.resolve.BindingContext;
 
 import java.util.Collection;
 
@@ -72,6 +74,9 @@ public abstract class LightClassGenerationSupport {
     @Nullable
     public abstract DeclarationDescriptor resolveToDescriptor(@NotNull KtDeclaration declaration);
 
+    @NotNull
+    public abstract BindingContext analyze(@NotNull KtElement element);
+
     @NotNull
     public abstract Collection<PsiClass> getFacadeClasses(@NotNull FqName facadeFqName, @NotNull GlobalSearchScope scope);