FIR LC: handle more kinds of annotation values

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/utils/firUtils.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeNullability
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.symbols.impl.FirEnumEntrySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.name.FqName
@@ -108,7 +109,8 @@ internal fun Map<String, FirExpression>.toNamedConstantValue(
         )
     }
 
-internal fun <T> FirConstExpression<T>.convertConstantExpression(): KtSimpleConstantValue<T> = KtSimpleConstantValue(kind, value)
+internal fun <T> FirConstExpression<T>.convertConstantExpression(): KtSimpleConstantValue<T> =
+    KtSimpleConstantValue(kind, value, realPsi as? KtElement)
 
 private fun Collection<FirExpression>.convertConstantExpression(
     session: FirSession,
@@ -116,11 +118,11 @@ private fun Collection<FirExpression>.convertConstantExpression(
 ): Collection<KtConstantValue> =
     mapNotNull { it.convertConstantExpression(session, firSymbolBuilder) }
 
-private fun Collection<KtConstantValue>.toArrayConstantValueIfNecessary(): KtConstantValue {
+private fun Collection<KtConstantValue>.toArrayConstantValueIfNecessary(kotlinOrigin: KtElement?): KtConstantValue {
     return if (size == 1)
         single()
     else
-        KtArrayConstantValue(this)
+        KtArrayConstantValue(this, kotlinOrigin)
 }
 
 internal fun FirExpression.convertConstantExpression(
@@ -133,10 +135,12 @@ internal fun FirExpression.convertConstantExpression(
             expression.convertConstantExpression(session, firSymbolBuilder)
         }
         is FirVarargArgumentsExpression -> {
-            arguments.convertConstantExpression(session, firSymbolBuilder).toArrayConstantValueIfNecessary()
+            arguments.convertConstantExpression(session, firSymbolBuilder)
+                .toArrayConstantValueIfNecessary(realPsi as? KtElement)
         }
         is FirArrayOfCall -> {
-            argumentList.arguments.convertConstantExpression(session, firSymbolBuilder).toArrayConstantValueIfNecessary()
+            argumentList.arguments.convertConstantExpression(session, firSymbolBuilder)
+                .toArrayConstantValueIfNecessary(realPsi as? KtElement)
         }
         is FirFunctionCall -> {
             val reference = calleeReference as? FirResolvedNamedReference ?: return null
@@ -150,13 +154,15 @@ internal fun FirExpression.convertConstantExpression(
                         }
                         KtAnnotationConstantValue(
                             resolvedSymbol.callableId.className?.asString(),
-                            resultMap.toNamedConstantValue(session, firSymbolBuilder)
+                            resultMap.toNamedConstantValue(session, firSymbolBuilder),
+                            this.realPsi as? KtCallElement
                         )
                     } else null
                 }
                 is FirNamedFunctionSymbol -> {
                     if (resolvedSymbol.callableId.asSingleFqName() in ArrayFqNames.ARRAY_CALL_FQ_NAMES)
-                        argumentList.arguments.convertConstantExpression(session, firSymbolBuilder).toArrayConstantValueIfNecessary()
+                        argumentList.arguments.convertConstantExpression(session, firSymbolBuilder)
+                            .toArrayConstantValueIfNecessary(realPsi as? KtElement)
                     else null
                 }
                 else -> null
@@ -166,7 +172,7 @@ internal fun FirExpression.convertConstantExpression(
             val reference = calleeReference as? FirResolvedNamedReference ?: return null
             when (val resolvedSymbol = reference.resolvedSymbol) {
                 is FirEnumEntrySymbol -> {
-                    KtEnumEntryValue(resolvedSymbol.fir.buildSymbol(firSymbolBuilder) as KtEnumEntrySymbol)
+                    KtEnumEntryValue(resolvedSymbol.fir.buildSymbol(firSymbolBuilder) as KtEnumEntrySymbol, realPsi as? KtElement)
                 }
                 else -> null
             }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/symbols/markers/KtConstantValue.kt

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.analysis.api.symbols.markers
 
 import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
+import org.jetbrains.kotlin.psi.KtCallElement
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.types.ConstantValueKind
 
 /**
@@ -23,23 +25,32 @@ import org.jetbrains.kotlin.types.ConstantValueKind
  *  [KtAnnotationConstantValue] represents annotation types (with annotation fq name and arguments); and
  *  [KtArrayConstantValue] abstracts an array of [KtConstantValue]s.
  */
-public sealed class KtConstantValue
+public sealed class KtConstantValue(
+    public open val kotlinOrigin: KtElement? = null
+)
 public object KtUnsupportedConstantValue : KtConstantValue()
 
 public class KtArrayConstantValue(
-    public val values: Collection<KtConstantValue>
+    public val values: Collection<KtConstantValue>,
+    override val kotlinOrigin: KtElement?,
 ) : KtConstantValue()
 
 public class KtAnnotationConstantValue(
     public val fqName: String?,
-    public val arguments: List<KtNamedConstantValue>
+    public val arguments: List<KtNamedConstantValue>,
+    override val kotlinOrigin: KtCallElement?,
 ) : KtConstantValue()
 
 public class KtEnumEntryValue(
-    public val enumEntrySymbol: KtEnumEntrySymbol
+    public val enumEntrySymbol: KtEnumEntrySymbol,
+    override val kotlinOrigin: KtElement?,
 ) : KtConstantValue()
 
-public data class KtSimpleConstantValue<T>(val constantValueKind: ConstantValueKind<T>, val value: T) : KtConstantValue() {
+public class KtSimpleConstantValue<T>(
+    public val constantValueKind: ConstantValueKind<T>,
+    public val value: T,
+    override val kotlinOrigin: KtElement?,
+) : KtConstantValue() {
     public fun toConst(): Any? {
         return (value as? Long)?.let {
             when (constantValueKind) {
@@ -57,4 +68,8 @@ public data class KtSimpleConstantValue<T>(val constantValueKind: ConstantValueK
             }
         } ?: value
     }
+
+    override fun toString(): String {
+        return "KtSimpleConstantValue(constantValueKind=${constantValueKind}, value=${value})"
+    }
 }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/FirAnnotationParameterList.kt

@@ -17,11 +17,11 @@ import org.jetbrains.kotlin.psi.KtElement
 
 internal class FirAnnotationParameterList(
     parent: FirLightAbstractAnnotation,
-    private val annotationCall: KtAnnotationCall,
+    private val arguments: List<KtNamedConstantValue>,
 ) : KtLightElementBase(parent), PsiAnnotationParameterList {
 
     private val _attributes: Array<PsiNameValuePair> by lazyPub {
-        annotationCall.arguments.map {
+        arguments.map {
             FirNameValuePairForAnnotationArgument(it, this)
         }.toTypedArray()
     }
@@ -41,7 +41,7 @@ private class FirNameValuePairForAnnotationArgument(
     override val kotlinOrigin: KtElement? get() = null
 
     private val _value by lazyPub {
-        (constantValue.expression as? KtSimpleConstantValue<*>)?.createPsiLiteral(this)
+        constantValue.expression.toAnnotationMemberValue(this)
     }
 
     override fun setValue(p0: PsiAnnotationMemberValue) = cannotModify()
@@ -57,4 +57,4 @@ private class FirNameValuePairForAnnotationArgument(
     override fun getLiteralValue(): String? = (value as? PsiLiteralExpression)?.value?.toString()
 
     override fun getName(): String = constantValue.name
-}
+}

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/FirAnnotationValues.kt

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.light.classes.symbol
+
+import com.intellij.psi.*
+import org.jetbrains.kotlin.asJava.elements.KtLightElementBase
+import org.jetbrains.kotlin.psi.KtElement
+
+internal class FirPsiArrayInitializerMemberValue(
+    override val kotlinOrigin: KtElement?,
+    private val lightParent: PsiElement,
+    private val arguments: (FirPsiArrayInitializerMemberValue) -> List<PsiAnnotationMemberValue>
+) : KtLightElementBase(lightParent), PsiArrayInitializerMemberValue {
+
+    override fun getInitializers(): Array<PsiAnnotationMemberValue> = arguments(this).toTypedArray()
+
+    override fun getParent(): PsiElement = lightParent
+    override fun isPhysical(): Boolean = false
+
+    override fun getText(): String = "{" + initializers.joinToString { it.text } + "}"
+}
+
+internal abstract class FirPsiAnnotationMemberValue(
+    override val kotlinOrigin: KtElement?,
+    private val lightParent: PsiElement,
+) : KtLightElementBase(lightParent), PsiAnnotationMemberValue {
+
+    override fun getParent(): PsiElement = lightParent
+    override fun isPhysical(): Boolean = false
+}
+
+internal class FirPsiExpression(
+    override val kotlinOrigin: KtElement?,
+    lightParent: PsiElement,
+    private val psiExpression: PsiExpression,
+) : FirPsiAnnotationMemberValue(kotlinOrigin, lightParent), PsiExpression {
+    override fun getType(): PsiType? = psiExpression.type
+    override fun getText(): String = psiExpression.text
+}
+
+internal class FirPsiLiteral(
+    override val kotlinOrigin: KtElement?,
+    lightParent: PsiElement,
+    private val psiLiteral: PsiLiteral,
+) : FirPsiAnnotationMemberValue(kotlinOrigin, lightParent), PsiLiteral {
+    override fun getValue(): Any? = psiLiteral.value
+    override fun getText(): String = psiLiteral.text
+}

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/FirLightAnnotationForAnnotationCall.kt

@@ -26,7 +26,7 @@ internal class FirLightAnnotationForAnnotationCall(
         PsiImplUtil.findDeclaredAttributeValue(this, attributeName)
 
     private val _parameterList: PsiAnnotationParameterList by lazyPub {
-        FirAnnotationParameterList(this@FirLightAnnotationForAnnotationCall, annotationCall)
+        FirAnnotationParameterList(this@FirLightAnnotationForAnnotationCall, annotationCall.arguments)
     }
 
     override fun getParameterList(): PsiAnnotationParameterList = _parameterList
@@ -47,4 +47,4 @@ internal class FirLightAnnotationForAnnotationCall(
 
     override fun isEquivalentTo(another: PsiElement?): Boolean =
         basicIsEquivalentTo(this, another as? PsiAnnotation)
-}
+}

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/FirLightSimpleAnnotation.kt

@@ -6,14 +6,19 @@
 package org.jetbrains.kotlin.light.classes.symbol
 
 import com.intellij.psi.PsiAnnotationMemberValue
+import com.intellij.psi.PsiAnnotationParameterList
 import com.intellij.psi.PsiElement
+import com.intellij.psi.impl.PsiImplUtil
+import org.jetbrains.kotlin.analysis.api.symbols.markers.KtNamedConstantValue
+import org.jetbrains.kotlin.asJava.classes.lazyPub
 import org.jetbrains.kotlin.psi.KtCallElement
 
 internal class FirLightSimpleAnnotation(
     private val fqName: String?,
-    parent: PsiElement
+    parent: PsiElement,
+    private val arguments: List<KtNamedConstantValue> = listOf(),
+    override val kotlinOrigin: KtCallElement? = null,
 ) : FirLightAbstractAnnotation(parent) {
-    override val kotlinOrigin: KtCallElement? = null
 
     override fun getQualifiedName(): String? = fqName
 
@@ -25,7 +30,16 @@ internal class FirLightSimpleAnnotation(
 
     override fun hashCode(): Int = fqName.hashCode()
 
-    override fun findAttributeValue(attributeName: String?): PsiAnnotationMemberValue? = null
+    override fun findAttributeValue(attributeName: String?): PsiAnnotationMemberValue? =
+        PsiImplUtil.findAttributeValue(this, attributeName)
+
+    override fun findDeclaredAttributeValue(attributeName: String?) =
+        PsiImplUtil.findDeclaredAttributeValue(this, attributeName)
+
+    private val _parameterList: PsiAnnotationParameterList by lazyPub {
+        FirAnnotationParameterList(this@FirLightSimpleAnnotation, arguments)
+    }
+
+    override fun getParameterList(): PsiAnnotationParameterList = _parameterList
 
-    override fun findDeclaredAttributeValue(attributeName: String?): PsiAnnotationMemberValue? = null
 }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/firLightUtils.kt

@@ -16,10 +16,7 @@ import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.components.DefaultTypeClassIds
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtFunctionSymbol
-import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSimpleConstantValue
+import org.jetbrains.kotlin.analysis.api.symbols.markers.*
-import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithModality
-import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithVisibility
-import org.jetbrains.kotlin.analysis.api.symbols.markers.KtTypeAndAnnotations
 import org.jetbrains.kotlin.analysis.api.types.*
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.name.ClassId
@@ -182,16 +179,41 @@ private fun escapeString(str: String): String = buildString {
     }
 }
 
-private fun KtSimpleConstantValue<*>.asStringForPsiLiteral(): String =
+internal fun KtConstantValue.toAnnotationMemberValue(parent: PsiElement): PsiAnnotationMemberValue? {
-    when (val value = this.value) {
+    return when (this) {
-        is String -> "\"${escapeString(value)}\""
+        is KtArrayConstantValue ->
-        is Long -> "${value}L"
+            FirPsiArrayInitializerMemberValue(kotlinOrigin, parent) { arrayLiteralParent ->
-        is Float -> "${value}f"
+                values.mapNotNull { element -> element.toAnnotationMemberValue(arrayLiteralParent) }
-        else -> value?.toString() ?: "null"
+            }
+        is KtAnnotationConstantValue ->
+            FirLightSimpleAnnotation(fqName, parent, arguments, kotlinOrigin)
+        else ->
+            createPsiLiteral(parent)?.let {
+                when (it) {
+                    is PsiLiteral -> FirPsiLiteral(kotlinOrigin, parent, it)
+                    else -> FirPsiExpression(kotlinOrigin, parent, it)
+                }
+            }
+    }
+}
+
+private fun KtConstantValue.asStringForPsiLiteral(): String? =
+    when (this) {
+        is KtEnumEntryValue ->
+            "${enumEntrySymbol.containingEnumClassIdIfNonLocal?.asSingleFqName()?.asString() ?: ""}.${enumEntrySymbol.name}"
+        is KtSimpleConstantValue<*> -> {
+            when (val value = this.value) {
+                is String -> "\"${escapeString(value)}\""
+                is Long -> "${value}L"
+                is Float -> "${value}f"
+                else -> value?.toString() ?: "null"
+            }
+        }
+        else -> null
     }
 
-internal fun KtSimpleConstantValue<*>.createPsiLiteral(parent: PsiElement): PsiExpression? {
+internal fun KtConstantValue.createPsiLiteral(parent: PsiElement): PsiExpression? {
-    val asString = asStringForPsiLiteral()
+    val asString = asStringForPsiLiteral() ?: return null
     val instance = PsiElementFactory.getInstance(parent.project)
     return try {
         instance.createExpressionFromText(asString, parent)
@@ -200,5 +222,4 @@ internal fun KtSimpleConstantValue<*>.createPsiLiteral(parent: PsiElement): PsiE
     }
 }
 
-
 internal fun BitSet.copy(): BitSet = clone() as BitSet