Light Classes: Constant expression evaluator for light annotation arguments

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/CliLightClassGenerationSupport.kt

@@ -33,10 +33,7 @@ import org.jetbrains.kotlin.descriptors.PackageViewDescriptor
 import org.jetbrains.kotlin.fileClasses.JvmFileClassUtil
 import org.jetbrains.kotlin.load.kotlin.PackagePartClassUtils
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtPsiUtil
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.lazy.KotlinCodeAnalyzer
 import org.jetbrains.kotlin.resolve.lazy.ResolveSessionUtils
@@ -141,6 +138,8 @@ class CliLightClassGenerationSupport(project: Project) : LightClassGenerationSup
         return bindingContext.get(BindingContext.DECLARATION_TO_DESCRIPTOR, declaration)
     }
 
+    override fun analyze(element: KtElement) = bindingContext
+
     override fun getFacadeClasses(facadeFqName: FqName, scope: GlobalSearchScope): Collection<PsiClass> {
         val filesForFacade = findFilesForFacade(facadeFqName, scope)
         if (filesForFacade.isEmpty()) return emptyList()

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.*
-import com.intellij.psi.PsiAnnotationOwner
-import com.intellij.psi.PsiElement
 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);
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/IDELightClassGenerationSupport.kt

@@ -44,6 +44,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getElementTextWithContext
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil
 import org.jetbrains.kotlin.resolve.lazy.NoDescriptorForDeclarationException
 import org.jetbrains.kotlin.resolve.lazy.ResolveSession
@@ -202,6 +203,8 @@ class IDELightClassGenerationSupport(private val project: Project) : LightClassG
         }
     }
 
+    override fun analyze(element: KtElement) = element.analyze(BodyResolveMode.PARTIAL)
+
     override fun getFacadeNames(packageFqName: FqName, scope: GlobalSearchScope): Collection<String> {
         val facadeFilesInPackage = KotlinFileFacadeClassByPackageIndex.getInstance().get(packageFqName.asString(), project, scope)
         return facadeFilesInPackage.map { it.javaFileFacadeFqName.shortName().asString() }

idea/src/META-INF/plugin.xml

@@ -640,6 +640,10 @@
     <testCreator language="kotlin" implementationClass="org.jetbrains.kotlin.idea.testIntegration.KotlinTestCreator"/>
     <testFinder implementation="org.jetbrains.kotlin.idea.testIntegration.KotlinTestFinder"/>
 
+    <constantExpressionEvaluator
+        language="kotlin"
+        implementationClass="org.jetbrains.kotlin.idea.KotlinLightConstantExpressionEvaluator"/>
+
     <intentionAction>
       <className>org.jetbrains.kotlin.idea.intentions.IfNullToElvisIntention</className>
       <category>Kotlin</category>

idea/src/org/jetbrains/kotlin/idea/KotlinLightConstantExpressionEvaluator.kt

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2010-2016 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.idea
+
+import com.intellij.psi.PsiConstantEvaluationHelper
+import com.intellij.psi.PsiElement
+import com.intellij.psi.impl.ConstantExpressionEvaluator
+import org.jetbrains.kotlin.asJava.KtLightAnnotation
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.resolve.DelegatingBindingTrace
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator as FrontendConstantExpressionEvaluator
+
+class KotlinLightConstantExpressionEvaluator : ConstantExpressionEvaluator {
+    override fun computeConstantExpression(expression: PsiElement, throwExceptionOnOverflow: Boolean): Any? {
+        return computeExpression(expression, throwExceptionOnOverflow, null)
+    }
+
+    override fun computeExpression(
+            expression: PsiElement,
+            throwExceptionOnOverflow: Boolean,
+            auxEvaluator: PsiConstantEvaluationHelper.AuxEvaluator?
+    ): Any? {
+        if (expression !is KtLightAnnotation.LightExpressionValue) return null
+        val expressionToCompute = expression.originalExpression ?: return null
+        val resolutionFacade = expressionToCompute.getResolutionFacade()
+        val evaluator = FrontendConstantExpressionEvaluator(resolutionFacade.moduleDescriptor.builtIns)
+        val evaluatorTrace = DelegatingBindingTrace(resolutionFacade.analyze(expressionToCompute), "Evaluating annotation argument")
+        val constant = evaluator.evaluateExpression(expressionToCompute, evaluatorTrace) ?: return null
+        if (constant.isError) return null
+        return constant.getValue(TypeUtils.NO_EXPECTED_TYPE)
+    }
+}