[K2] Return from toConstantValue only if expected type match

There was a bug with early return from `toConstantValue`
method when we found any const value in `ConstValueProvider`.
Actually we should return only if expected type match. This error can
occur if annotation and constant inside it are fakes and
have the same offset.

#KT-57812

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

@@ -25,7 +25,7 @@ class FirAnnotationSerializer(
 ) {
     fun serializeAnnotation(annotation: FirAnnotation): ProtoBuf.Annotation {
         // 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
+        val annotationValue = annotation.toConstantValue<AnnotationValue>(session, constValueProvider)
             ?: error("Cannot serialize annotation ${annotation.render()}")
         return serializeAnnotation(annotationValue)
     }

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

@@ -1091,8 +1091,7 @@ class FirElementSerializer private constructor(
     private fun serializeVersionRequirementFromRequireKotlin(annotation: FirAnnotation): ProtoBuf.VersionRequirement.Builder? {
         val argumentMapping = annotation.argumentMapping.mapping
 
-        val versionString =
+        val versionString = argumentMapping[RequireKotlinConstants.VERSION]?.toConstantValue<StringValue>(session)?.value ?: return null
-            (argumentMapping[RequireKotlinConstants.VERSION]?.toConstantValue(session) as? StringValue)?.value ?: return null
         val matchResult = RequireKotlinConstants.VERSION_REGEX.matchEntire(versionString) ?: return null
 
         val major = matchResult.groupValues.getOrNull(1)?.toIntOrNull() ?: return null
@@ -1105,12 +1104,12 @@ class FirElementSerializer private constructor(
             writeVersionFull = { proto.versionFull = it }
         )
 
-        val message = (argumentMapping[RequireKotlinConstants.MESSAGE]?.toConstantValue(session) as? StringValue)?.value
+        val message = argumentMapping[RequireKotlinConstants.MESSAGE]?.toConstantValue<StringValue>(session)?.value
         if (message != null) {
             proto.message = stringTable.getStringIndex(message)
         }
 
-        when ((argumentMapping[RequireKotlinConstants.LEVEL]?.toConstantValue(session) as? EnumValue)?.enumEntryName?.asString()) {
+        when (argumentMapping[RequireKotlinConstants.LEVEL]?.toConstantValue<EnumValue>(session)?.enumEntryName?.asString()) {
             DeprecationLevel.ERROR.name -> {
                 // ERROR is the default level
             }
@@ -1118,7 +1117,7 @@ class FirElementSerializer private constructor(
             DeprecationLevel.HIDDEN.name -> proto.level = ProtoBuf.VersionRequirement.Level.HIDDEN
         }
 
-        when ((argumentMapping[RequireKotlinConstants.VERSION_KIND]?.toConstantValue(session) as? EnumValue)?.enumEntryName?.asString()) {
+        when (argumentMapping[RequireKotlinConstants.VERSION_KIND]?.toConstantValue<EnumValue>(session)?.enumEntryName?.asString()) {
             ProtoBuf.VersionRequirement.VersionKind.LANGUAGE_VERSION.name -> {
                 // LANGUAGE_VERSION is the default kind
             }
@@ -1128,7 +1127,7 @@ class FirElementSerializer private constructor(
                 proto.versionKind = ProtoBuf.VersionRequirement.VersionKind.API_VERSION
         }
 
-        val errorCode = (argumentMapping[RequireKotlinConstants.ERROR_CODE]?.toConstantValue(session) as? IntValue)?.value
+        val errorCode = argumentMapping[RequireKotlinConstants.ERROR_CODE]?.toConstantValue<IntValue>(session)?.value
         if (errorCode != null && errorCode != -1) {
             proto.errorCode = errorCode
         }

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

@@ -30,11 +30,12 @@ abstract class FirSerializerExtension {
 
     @OptIn(ConstValueProviderInternals::class)
     internal inline fun <T> processFile(firFile: FirFile, action: () -> T): T {
+        val previousFile = constValueProvider?.processingFirFile
         constValueProvider?.processingFirFile = firFile
         return try {
             action()
         } finally {
-            constValueProvider?.processingFirFile = null
+            constValueProvider?.processingFirFile = previousFile
         }
     }
 

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.serialization
 
+import org.jetbrains.kotlin.constant.ConstantValue
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.declarations.*
@@ -79,7 +80,7 @@ abstract class FirSerializerExtensionBase(
         property.receiverParameter?.serializeAnnotations(proto, protocol.propertyExtensionReceiverAnnotation)
 
         if (!Flags.HAS_CONSTANT.get(proto.flags)) return
-        property.initializer?.toConstantValue(session, constValueProvider)?.let {
+        property.initializer?.toConstantValue<ConstantValue<*>>(session, constValueProvider)?.let {
             proto.setExtension(protocol.compileTimeValue, annotationSerializer.valueProto(it).build())
         }
     }

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

@@ -25,25 +25,28 @@ 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, constValueProvider: ConstValueProvider? = null): ConstantValue<*>? {
+internal inline fun <reified T : ConstantValue<*>> FirExpression.toConstantValue(
-    constValueProvider?.findConstantValueFor(this)?.let { return it }
+    session: FirSession,
-    return accept(FirToConstantValueTransformerUnsafe(), FirToConstantValueTransformerData(session, constValueProvider))
+    constValueProvider: ConstValueProvider? = null
+): T? {
+    return constValueProvider?.findConstantValueFor(this) as? T
+        ?: accept(FirToConstantValueTransformerUnsafe(), FirToConstantValueTransformerData(session, constValueProvider)) as? T
 }
 
 internal fun FirExpression?.hasConstantValue(session: FirSession): Boolean {
     return this?.accept(FirToConstantValueChecker, session) == true
 }
 
-private class FirToConstantValueTransformerSafe : FirToConstantValueTransformer(failOnNonConst = false)
+internal class FirToConstantValueTransformerSafe : FirToConstantValueTransformer(failOnNonConst = false)
 
-private class FirToConstantValueTransformerUnsafe : FirToConstantValueTransformer(failOnNonConst = true)
+internal class FirToConstantValueTransformerUnsafe : FirToConstantValueTransformer(failOnNonConst = true)
 
-private data class FirToConstantValueTransformerData(
+internal data class FirToConstantValueTransformerData(
     val session: FirSession,
     val constValueProvider: ConstValueProvider?,
 )
 
-private abstract class FirToConstantValueTransformer(
+internal abstract class FirToConstantValueTransformer(
     private val failOnNonConst: Boolean,
 ) : FirDefaultVisitor<ConstantValue<*>?, FirToConstantValueTransformerData>() {
     override fun visitElement(