[K2] Fix annotation serialization for simple cases

#KT-57611

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirAnnotationSerializer.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.serialization
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.serialization.constant.ConstantValue
 import org.jetbrains.kotlin.fir.serialization.constant.toConstantValue
@@ -28,7 +29,9 @@ class FirAnnotationSerializer(private val session: FirSession, internal val stri
         fun addArgument(argumentExpression: FirExpression, parameterName: Name) {
             val argument = ProtoBuf.Annotation.Argument.newBuilder()
             argument.nameId = stringTable.getStringIndex(parameterName.asString())
-            argument.setValue(valueProto(argumentExpression.toConstantValue(session) ?: return))
+            val constantValue = argumentExpression.toConstantValue(session)
+                ?: error("Cannot convert expression ${argumentExpression.render()} to constant")
+            argument.setValue(valueProto(constantValue))
             addArgument(argument)
         }
 

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirElementSerializer.kt

@@ -573,10 +573,6 @@ class FirElementSerializer private constructor(
             builder.addAllVersionRequirement(serializeVersionRequirements(typeAlias))
         }
 
-        for (annotation in typeAlias.nonSourceAnnotations(session)) {
-            builder.addAnnotation(extension.annotationSerializer.serializeAnnotation(annotation))
-        }
-
         extension.serializeTypeAlias(typeAlias, builder)
 
         return builder

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirKLibSerializerExtension.kt

@@ -24,10 +24,10 @@ import org.jetbrains.kotlin.serialization.deserialization.DYNAMIC_TYPE_DESERIALI
 class FirKLibSerializerExtension(
     override val session: FirSession,
     override val metadataVersion: BinaryVersion,
-    constValueProvider: ConstValueProvider?,
+    override val constValueProvider: ConstValueProvider?,
     private val allowErrorTypes: Boolean,
     private val exportKDoc: Boolean
-) : FirSerializerExtensionBase(KlibMetadataSerializerProtocol, constValueProvider) {
+) : FirSerializerExtensionBase(KlibMetadataSerializerProtocol) {
     override fun shouldUseTypeTable(): Boolean = true
 
     override fun serializeFlexibleType(type: ConeFlexibleType, lowerProto: ProtoBuf.Type.Builder, upperProto: ProtoBuf.Type.Builder) {

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirSerializerExtension.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.serialization
 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.types.ConeErrorType
 import org.jetbrains.kotlin.fir.types.ConeFlexibleType
 import org.jetbrains.kotlin.metadata.ProtoBuf
@@ -24,6 +25,8 @@ abstract class FirSerializerExtension {
 
     val annotationSerializer by lazy { FirAnnotationSerializer(session, stringTable) }
 
+    protected abstract val constValueProvider: ConstValueProvider?
+
     open fun shouldUseTypeTable(): Boolean = false
     open fun shouldUseNormalizedVisibility(): Boolean = false
 
@@ -77,6 +80,10 @@ 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))
+        }
     }
 
     open fun serializeErrorType(type: ConeErrorType, builder: ProtoBuf.Type.Builder) {

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirSerializerExtensionBase.kt

@@ -10,6 +10,8 @@ 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.ConstValueProvider
+import org.jetbrains.kotlin.fir.serialization.constant.buildValueProtoBufIfPropertyIsConst
+import org.jetbrains.kotlin.library.metadata.KlibMetadataProtoBuf
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
 import org.jetbrains.kotlin.name.FqName
@@ -18,7 +20,6 @@ import org.jetbrains.kotlin.serialization.SerializerExtensionProtocol
 
 abstract class FirSerializerExtensionBase(
     val protocol: SerializerExtensionProtocol,
-    private val constValueProvider: ConstValueProvider?
 ) : FirSerializerExtension() {
     final override val stringTable = FirElementAwareSerializableStringTable()
 
@@ -50,7 +51,7 @@ abstract class FirSerializerExtensionBase(
         childSerializer: FirElementSerializer
     ) {
         function.serializeAnnotations(proto, protocol.functionAnnotation)
-        function.receiverParameter?.serializeAnnotations(proto, protocol.functionExtensionReceiverAnnotation)
+        function.receiverParameter?.serializeAnnotations(proto, protocol.functionExtensionReceiverAnnotation, function)
     }
 
     override fun serializeProperty(
@@ -72,13 +73,13 @@ abstract class FirSerializerExtensionBase(
             destination += annotation
         }
 
-        regularPropertyAnnotations.serializeAnnotations(proto, protocol.propertyAnnotation)
-        fieldPropertyAnnotations.serializeAnnotations(proto, protocol.propertyBackingFieldAnnotation)
-        delegatePropertyAnnotations.serializeAnnotations(proto, protocol.propertyDelegatedFieldAnnotation)
+        regularPropertyAnnotations.serializeAnnotations(proto, protocol.propertyAnnotation, property)
+        fieldPropertyAnnotations.serializeAnnotations(proto, protocol.propertyBackingFieldAnnotation, property)
+        delegatePropertyAnnotations.serializeAnnotations(proto, protocol.propertyDelegatedFieldAnnotation, property)
 
         property.getter?.serializeAnnotations(proto, protocol.propertyGetterAnnotation)
         property.setter?.serializeAnnotations(proto, protocol.propertySetterAnnotation)
-        property.receiverParameter?.serializeAnnotations(proto, protocol.propertyExtensionReceiverAnnotation)
+        property.receiverParameter?.serializeAnnotations(proto, protocol.propertyExtensionReceiverAnnotation, property)
 
         constValueProvider?.buildValueProtoBufIfPropertyIsConst(property, annotationSerializer)?.let { constProtoBuf ->
             proto.setExtension(protocol.compileTimeValue, constProtoBuf)
@@ -94,7 +95,8 @@ abstract class FirSerializerExtensionBase(
     }
 
     override fun serializeTypeAnnotations(annotations: List<FirAnnotation>, proto: ProtoBuf.Type.Builder) {
-        annotations.serializeAnnotations(proto, protocol.typeAnnotation)
+        // TODO support const extraction for type annotations
+        annotations.serializeAnnotations(proto, protocol.typeAnnotation, container = null)
     }
 
     override fun serializeTypeParameter(typeParameter: FirTypeParameter, proto: ProtoBuf.TypeParameter.Builder) {
@@ -111,9 +113,10 @@ abstract class FirSerializerExtensionBase(
     > FirAnnotationContainer.serializeAnnotations(
         proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
         extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?,
+        container: FirAnnotationContainer? = this
     ) {
         if (extension == null) return
-        this.nonSourceAnnotations(session).serializeAnnotations(proto, extension)
+        this.nonSourceAnnotations(session).serializeAnnotations(proto, extension, container)
     }
 
     @Suppress("Reformat")
@@ -123,10 +126,37 @@ 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) {
-            proto.addExtensionOrNull(extension, annotationSerializer.serializeAnnotation(annotation))
+            val annotationWithConstants = when {
+                container == null -> null
+                container is FirPropertyAccessor ->
+                    constValueProvider?.getNewFirAnnotationWithConstantValues(
+                        container.propertySymbol.fir,
+                        annotation,
+                        container,
+                        isGetter = container.isGetter
+                    )
+                container is FirValueParameter ->
+                    constValueProvider?.getNewFirAnnotationWithConstantValues(
+                        container,
+                        annotation,
+                    )
+                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))
         }
     }
 

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/constant/ConstValueProvider.kt

@@ -5,8 +5,13 @@
 
 package org.jetbrains.kotlin.fir.serialization.constant
 
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
 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.declarations.FirValueParameter
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirConstExpression
 import org.jetbrains.kotlin.fir.serialization.FirAnnotationSerializer
 import org.jetbrains.kotlin.metadata.ProtoBuf
@@ -16,17 +21,37 @@ abstract class ConstValueProvider {
 
     abstract fun getConstantValueForProperty(firProperty: FirProperty): FirConstExpression<*>?
 
-    fun buildValueProtoBufIfPropertyIsConst(
-        firProperty: FirProperty, annotationSerializer: FirAnnotationSerializer
-    ): ProtoBuf.Annotation.Argument.Value? {
-        return getConstantValueForProperty(firProperty).toProtoBuf(annotationSerializer)
-    }
+    abstract fun getNewFirAnnotationWithConstantValues(
+        firAnnotationContainer: FirAnnotationContainer,
+        firAnnotation: FirAnnotation,
+    ): FirAnnotation
 
-    private fun FirConstExpression<*>?.toProtoBuf(
-        annotationSerializer: FirAnnotationSerializer
-    ): ProtoBuf.Annotation.Argument.Value? {
+    abstract fun getNewFirAnnotationWithConstantValues(
+        firProperty: FirProperty,
+        firAnnotation: FirAnnotation,
+        firPropertyAccessor: FirPropertyAccessor,
+        isGetter: Boolean,
+    ): FirAnnotation
+
+    abstract fun getNewFirAnnotationWithConstantValues(
+        firExtensionReceiverContainer: FirAnnotationContainer,
+        firAnnotation: FirAnnotation,
+        receiverParameter: FirReceiverParameter,
+    ): FirAnnotation
+
+    abstract fun getNewFirAnnotationWithConstantValues(
+        valueParameter: FirValueParameter,
+        firAnnotation: FirAnnotation,
+    ): FirAnnotation
+
+    fun FirConstExpression<*>?.toProtoBuf(annotationSerializer: FirAnnotationSerializer): ProtoBuf.Annotation.Argument.Value? {
         val constantValue = this?.toConstantValue(session) ?: return null
         return annotationSerializer.valueProto(constantValue).build()
     }
 }
 
+fun ConstValueProvider.buildValueProtoBufIfPropertyIsConst(
+    firProperty: FirProperty, annotationSerializer: FirAnnotationSerializer
+): ProtoBuf.Annotation.Argument.Value? {
+    return getConstantValueForProperty(firProperty).toProtoBuf(annotationSerializer)
+}