[K2] Rewrite the way of extraction const values from IR for FIR

Right now it is much easier to extract constant values from `EvaluatedConstTracker` #KT-57928 #KT-57929
2023-04-08 14:55:56 +02:00
parent 1ddcdcfc39
commit 45aacae945
10 changed files with 82 additions and 299 deletions
@@ -9,27 +9,32 @@ import org.jetbrains.kotlin.constant.AnnotationValue
 import org.jetbrains.kotlin.constant.ConstantValue
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.serialization.constant.ConstValueProvider
 import org.jetbrains.kotlin.fir.serialization.constant.coneTypeSafe
-import org.jetbrains.kotlin.fir.serialization.constant.convertToConstantValues
+import org.jetbrains.kotlin.fir.serialization.constant.toConstantValue
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
-import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.name.Name

-class FirAnnotationSerializer(private val session: FirSession, internal val stringTable: FirElementAwareStringTable) {
+class FirAnnotationSerializer(
+    private val session: FirSession,
+    internal val stringTable: FirElementAwareStringTable,
+    private val constValueProvider: ConstValueProvider?
+) {
    fun serializeAnnotation(annotation: FirAnnotation): ProtoBuf.Annotation {
-        return serializeAnnotation(
-            annotation.typeRef.coneTypeSafe<ConeClassLikeType>(),
-            annotation.argumentMapping.mapping.convertToConstantValues(session)
-        )
+        // TODO this logic can be significantly simplified if we will find the way to convert `IrAnnotation` to `AnnotationValue`
+        val annotationValue = annotation.toConstantValue(session, constValueProvider) as? AnnotationValue
+            ?: error("Cannot serialize annotation ${annotation.render()}")
+        return serializeAnnotation(annotationValue)
    }

    fun serializeAnnotation(annotation: AnnotationValue): ProtoBuf.Annotation {
        return serializeAnnotation(annotation.coneTypeSafe<ConeClassLikeType>(), annotation.value.argumentsMapping)
    }

-    private fun serializeAnnotation(coneType: ConeClassLikeType?, argumentsMapping: Map<Name, ConstantValue<*>?>): ProtoBuf.Annotation {
+    private fun serializeAnnotation(coneType: ConeClassLikeType?, argumentsMapping: Map<Name, ConstantValue<*>>): ProtoBuf.Annotation {
        return ProtoBuf.Annotation.newBuilder().apply {
            val lookupTag = coneType?.lookupTag
                ?: error { "Annotation without proper lookup tag: $coneType" }
@@ -37,10 +42,7 @@ class FirAnnotationSerializer(private val session: FirSession, internal val stri
            id = lookupTag.toSymbol(session)?.let { stringTable.getFqNameIndex(it.fir) }
                ?: stringTable.getQualifiedClassNameIndex(lookupTag.classId)

-            fun addArgument(argumentExpression: ConstantValue<*>?, parameterName: Name) {
-                if (argumentExpression == null) {
-                    error("Cannot use null argument expression for parameter $parameterName")
-                }
+            fun addArgument(argumentExpression: ConstantValue<*>, parameterName: Name) {
                val argument = ProtoBuf.Annotation.Argument.newBuilder()
                argument.nameId = stringTable.getStringIndex(parameterName.asString())
                argument.setValue(valueProto(argumentExpression))
@@ -9,12 +9,10 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.serialization.constant.ConstValueProvider
-import org.jetbrains.kotlin.fir.serialization.constant.buildValueProtoBufIfPropertyHasConst
 import org.jetbrains.kotlin.fir.types.ConeErrorType
 import org.jetbrains.kotlin.fir.types.ConeFlexibleType
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
-import org.jetbrains.kotlin.metadata.deserialization.Flags
 import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
 import org.jetbrains.kotlin.name.FqName

@@ -25,7 +23,7 @@ abstract class FirSerializerExtension {

    abstract val metadataVersion: BinaryVersion

-    val annotationSerializer by lazy { FirAnnotationSerializer(session, stringTable) }
+    val annotationSerializer by lazy { FirAnnotationSerializer(session, stringTable, constValueProvider) }

    protected abstract val constValueProvider: ConstValueProvider?

@@ -83,8 +81,7 @@ abstract class FirSerializerExtension {

    open fun serializeTypeAlias(typeAlias: FirTypeAlias, proto: ProtoBuf.TypeAlias.Builder) {
        for (annotation in typeAlias.nonSourceAnnotations(session)) {
-            val annotationWithConstants = constValueProvider?.getNewFirAnnotationWithConstantValues(typeAlias, annotation) ?: annotation
-            proto.addAnnotation(annotationSerializer.serializeAnnotation(annotationWithConstants))
+            proto.addAnnotation(annotationSerializer.serializeAnnotation(annotation))
        }
    }

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.serialization.constant.buildValueProtoBufIfPropertyHasConst
+import org.jetbrains.kotlin.fir.serialization.constant.toConstantValue
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.Flags
 import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
@@ -50,7 +50,7 @@ abstract class FirSerializerExtensionBase(
        childSerializer: FirElementSerializer
    ) {
        function.serializeAnnotations(proto, protocol.functionAnnotation)
-        function.receiverParameter?.serializeAnnotations(proto, protocol.functionExtensionReceiverAnnotation, function)
+        function.receiverParameter?.serializeAnnotations(proto, protocol.functionExtensionReceiverAnnotation)
    }

    override fun serializeProperty(
@@ -72,17 +72,17 @@ abstract class FirSerializerExtensionBase(
            destination += annotation
        }

-        regularPropertyAnnotations.serializeAnnotations(proto, protocol.propertyAnnotation, property)
-        fieldPropertyAnnotations.serializeAnnotations(proto, protocol.propertyBackingFieldAnnotation, property)
-        delegatePropertyAnnotations.serializeAnnotations(proto, protocol.propertyDelegatedFieldAnnotation, property)
+        regularPropertyAnnotations.serializeAnnotations(proto, protocol.propertyAnnotation)
+        fieldPropertyAnnotations.serializeAnnotations(proto, protocol.propertyBackingFieldAnnotation)
+        delegatePropertyAnnotations.serializeAnnotations(proto, protocol.propertyDelegatedFieldAnnotation)

        property.getter?.serializeAnnotations(proto, protocol.propertyGetterAnnotation)
        property.setter?.serializeAnnotations(proto, protocol.propertySetterAnnotation)
-        property.receiverParameter?.serializeAnnotations(proto, protocol.propertyExtensionReceiverAnnotation, property)
+        property.receiverParameter?.serializeAnnotations(proto, protocol.propertyExtensionReceiverAnnotation)

        if (!Flags.HAS_CONSTANT.get(proto.flags)) return
-        constValueProvider?.buildValueProtoBufIfPropertyHasConst(property, annotationSerializer)?.let { constProtoBuf ->
-            proto.setExtension(protocol.compileTimeValue, constProtoBuf)
+        property.initializer?.toConstantValue(session, constValueProvider)?.let {
+            proto.setExtension(protocol.compileTimeValue, annotationSerializer.valueProto(it).build())
        }
    }

@@ -96,7 +96,7 @@ abstract class FirSerializerExtensionBase(

    override fun serializeTypeAnnotations(annotations: List<FirAnnotation>, proto: ProtoBuf.Type.Builder) {
        // TODO support const extraction for type annotations
-        annotations.serializeAnnotations(proto, protocol.typeAnnotation, container = null)
+        annotations.serializeAnnotations(proto, protocol.typeAnnotation)
    }

    override fun serializeTypeParameter(typeParameter: FirTypeParameter, proto: ProtoBuf.TypeParameter.Builder) {
@@ -112,11 +112,10 @@ abstract class FirSerializerExtensionBase(
        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
    > FirAnnotationContainer.serializeAnnotations(
        proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
-        extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?,
-        container: FirAnnotationContainer? = this
+        extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?
    ) {
        if (extension == null) return
-        this.nonSourceAnnotations(session).serializeAnnotations(proto, extension, container)
+        this.nonSourceAnnotations(session).serializeAnnotations(proto, extension)
    }

    @Suppress("Reformat")
@@ -126,25 +125,10 @@ abstract class FirSerializerExtensionBase(
    > List<FirAnnotation>.serializeAnnotations(
        proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
        extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?,
-        container: FirAnnotationContainer?,
    ) {
        if (extension == null) return
        for (annotation in this) {
-            val annotationWithConstants = when {
-                container == null -> null
-                extension == protocol.propertyExtensionReceiverAnnotation || extension == protocol.functionExtensionReceiverAnnotation  ->
-                    constValueProvider?.getNewFirAnnotationWithConstantValues(
-                        container,
-                        annotation,
-                        (container as FirCallableDeclaration).receiverParameter!!,
-                    )
-                else ->
-                    constValueProvider?.getNewFirAnnotationWithConstantValues(
-                        container,
-                        annotation,
-                    )
-            } ?: annotation
-            proto.addExtensionOrNull(extension, annotationSerializer.serializeAnnotation(annotationWithConstants))
+            proto.addExtensionOrNull(extension, annotationSerializer.serializeAnnotation(annotation))
        }
    }

@@ -5,40 +5,14 @@

 package org.jetbrains.kotlin.fir.serialization.constant

-import org.jetbrains.kotlin.fir.FirAnnotationContainer
+import org.jetbrains.kotlin.constant.ConstantValue
+import org.jetbrains.kotlin.constant.EvaluatedConstTracker
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
-import org.jetbrains.kotlin.fir.declarations.FirReceiverParameter
-import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.serialization.FirAnnotationSerializer
-import org.jetbrains.kotlin.metadata.ProtoBuf

 abstract class ConstValueProvider {
    abstract val session: FirSession
+    abstract val evaluatedConstTracker: EvaluatedConstTracker

-    abstract fun getConstantValueForProperty(firProperty: FirProperty): FirExpression?
-
-    abstract fun getNewFirAnnotationWithConstantValues(
-        firAnnotationContainer: FirAnnotationContainer,
-        firAnnotation: FirAnnotation,
-    ): FirAnnotation
-
-    abstract fun getNewFirAnnotationWithConstantValues(
-        firExtensionReceiverContainer: FirAnnotationContainer,
-        firAnnotation: FirAnnotation,
-        receiverParameter: FirReceiverParameter,
-    ): FirAnnotation
-
-    fun FirExpression?.toProtoBuf(annotationSerializer: FirAnnotationSerializer): ProtoBuf.Annotation.Argument.Value? {
-        val constantValue = this?.toConstantValue(session) ?: return null
-        return annotationSerializer.valueProto(constantValue).build()
-    }
-}
-
-fun ConstValueProvider.buildValueProtoBufIfPropertyHasConst(
-    firProperty: FirProperty, annotationSerializer: FirAnnotationSerializer
-): ProtoBuf.Annotation.Argument.Value? {
-    return getConstantValueForProperty(firProperty).toProtoBuf(annotationSerializer)
+    abstract fun findConstantValueFor(firExpression: FirExpression?): ConstantValue<*>?
 }
@@ -12,20 +12,30 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.name.Name

-internal fun Map<Name, FirExpression>.convertToConstantValues(session: FirSession): Map<Name, ConstantValue<*>?> {
+internal fun Map<Name, FirExpression>.convertToConstantValues(
+    session: FirSession,
+    constValueProvider: ConstValueProvider?
+): Map<Name, ConstantValue<*>> {
    return this.map { (name, firExpression) ->
-        name to firExpression.toConstantValue(session)
+        val constantValue = constValueProvider?.findConstantValueFor(firExpression)
+            ?: firExpression.toConstantValue(session, constValueProvider)
+            ?: error("Cannot convert expression ${firExpression.render()} to constant")
+        name to constantValue
    }.toMap()
 }

-internal fun LinkedHashMap<FirExpression, FirValueParameter>.convertToConstantValues(session: FirSession): Map<Name, ConstantValue<*>?> {
+internal fun LinkedHashMap<FirExpression, FirValueParameter>.convertToConstantValues(
+    session: FirSession,
+    constValueProvider: ConstValueProvider?,
+): Map<Name, ConstantValue<*>> {
    return this.map { (firExpression, firValueParameter) -> firValueParameter.name to firExpression }
-        .toMap().convertToConstantValues(session)
+        .toMap().convertToConstantValues(session, constValueProvider)
 }

 inline fun <reified T : ConeKotlinType> AnnotationValue.coneTypeSafe(): T? {
@@ -25,19 +25,26 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 import org.jetbrains.kotlin.types.ConstantValueKind

-internal fun FirExpression.toConstantValue(session: FirSession): ConstantValue<*>? {
-    return accept(FirToConstantValueTransformerUnsafe, session)
+internal fun FirExpression.toConstantValue(session: FirSession, constValueProvider: ConstValueProvider? = null): ConstantValue<*>? {
+    constValueProvider?.findConstantValueFor(this)?.let { return it }
+    return accept(FirToConstantValueTransformerUnsafe(constValueProvider), session)
 }

 internal fun FirExpression?.hasConstantValue(session: FirSession): Boolean {
    return this?.accept(FirToConstantValueChecker, session) == true
 }

-private object FirToConstantValueTransformerSafe : FirToConstantValueTransformer(failOnNonConst = false)
-private object FirToConstantValueTransformerUnsafe : FirToConstantValueTransformer(failOnNonConst = true)
+private class FirToConstantValueTransformerSafe(
+    constValueProvider: ConstValueProvider?
+) : FirToConstantValueTransformer(failOnNonConst = false, constValueProvider)
+
+private class FirToConstantValueTransformerUnsafe(
+    constValueProvider: ConstValueProvider?
+) : FirToConstantValueTransformer(failOnNonConst = true, constValueProvider)

 private abstract class FirToConstantValueTransformer(
-    private val failOnNonConst: Boolean
+    private val failOnNonConst: Boolean,
+    private val constValueProvider: ConstValueProvider?
 ) : FirDefaultVisitor<ConstantValue<*>?, FirSession>() {
    override fun visitElement(
        element: FirElement,
@@ -90,7 +97,7 @@ private abstract class FirToConstantValueTransformer(
        annotation: FirAnnotation,
        data: FirSession
    ): ConstantValue<*> {
-        val mapping = annotation.argumentMapping.mapping.convertToConstantValues(data)
+        val mapping = annotation.argumentMapping.mapping.convertToConstantValues(data, constValueProvider)
        return AnnotationValue.create(annotation.annotationTypeRef.coneType, mapping)
    }

@@ -135,7 +142,7 @@ private abstract class FirToConstantValueTransformer(
                val constructedClassSymbol = symbol.containingClassLookupTag()?.toFirRegularClassSymbol(data) ?: return null
                if (constructedClassSymbol.classKind != ClassKind.ANNOTATION_CLASS) return null

-                val mapping = constructorCall.resolvedArgumentMapping?.convertToConstantValues(data) ?: return null
+                val mapping = constructorCall.resolvedArgumentMapping?.convertToConstantValues(data, constValueProvider) ?: return null
                return AnnotationValue.create(qualifiedAccessExpression.typeRef.coneType, mapping)
            }