Serialize and deserialize const val's properly for K2 klib

#KT-57312 Fixed

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

@@ -8,8 +8,10 @@ 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.library.metadata.KlibMetadataProtoBuf
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
 import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
@@ -24,6 +26,8 @@ abstract class FirSerializerExtension {
 
     val annotationSerializer by lazy { FirAnnotationSerializer(session, stringTable) }
 
+    abstract val constValueProvider: ConstValueProvider?
+
     open fun shouldUseTypeTable(): Boolean = false
     open fun shouldUseNormalizedVisibility(): Boolean = false
 
@@ -83,6 +87,11 @@ abstract class FirSerializerExtension {
         throw IllegalStateException("Cannot serialize error type: $type")
     }
 
+    protected fun serializeConstant(property: FirProperty, proto: ProtoBuf.Property.Builder) {
+        val constProtoBuf = constValueProvider?.buildValueProtoBufIfPropertyIsConst(property, annotationSerializer) ?: return
+        proto.setExtension(KlibMetadataProtoBuf.compileTimeValue, constProtoBuf)
+    }
+
     open val customClassMembersProducer: ClassMembersProducer?
         get() = null
 

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2010-2023 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.fir.serialization.constant
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.expressions.FirConstExpression
+import org.jetbrains.kotlin.fir.serialization.FirAnnotationSerializer
+import org.jetbrains.kotlin.metadata.ProtoBuf
+
+abstract class ConstValueProvider {
+    abstract val session: FirSession
+
+    abstract fun getConstantValueForProperty(firProperty: FirProperty): FirConstExpression<*>?
+
+    fun buildValueProtoBufIfPropertyIsConst(
+        firProperty: FirProperty, annotationSerializer: FirAnnotationSerializer
+    ): ProtoBuf.Annotation.Argument.Value? {
+        return getConstantValueForProperty(firProperty).toProtoBuf(annotationSerializer)
+    }
+
+    private fun FirConstExpression<*>?.toProtoBuf(
+        annotationSerializer: FirAnnotationSerializer
+    ): ProtoBuf.Annotation.Argument.Value? {
+        val constantValue = this?.toConstantValue(session) ?: return null
+        return annotationSerializer.valueProto(constantValue).build()
+    }
+}
+

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

@@ -49,11 +49,11 @@ internal object FirToConstantValueTransformer : FirDefaultVisitor<ConstantValue<
             ConstantValueKind.UnsignedShort -> UShortValue((value as Number).toShort())
             ConstantValueKind.Int -> IntValue((value as Number).toInt())
             ConstantValueKind.UnsignedInt -> UIntValue((value as Number).toInt())
-            ConstantValueKind.Long -> LongValue(value as Long)
-            ConstantValueKind.UnsignedLong -> ULongValue(value as Long)
+            ConstantValueKind.Long -> LongValue((value as Number).toLong())
+            ConstantValueKind.UnsignedLong -> ULongValue((value as Number).toLong())
             ConstantValueKind.String -> StringValue(value as String)
-            ConstantValueKind.Float -> FloatValue(value as Float)
-            ConstantValueKind.Double -> DoubleValue(value as Double)
+            ConstantValueKind.Float -> FloatValue((value as Number).toFloat())
+            ConstantValueKind.Double -> DoubleValue((value as Number).toDouble())
             ConstantValueKind.Null -> NullValue
             else -> null
         }
@@ -99,24 +99,22 @@ internal object FirToConstantValueTransformer : FirDefaultVisitor<ConstantValue<
             symbol is FirConstructorSymbol -> {
                 val constructorCall = qualifiedAccessExpression as FirFunctionCall
                 val constructedClassSymbol = symbol.containingClassLookupTag()?.toFirRegularClassSymbol(data) ?: return null
-                return if (constructedClassSymbol.classKind == ClassKind.ANNOTATION_CLASS) {
-                    AnnotationValue(
-                        buildAnnotationCall {
-                            argumentMapping = buildAnnotationArgumentMapping {
-                                constructorCall.resolvedArgumentMapping?.forEach { (firExpression, firValueParameter) ->
-                                    mapping[firValueParameter.name] = firExpression
-                                }
-                            }
-                            annotationTypeRef = qualifiedAccessExpression.typeRef
-                            calleeReference = buildSimpleNamedReference {
-                                source = qualifiedAccessExpression.source
-                                name = qualifiedAccessExpression.calleeReference.name
+                if (constructedClassSymbol.classKind != ClassKind.ANNOTATION_CLASS) return null
+
+                return AnnotationValue(
+                    buildAnnotationCall {
+                        argumentMapping = buildAnnotationArgumentMapping {
+                            constructorCall.resolvedArgumentMapping?.forEach { (firExpression, firValueParameter) ->
+                                mapping[firValueParameter.name] = firExpression
                             }
                         }
-                    )
-                } else {
-                    null
-                }
+                        annotationTypeRef = qualifiedAccessExpression.typeRef
+                        calleeReference = buildSimpleNamedReference {
+                            source = qualifiedAccessExpression.source
+                            name = qualifiedAccessExpression.calleeReference.name
+                        }
+                    }
+                )
             }
 
             symbol.callableId.packageName.asString() == "kotlin" -> {

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

@@ -13,8 +13,10 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.serialization.constant.ConstValueProvider
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
+import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.serialization.SerializableStringTable
 
@@ -66,8 +68,19 @@ fun serializeSingleFirFile(
 open class FirKLibSerializerExtension(
     override val session: FirSession,
     override val metadataVersion: BinaryVersion,
-    override val stringTable: FirElementAwareSerializableStringTable
-) : FirSerializerExtension()
+    override val stringTable: FirElementAwareSerializableStringTable,
+    override val constValueProvider: ConstValueProvider?
+) : FirSerializerExtension() {
+    override fun serializeProperty(
+        property: FirProperty,
+        proto: ProtoBuf.Property.Builder,
+        versionRequirementTable: MutableVersionRequirementTable?,
+        childSerializer: FirElementSerializer
+    ) {
+        super.serializeProperty(property, proto, versionRequirementTable, childSerializer)
+        serializeConstant(property, proto)
+    }
+}
 
 class FirElementAwareSerializableStringTable() : FirElementAwareStringTable, SerializableStringTable() {
     override fun getLocalClassIdReplacement(firClass: FirClass): ClassId? = ClassId.topLevel(StandardNames.FqNames.any.toSafe())