FIR LC: handle more kinds of annotation values

2021-09-16 14:54:17 -07:00
parent 6ef2dad895
commit 4c8d975ae5
7 changed files with 143 additions and 36 deletions
@@ -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
            }
@@ -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})"
+    }
 }
@@ -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
-}
+}
@@ -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
+}
@@ -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)
-}
+}
@@ -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
 }
@@ -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.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.symbols.markers.*
 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 =
-    when (val value = this.value) {
-        is String -> "\"${escapeString(value)}\""
-        is Long -> "${value}L"
-        is Float -> "${value}f"
-        else -> value?.toString() ?: "null"
+internal fun KtConstantValue.toAnnotationMemberValue(parent: PsiElement): PsiAnnotationMemberValue? {
+    return when (this) {
+        is KtArrayConstantValue ->
+            FirPsiArrayInitializerMemberValue(kotlinOrigin, parent) { arrayLiteralParent ->
+                values.mapNotNull { element -> element.toAnnotationMemberValue(arrayLiteralParent) }
+            }
+        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? {
-    val asString = asStringForPsiLiteral()
+internal fun KtConstantValue.createPsiLiteral(parent: PsiElement): PsiExpression? {
+    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