FIR2IR: properly set constant expression type

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/evaluate/FirCompileTimeConstantEvaluator.kt

@@ -390,7 +390,7 @@ internal object FirCompileTimeConstantEvaluator {
 
     private fun <T> ConstantValueKind<T>.toConstExpression(source: KtSourceElement?, value: Any?): FirConstExpression<T> =
         @Suppress("UNCHECKED_CAST")
-        buildConstExpression(source, this, value as T)
+        buildConstExpression(source, this, value as T, setType = false)
 
     private fun FirFunctionCall.getOriginalFunction(): FirCallableDeclaration? {
         val symbol: FirBasedSymbol<*>? = when (val reference = calleeReference) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FirElementsRecorder.kt

@@ -215,7 +215,8 @@ internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement,
         return buildConstExpression(
             original.ktConstantExpression?.toKtPsiSourceElement(),
             this,
-            convertedValue as T
+            convertedValue as T,
+            setType = false
         ).also {
             it.replaceTypeRef(original.typeRef)
         }

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirConstDeserializer.kt

@@ -35,16 +35,33 @@ fun buildFirConstant(
     protoValue: ProtoBuf.Annotation.Argument.Value?, sourceValue: Any?, constKind: String, nameResolver: NameResolver
 ): FirExpression? {
     return when (constKind) {
-        "BYTE", "B" -> buildConstExpression(null, ConstantValueKind.Byte, ((protoValue?.intValue ?: sourceValue) as Number).toByte())
-        "CHAR", "C" -> buildConstExpression(null, ConstantValueKind.Char, ((protoValue?.intValue ?: sourceValue) as Number).toInt().toChar())
-        "SHORT", "S" -> buildConstExpression(null, ConstantValueKind.Short, ((protoValue?.intValue ?: sourceValue) as Number).toShort())
-        "INT", "I" -> buildConstExpression(null, ConstantValueKind.Int, protoValue?.intValue?.toInt() ?: sourceValue as Int)
-        "LONG", "J" -> buildConstExpression(null, ConstantValueKind.Long, protoValue?.intValue ?: sourceValue as Long)
-        "FLOAT", "F" -> buildConstExpression(null, ConstantValueKind.Float, protoValue?.floatValue ?: sourceValue as Float)
-        "DOUBLE", "D" -> buildConstExpression(null, ConstantValueKind.Double, protoValue?.doubleValue ?: sourceValue as Double)
-        "BOOLEAN", "Z" -> buildConstExpression(null, ConstantValueKind.Boolean, (protoValue?.intValue?.toInt() ?: sourceValue) != 0)
+        "BYTE", "B" -> buildConstExpression(
+            null, ConstantValueKind.Byte, ((protoValue?.intValue ?: sourceValue) as Number).toByte(), setType = true
+        )
+        "CHAR", "C" -> buildConstExpression(
+            null, ConstantValueKind.Char, ((protoValue?.intValue ?: sourceValue) as Number).toInt().toChar(), setType = true
+        )
+        "SHORT", "S" -> buildConstExpression(
+            null, ConstantValueKind.Short, ((protoValue?.intValue ?: sourceValue) as Number).toShort(), setType = true
+        )
+        "INT", "I" -> buildConstExpression(
+            null, ConstantValueKind.Int, protoValue?.intValue?.toInt() ?: sourceValue as Int, setType = true
+        )
+        "LONG", "J" -> buildConstExpression(
+            null, ConstantValueKind.Long, protoValue?.intValue ?: sourceValue as Long, setType = true
+        )
+        "FLOAT", "F" -> buildConstExpression(
+            null, ConstantValueKind.Float, protoValue?.floatValue ?: sourceValue as Float, setType = true
+        )
+        "DOUBLE", "D" -> buildConstExpression(
+            null, ConstantValueKind.Double, protoValue?.doubleValue ?: sourceValue as Double, setType = true
+        )
+        "BOOLEAN", "Z" -> buildConstExpression(
+            null, ConstantValueKind.Boolean, (protoValue?.intValue?.toInt() ?: sourceValue) != 0, setType = true
+        )
         "STRING", "Ljava/lang/String;" -> buildConstExpression(
-            null, ConstantValueKind.String, protoValue?.stringValue?.let { nameResolver.getString(it) } ?: sourceValue as String
+            null, ConstantValueKind.String,
+            protoValue?.stringValue?.let { nameResolver.getString(it) } ?: sourceValue as String, setType = true
         )
         else -> null
     }

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirMemberDeserializer.kt

@@ -467,6 +467,7 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
 
             proto.contextReceiverTypes(c.typeTable).mapTo(contextReceivers, ::loadContextReceiver)
         }.apply {
+            initializer?.replaceTypeRef(returnTypeRef)
             versionRequirementsTable = c.versionRequirementTable
             setLazyPublishedVisibility(c.session)
             getter?.setLazyPublishedVisibility(annotations, this, c.session)

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

@@ -600,7 +600,9 @@ class FirElementSerializer private constructor(
         }
     }
 
-    private fun typeAliasProto(typeAlias: FirTypeAlias): ProtoBuf.TypeAlias.Builder? = whileAnalysing<ProtoBuf.TypeAlias.Builder?>(session, typeAlias) {
+    private fun typeAliasProto(typeAlias: FirTypeAlias): ProtoBuf.TypeAlias.Builder? = whileAnalysing<ProtoBuf.TypeAlias.Builder?>(
+        session, typeAlias
+    ) {
         val builder = ProtoBuf.TypeAlias.newBuilder()
         val local = createChildSerializer(typeAlias)
 
@@ -851,7 +853,9 @@ class FirElementSerializer private constructor(
                             correspondingTypeRef?.annotations, CompilerConeAttributes.ContextFunctionTypeParams.ANNOTATION_CLASS_ID,
                             argumentMapping = buildAnnotationArgumentMapping {
                                 this.mapping[StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME] =
-                                    buildConstExpression(source = null, ConstantValueKind.Int, type.contextReceiversNumberForFunctionType)
+                                    buildConstExpression(
+                                        source = null, ConstantValueKind.Int, type.contextReceiversNumberForFunctionType, setType = false
+                                    )
                             }
                         )
                     )

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrAnnotationsFromPluginRegistrar.kt

@@ -117,16 +117,66 @@ class Fir2IrAnnotationsFromPluginRegistrar(private val components: Fir2IrCompone
                         val name = this@toFirAnnotation.symbol.owner.valueParameters[i].name
                         val argument = this@toFirAnnotation.getValueArgument(i) as IrConst<*>
                         this.mapping[name] = when (argument.kind) {
-                            IrConstKind.Boolean -> buildConstExpression(source = null, ConstantValueKind.Boolean, argument.value as Boolean)
-                            IrConstKind.Byte -> buildConstExpression(source = null, ConstantValueKind.Byte, argument.value as Byte)
-                            IrConstKind.Char -> buildConstExpression(source = null, ConstantValueKind.Char, argument.value as Char)
-                            IrConstKind.Double -> buildConstExpression(source = null, ConstantValueKind.Double, argument.value as Double)
-                            IrConstKind.Float -> buildConstExpression(source = null, ConstantValueKind.Float, argument.value as Float)
-                            IrConstKind.Int -> buildConstExpression(source = null, ConstantValueKind.Int, argument.value as Int)
-                            IrConstKind.Long -> buildConstExpression(source = null, ConstantValueKind.Long, argument.value as Long)
-                            IrConstKind.Null -> buildConstExpression(source = null, ConstantValueKind.Null, null)
-                            IrConstKind.Short -> buildConstExpression(source = null, ConstantValueKind.Short, argument.value as Short)
-                            IrConstKind.String -> buildConstExpression(source = null, ConstantValueKind.String, argument.value as String)
+                            IrConstKind.Boolean -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Boolean,
+                                argument.value as Boolean,
+                                setType = true
+                            )
+                            IrConstKind.Byte -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Byte,
+                                argument.value as Byte,
+                                setType = true
+                            )
+                            IrConstKind.Char -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Char,
+                                argument.value as Char,
+                                setType = true
+                            )
+                            IrConstKind.Double -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Double,
+                                argument.value as Double,
+                                setType = true
+                            )
+                            IrConstKind.Float -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Float,
+                                argument.value as Float,
+                                setType = true
+                            )
+                            IrConstKind.Int -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Int,
+                                argument.value as Int,
+                                setType = true
+                            )
+                            IrConstKind.Long -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Long,
+                                argument.value as Long,
+                                setType = true
+                            )
+                            IrConstKind.Null -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Null,
+                                value = null,
+                                setType = true
+                            )
+                            IrConstKind.Short -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.Short,
+                                argument.value as Short,
+                                setType = false
+                            )
+                            IrConstKind.String -> buildConstExpression(
+                                source = null,
+                                ConstantValueKind.String,
+                                argument.value as String,
+                                setType = false
+                            )
                         }
                     }
                 }

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaUtils.kt

@@ -63,15 +63,33 @@ internal fun Any?.createConstantOrError(session: FirSession): FirExpression {
 
 internal fun Any?.createConstantIfAny(session: FirSession, unsigned: Boolean = false): FirExpression? {
     return when (this) {
-        is Byte -> buildConstExpression(null, if (unsigned) ConstantValueKind.UnsignedByte else ConstantValueKind.Byte, this).setProperType(session)
-        is Short -> buildConstExpression(null, if (unsigned) ConstantValueKind.UnsignedShort else ConstantValueKind.Short, this).setProperType(session)
-        is Int -> buildConstExpression(null, if (unsigned) ConstantValueKind.UnsignedInt else ConstantValueKind.Int, this).setProperType(session)
-        is Long -> buildConstExpression(null, if (unsigned) ConstantValueKind.UnsignedLong else ConstantValueKind.Long, this).setProperType(session)
-        is Char -> buildConstExpression(null, ConstantValueKind.Char, this).setProperType(session)
-        is Float -> buildConstExpression(null, ConstantValueKind.Float, this).setProperType(session)
-        is Double -> buildConstExpression(null, ConstantValueKind.Double, this).setProperType(session)
-        is Boolean -> buildConstExpression(null, ConstantValueKind.Boolean, this).setProperType(session)
-        is String -> buildConstExpression(null, ConstantValueKind.String, this).setProperType(session)
+        is Byte -> buildConstExpression(
+            null, if (unsigned) ConstantValueKind.UnsignedByte else ConstantValueKind.Byte, this, setType = true
+        )
+        is Short -> buildConstExpression(
+            null, if (unsigned) ConstantValueKind.UnsignedShort else ConstantValueKind.Short, this, setType = true
+        )
+        is Int -> buildConstExpression(
+            null, if (unsigned) ConstantValueKind.UnsignedInt else ConstantValueKind.Int, this, setType = true
+        )
+        is Long -> buildConstExpression(
+            null, if (unsigned) ConstantValueKind.UnsignedLong else ConstantValueKind.Long, this, setType = true
+        )
+        is Char -> buildConstExpression(
+            null, ConstantValueKind.Char, this, setType = true
+        )
+        is Float -> buildConstExpression(
+            null, ConstantValueKind.Float, this, setType = true
+        )
+        is Double -> buildConstExpression(
+            null, ConstantValueKind.Double, this, setType = true
+        )
+        is Boolean -> buildConstExpression(
+            null, ConstantValueKind.Boolean, this, setType = true
+        )
+        is String -> buildConstExpression(
+            null, ConstantValueKind.String, this, setType = true
+        )
         is ByteArray -> toList().createArrayOfCall(session, ConstantValueKind.Byte)
         is ShortArray -> toList().createArrayOfCall(session, ConstantValueKind.Short)
         is IntArray -> toList().createArrayOfCall(session, ConstantValueKind.Int)
@@ -80,7 +98,9 @@ internal fun Any?.createConstantIfAny(session: FirSession, unsigned: Boolean = f
         is FloatArray -> toList().createArrayOfCall(session, ConstantValueKind.Float)
         is DoubleArray -> toList().createArrayOfCall(session, ConstantValueKind.Double)
         is BooleanArray -> toList().createArrayOfCall(session, ConstantValueKind.Boolean)
-        null -> buildConstExpression(null, ConstantValueKind.Null, null).setProperType(session)
+        null -> buildConstExpression(
+            null, ConstantValueKind.Null, null, setType = true
+        )
 
         else -> null
     }
@@ -99,14 +119,6 @@ private fun <T> List<T>.createArrayOfCall(session: FirSession, kind: ConstantVal
     }
 }
 
-private fun FirConstExpression<*>.setProperType(session: FirSession): FirConstExpression<*> {
-    val typeRef = buildResolvedTypeRef {
-        type = kind.expectedConeType(session)
-    }
-    replaceTypeRef(typeRef)
-    return this
-}
-
 // For now, it's supported only for RxJava3 annotations, see KT-53041
 fun extractNullabilityAnnotationOnBoundedWildcard(wildcardType: JavaWildcardType): JavaAnnotation? {
     require(wildcardType.bound != null) { "Nullability annotations on unbounded wildcards aren't supported" }

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilder.kt

@@ -406,13 +406,15 @@ abstract class AbstractRawFirBuilder<T>(val baseSession: FirSession, val context
                 buildConstExpression(
                     sourceElement,
                     ConstantValueKind.Boolean,
-                    convertedText as Boolean
+                    convertedText as Boolean,
+                    setType = false
                 )
             NULL ->
                 buildConstExpression(
                     sourceElement,
                     ConstantValueKind.Null,
-                    null
+                    null,
+                    setType = false
                 )
             else ->
                 throw AssertionError("Unknown literal type: $type, $text")
@@ -438,7 +440,8 @@ abstract class AbstractRawFirBuilder<T>(val baseSession: FirSession, val context
         return buildConstExpression(
             source.toFirSourceElement(),
             ConstantValueKind.IntegerLiteral,
-            convertedValue
+            convertedValue,
+            setType = false
         )
     }
 
@@ -457,11 +460,15 @@ abstract class AbstractRawFirBuilder<T>(val baseSession: FirSession, val context
                         OPEN_QUOTE, CLOSING_QUOTE -> continue@L
                         LITERAL_STRING_TEMPLATE_ENTRY -> {
                             sb.append(entry.asText)
-                            buildConstExpression(entry.toFirSourceElement(), ConstantValueKind.String, entry.asText)
+                            buildConstExpression(
+                                entry.toFirSourceElement(), ConstantValueKind.String, entry.asText, setType = false
+                            )
                         }
                         ESCAPE_STRING_TEMPLATE_ENTRY -> {
                             sb.append(entry.unescapedValue)
-                            buildConstExpression(entry.toFirSourceElement(), ConstantValueKind.String, entry.unescapedValue)
+                            buildConstExpression(
+                                entry.toFirSourceElement(), ConstantValueKind.String, entry.unescapedValue, setType = false
+                            )
                         }
                         SHORT_STRING_TEMPLATE_ENTRY, LONG_STRING_TEMPLATE_ENTRY -> {
                             hasExpressions = true
@@ -479,7 +486,7 @@ abstract class AbstractRawFirBuilder<T>(val baseSession: FirSession, val context
             }
             source = base?.toFirSourceElement()
             // Fast-pass if there is no non-const string expressions
-            if (!hasExpressions) return buildConstExpression(source, ConstantValueKind.String, sb.toString())
+            if (!hasExpressions) return buildConstExpression(source, ConstantValueKind.String, sb.toString(), setType = false)
         }
     }
 

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -355,7 +355,7 @@ fun <T> FirPropertyBuilder.generateAccessorsByDelegate(
                     type = context.dispatchReceiverTypesStack.last()
                 }
             }
-            else -> buildConstExpression(null, ConstantValueKind.Null, null)
+            else -> buildConstExpression(null, ConstantValueKind.Null, null, setType = false)
         }
 
     fun delegateAccess() = buildPropertyAccessExpression {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt

@@ -38,7 +38,7 @@ inline val FirAnnotation.unexpandedClassId: ClassId?
 
 fun <T> buildConstOrErrorExpression(source: KtSourceElement?, kind: ConstantValueKind<T>, value: T?, diagnostic: ConeDiagnostic): FirExpression =
     value?.let {
-        buildConstExpression(source, kind, it)
+        buildConstExpression(source, kind, it, setType = false)
     } ?: buildErrorExpression {
         this.source = source
         this.diagnostic = diagnostic

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/builder/FirConstExpressionBuilder.kt

@@ -18,18 +18,30 @@ fun <T> buildConstExpression(
     kind: ConstantValueKind<T>,
     value: T,
     annotations: MutableList<FirAnnotation> = mutableListOf(),
-    setType: Boolean = false
+    setType: Boolean
 ): FirConstExpression<T> {
     return FirConstExpressionImpl(source, annotations.toMutableOrEmpty(), kind, value).also {
         if (setType) {
-            when (value) {
-                null -> it.typeRef = FirImplicitNullableAnyTypeRef(null)
-                is String -> it.typeRef = FirImplicitStringTypeRef(null)
-                is Long -> it.typeRef = FirImplicitLongTypeRef(null)
-                is Int -> it.typeRef = FirImplicitIntTypeRef(null)
-                is Boolean -> it.typeRef = FirImplicitBooleanTypeRef(null)
-                is Double -> it.typeRef = FirImplicitDoubleTypeRef(null)
-                is Float -> it.typeRef = FirImplicitFloatTypeRef(null)
+            when (kind) {
+                ConstantValueKind.Boolean -> it.typeRef = FirImplicitBooleanTypeRef(null)
+                ConstantValueKind.Byte -> it.typeRef = FirImplicitByteTypeRef(null)
+                ConstantValueKind.Char -> it.typeRef = FirImplicitCharTypeRef(null)
+                ConstantValueKind.Double -> it.typeRef = FirImplicitDoubleTypeRef(null)
+                ConstantValueKind.Float -> it.typeRef = FirImplicitFloatTypeRef(null)
+                ConstantValueKind.Int -> it.typeRef = FirImplicitIntTypeRef(null)
+                ConstantValueKind.Long -> it.typeRef = FirImplicitLongTypeRef(null)
+                ConstantValueKind.Null -> it.typeRef = FirImplicitNullableAnyTypeRef(null)
+                ConstantValueKind.Short -> it.typeRef = FirImplicitShortTypeRef(null)
+                ConstantValueKind.String -> it.typeRef = FirImplicitStringTypeRef(null)
+                ConstantValueKind.UnsignedByte -> it.typeRef = FirImplicitUByteTypeRef(null)
+                ConstantValueKind.UnsignedInt -> it.typeRef = FirImplicitUIntTypeRef(null)
+                ConstantValueKind.UnsignedLong -> it.typeRef = FirImplicitULongTypeRef(null)
+                ConstantValueKind.UnsignedShort -> it.typeRef = FirImplicitUShortTypeRef(null)
+                ConstantValueKind.IntegerLiteral,
+                ConstantValueKind.UnsignedIntegerLiteral,
+                ConstantValueKind.Error,
+                -> {
+                }
             }
         }
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirImplicitBuiltinTypeRef.kt

@@ -106,6 +106,14 @@ class FirImplicitULongTypeRef(
     source: KtSourceElement?
 ) : FirImplicitBuiltinTypeRef(source, StandardClassIds.ULong)
 
+class FirImplicitUShortTypeRef(
+    source: KtSourceElement?
+) : FirImplicitBuiltinTypeRef(source, StandardClassIds.UShort)
+
+class FirImplicitUByteTypeRef(
+    source: KtSourceElement?
+) : FirImplicitBuiltinTypeRef(source, StandardClassIds.UByte)
+
 class FirImplicitNothingTypeRef(
     source: KtSourceElement?
 ) : FirImplicitBuiltinTypeRef(source, StandardClassIds.Nothing)
@@ -241,5 +249,7 @@ fun FirImplicitBuiltinTypeRef.withNewSource(newSource: KtSourceElement?): FirImp
             extensionReceiverTypeArgument = type.typeArguments[1],
             propertyTypeArgument = type.typeArguments[2]
         )
+        is FirImplicitUByteTypeRef -> FirImplicitUByteTypeRef(newSource)
+        is FirImplicitUShortTypeRef -> FirImplicitUShortTypeRef(newSource)
     }
 }