[FIR] Properly implement serializer extension for klib compilation

^KT-57654 Fixed

compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/klib/compilerPipeline.kt

@@ -31,7 +31,6 @@ import org.jetbrains.kotlin.fir.descriptors.FirModuleDescriptor
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 import org.jetbrains.kotlin.fir.pipeline.*
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.serialization.FirElementAwareSerializableStringTable
 import org.jetbrains.kotlin.fir.serialization.FirKLibSerializerExtension
 import org.jetbrains.kotlin.fir.serialization.serializeSingleFirFile
 import org.jetbrains.kotlin.incremental.components.LookupTracker
@@ -191,8 +190,9 @@ fun serializeFirKlib(
             scopeSession,
             actualizedExpectDeclarations,
             FirKLibSerializerExtension(
-                session, metadataVersion, FirElementAwareSerializableStringTable(),
+                session, metadataVersion,
                 ConstValueProviderImpl(fir2IrActualizedResult.components),
+                allowErrorTypes = false, exportKDoc = false
             ),
             moduleStructure.compilerConfiguration.languageVersionSettings,
         )

compiler/cli/src/org/jetbrains/kotlin/cli/metadata/FirMetadataSerializer.kt

@@ -36,7 +36,6 @@ import org.jetbrains.kotlin.fir.pipeline.ModuleCompilerAnalyzedOutput
 import org.jetbrains.kotlin.fir.pipeline.buildFirFromKtFiles
 import org.jetbrains.kotlin.fir.pipeline.buildFirViaLightTree
 import org.jetbrains.kotlin.fir.pipeline.resolveAndCheckFir
-import org.jetbrains.kotlin.fir.serialization.FirElementAwareSerializableStringTable
 import org.jetbrains.kotlin.fir.serialization.FirKLibSerializerExtension
 import org.jetbrains.kotlin.fir.serialization.serializeSingleFirFile
 import org.jetbrains.kotlin.library.SerializedMetadata
@@ -168,7 +167,8 @@ internal class FirMetadataSerializer(
                     scopeSession,
                     actualizedExpectDeclarations = null,
                     FirKLibSerializerExtension(
-                        session, metadataVersion, FirElementAwareSerializableStringTable(), constValueProvider = null
+                        session, metadataVersion, constValueProvider = null,
+                        allowErrorTypes = false, exportKDoc = false
                     ),
                     languageVersionSettings,
                 )

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

@@ -53,6 +53,8 @@ import org.jetbrains.kotlin.serialization.deserialization.ProtoEnumFlags
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
 import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
+import org.jetbrains.kotlin.utils.addIfNotNull
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class FirElementSerializer private constructor(
     private val session: FirSession,
@@ -60,7 +62,7 @@ class FirElementSerializer private constructor(
     private val containingDeclaration: FirDeclaration?,
     private val typeParameters: Interner<FirTypeParameter>,
     private val extension: FirSerializerExtension,
-    private val typeTable: MutableTypeTable,
+    val typeTable: MutableTypeTable,
     private val versionRequirementTable: MutableVersionRequirementTable?,
     private val serializeTypeTableToFunction: Boolean,
     private val typeApproximator: AbstractTypeApproximator,
@@ -738,6 +740,7 @@ class FirElementSerializer private constructor(
         isDefinitelyNotNullType: Boolean,
     ): ProtoBuf.Type.Builder {
         val builder = ProtoBuf.Type.newBuilder()
+        val typeAnnotations = mutableListOf<FirAnnotation>()
         when (type) {
             is ConeDefinitelyNotNullType -> return typeProto(type.original, toSuper, correspondingTypeRef, isDefinitelyNotNullType = true)
             is ConeErrorType -> {
@@ -766,12 +769,14 @@ class FirElementSerializer private constructor(
                 }
                 fillFromPossiblyInnerType(builder, type)
                 if (type.hasContextReceivers) {
-                    serializeAnnotationFromAttribute(
+                    typeAnnotations.addIfNotNull(
-                        correspondingTypeRef?.annotations, CompilerConeAttributes.ContextFunctionTypeParams.ANNOTATION_CLASS_ID, builder,
+                        createAnnotationFromAttribute(
-                        argumentMapping = buildAnnotationArgumentMapping {
+                            correspondingTypeRef?.annotations, CompilerConeAttributes.ContextFunctionTypeParams.ANNOTATION_CLASS_ID,
-                            this.mapping[StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME] =
+                            argumentMapping = buildAnnotationArgumentMapping {
-                                buildConstExpression(source = null, ConstantValueKind.Int, type.contextReceiversNumberForFunctionType)
+                                this.mapping[StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME] =
-                        }
+                                    buildConstExpression(source = null, ConstantValueKind.Int, type.contextReceiversNumberForFunctionType)
+                            }
+                        )
                     )
                 }
             }
@@ -816,23 +821,21 @@ class FirElementSerializer private constructor(
         val extensionAttributes = mutableListOf<ConeAttribute<*>>()
         for (attribute in type.attributes) {
             when {
-                attribute is CustomAnnotationTypeAttribute ->
+                attribute is CustomAnnotationTypeAttribute -> typeAnnotations.addAll(attribute.annotations.nonSourceAnnotations(session))
-                    for (annotation in attribute.annotations.nonSourceAnnotations(session)) {
-                        extension.serializeTypeAnnotation(annotation, builder)
-                    }
                 attribute.key in CompilerConeAttributes.classIdByCompilerAttributeKey ->
-                    serializeCompilerDefinedTypeAttribute(builder, attribute)
+                    typeAnnotations.add(createAnnotationForCompilerDefinedTypeAttribute(attribute))
                 else -> extensionAttributes += attribute
             }
         }
 
         for (attributeExtension in session.extensionService.typeAttributeExtensions) {
             for (attribute in extensionAttributes) {
-                val annotation = attributeExtension.convertAttributeToAnnotation(attribute) ?: continue
+                typeAnnotations.addIfNotNull(attributeExtension.convertAttributeToAnnotation(attribute))
-                extension.serializeTypeAnnotation(annotation, builder)
             }
         }
 
+        extension.serializeTypeAnnotations(typeAnnotations, builder)
+
         // TODO: abbreviated type
 //        val abbreviation = type.getAbbreviatedType()?.abbreviation
 //        if (abbreviation != null) {
@@ -846,11 +849,8 @@ class FirElementSerializer private constructor(
         return builder
     }
 
-    private fun serializeCompilerDefinedTypeAttribute(
+    private fun createAnnotationForCompilerDefinedTypeAttribute(attribute: ConeAttribute<*>): FirAnnotation {
-        builder: ProtoBuf.Type.Builder,
+        return buildAnnotation {
-        attribute: ConeAttribute<*>
-    ) {
-        val annotation = buildAnnotation {
             annotationTypeRef = buildResolvedTypeRef {
                 this.type = ConeClassLikeTypeImpl(
                     CompilerConeAttributes.classIdByCompilerAttributeKey.getValue(attribute.key).toLookupTag(),
@@ -860,26 +860,22 @@ class FirElementSerializer private constructor(
             }
             argumentMapping = FirEmptyAnnotationArgumentMapping
         }
-        extension.serializeTypeAnnotation(annotation, builder)
     }
 
-    private fun serializeAnnotationFromAttribute(
+    private fun createAnnotationFromAttribute(
         existingAnnotations: List<FirAnnotation>?,
         classId: ClassId,
-        builder: ProtoBuf.Type.Builder,
         argumentMapping: FirAnnotationArgumentMapping = FirEmptyAnnotationArgumentMapping,
-    ) {
+    ): FirAnnotation? {
-        if (existingAnnotations?.any { it.annotationTypeRef.coneTypeSafe<ConeClassLikeType>()?.classId == classId } != true) {
+        return runIf(existingAnnotations?.any { it.annotationTypeRef.coneTypeSafe<ConeClassLikeType>()?.classId == classId } != true) {
-            extension.serializeTypeAnnotation(
+            buildAnnotation {
-                buildAnnotation {
+                annotationTypeRef = buildResolvedTypeRef {
-                    annotationTypeRef = buildResolvedTypeRef {
+                    this.type = classId.constructClassLikeType(
-                        this.type = classId.constructClassLikeType(
+                        emptyArray(), isNullable = false
-                            emptyArray(), isNullable = false
+                    )
-                        )
+                }
-                    }
+                this.argumentMapping = argumentMapping
-                    this.argumentMapping = argumentMapping
+            }
-                }, builder
-            )
         }
     }
 

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

@@ -0,0 +1,123 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.providers.firProvider
+import org.jetbrains.kotlin.fir.serialization.constant.ConstValueProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+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.library.metadata.KlibMetadataSerializerProtocol
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
+import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
+import org.jetbrains.kotlin.protobuf.GeneratedMessageLite
+import org.jetbrains.kotlin.serialization.deserialization.DYNAMIC_TYPE_DESERIALIZER_ID
+
+class FirKLibSerializerExtension(
+    override val session: FirSession,
+    override val metadataVersion: BinaryVersion,
+    constValueProvider: ConstValueProvider?,
+    private val allowErrorTypes: Boolean,
+    private val exportKDoc: Boolean
+) : FirSerializerExtensionBase(KlibMetadataSerializerProtocol, constValueProvider) {
+    override fun shouldUseTypeTable(): Boolean = true
+
+    override fun serializeFlexibleType(type: ConeFlexibleType, lowerProto: ProtoBuf.Type.Builder, upperProto: ProtoBuf.Type.Builder) {
+        lowerProto.flexibleTypeCapabilitiesId = stringTable.getStringIndex(DYNAMIC_TYPE_DESERIALIZER_ID)
+    }
+
+    override fun serializeErrorType(type: ConeErrorType, builder: ProtoBuf.Type.Builder) {
+        if (!allowErrorTypes) super.serializeErrorType(type, builder)
+    }
+
+    override fun serializeClass(
+        klass: FirClass,
+        proto: ProtoBuf.Class.Builder,
+        versionRequirementTable: MutableVersionRequirementTable,
+        childSerializer: FirElementSerializer
+    ) {
+        klass.setFileId(proto, KlibMetadataProtoBuf.classFile)
+        klass.setKDoc(proto, KlibMetadataProtoBuf.classKdoc)
+        super.serializeClass(klass, proto, versionRequirementTable, childSerializer)
+        childSerializer.typeTable.serialize()?.let { proto.mergeTypeTable(it) }
+    }
+
+    override fun serializeConstructor(
+        constructor: FirConstructor,
+        proto: ProtoBuf.Constructor.Builder,
+        childSerializer: FirElementSerializer
+    ) {
+        constructor.setKDoc(proto, KlibMetadataProtoBuf.constructorKdoc)
+        super.serializeConstructor(constructor, proto, childSerializer)
+    }
+
+    override fun serializeProperty(
+        property: FirProperty,
+        proto: ProtoBuf.Property.Builder,
+        versionRequirementTable: MutableVersionRequirementTable?,
+        childSerializer: FirElementSerializer
+    ) {
+        property.setFileId(proto, KlibMetadataProtoBuf.propertyFile)
+        property.setKDoc(proto, KlibMetadataProtoBuf.propertyKdoc)
+        super.serializeProperty(property, proto, versionRequirementTable, childSerializer)
+    }
+
+    override fun serializeFunction(
+        function: FirFunction,
+        proto: ProtoBuf.Function.Builder,
+        versionRequirementTable: MutableVersionRequirementTable?,
+        childSerializer: FirElementSerializer
+    ) {
+        function.setFileId(proto, KlibMetadataProtoBuf.functionFile)
+        function.setKDoc(proto, KlibMetadataProtoBuf.functionKdoc)
+        super.serializeFunction(function, proto, versionRequirementTable, childSerializer)
+    }
+
+    private val firProvider = session.firProvider
+
+    @Suppress("Reformat")
+    private fun <
+        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
+        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+    > FirDeclaration.setKDoc(
+        proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+        extension: GeneratedMessageLite.GeneratedExtension<MessageType, String>,
+    ) {
+        if (exportKDoc) {
+            findKDocString()?.let { proto.setExtension(extension, it) }
+        }
+    }
+
+    private fun FirDeclaration.findKDocString(): String? {
+        // TODO: KT-56090
+        return null
+    }
+
+    @Suppress("Reformat")
+    private fun <
+        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
+        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+    > FirDeclaration.setFileId(
+        proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+        extension: GeneratedMessageLite.GeneratedExtension<MessageType, Int>,
+    ) {
+        declarationFileId(this)?.let { proto.setExtension(extension, it) }
+    }
+
+    private fun declarationFileId(declaration: FirDeclaration): Int? {
+        val file = when (val symbol = declaration.symbol) {
+            is FirCallableSymbol<*> -> firProvider.getFirCallableContainerFile(symbol)
+            is FirClassLikeSymbol<*> -> firProvider.getFirClassifierContainerFileIfAny(symbol)
+            else -> null
+        } ?: return null
+        return stringTable.getStringIndex(file.name)
+    }
+}

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

@@ -8,10 +8,8 @@ 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
@@ -26,8 +24,6 @@ abstract class FirSerializerExtension {
 
     val annotationSerializer by lazy { FirAnnotationSerializer(session, stringTable) }
 
-    abstract val constValueProvider: ConstValueProvider?
-
     open fun shouldUseTypeTable(): Boolean = false
     open fun shouldUseNormalizedVisibility(): Boolean = false
 
@@ -74,7 +70,7 @@ abstract class FirSerializerExtension {
     open fun serializeFlexibleType(type: ConeFlexibleType, lowerProto: ProtoBuf.Type.Builder, upperProto: ProtoBuf.Type.Builder) {
     }
 
-    open fun serializeTypeAnnotation(annotation: FirAnnotation, proto: ProtoBuf.Type.Builder) {
+    open fun serializeTypeAnnotations(annotations: List<FirAnnotation>, proto: ProtoBuf.Type.Builder) {
     }
 
     open fun serializeTypeParameter(typeParameter: FirTypeParameter, proto: ProtoBuf.TypeParameter.Builder) {
@@ -87,16 +83,10 @@ 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
 
     interface ClassMembersProducer {
         fun getCallableMembers(klass: FirClass): Collection<FirCallableDeclaration>
     }
-
 }

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

@@ -0,0 +1,146 @@
+/*
+ * 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
+
+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.ConstValueProvider
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.protobuf.GeneratedMessageLite
+import org.jetbrains.kotlin.serialization.SerializerExtensionProtocol
+
+abstract class FirSerializerExtensionBase(
+    val protocol: SerializerExtensionProtocol,
+    private val constValueProvider: ConstValueProvider?
+) : FirSerializerExtension() {
+    final override val stringTable = FirElementAwareSerializableStringTable()
+
+    override fun serializePackage(packageFqName: FqName, proto: ProtoBuf.Package.Builder) {
+        proto.setExtension(protocol.packageFqName, stringTable.getPackageFqNameIndex(packageFqName))
+    }
+
+    override fun serializeClass(
+        klass: FirClass,
+        proto: ProtoBuf.Class.Builder,
+        versionRequirementTable: MutableVersionRequirementTable,
+        childSerializer: FirElementSerializer
+    ) {
+        klass.serializeAnnotations(proto, protocol.classAnnotation)
+    }
+
+    override fun serializeConstructor(
+        constructor: FirConstructor,
+        proto: ProtoBuf.Constructor.Builder,
+        childSerializer: FirElementSerializer
+    ) {
+        constructor.serializeAnnotations(proto, protocol.constructorAnnotation)
+    }
+
+    override fun serializeFunction(
+        function: FirFunction,
+        proto: ProtoBuf.Function.Builder,
+        versionRequirementTable: MutableVersionRequirementTable?,
+        childSerializer: FirElementSerializer
+    ) {
+        function.serializeAnnotations(proto, protocol.functionAnnotation)
+        function.receiverParameter?.serializeAnnotations(proto, protocol.functionExtensionReceiverAnnotation)
+    }
+
+    override fun serializeProperty(
+        property: FirProperty,
+        proto: ProtoBuf.Property.Builder,
+        versionRequirementTable: MutableVersionRequirementTable?,
+        childSerializer: FirElementSerializer
+    ) {
+        val regularPropertyAnnotations = mutableListOf<FirAnnotation>()
+        val fieldPropertyAnnotations = mutableListOf<FirAnnotation>()
+        val delegatePropertyAnnotations = mutableListOf<FirAnnotation>()
+
+        for (annotation in property.nonSourceAnnotations(session)) {
+            val destination = when (annotation.useSiteTarget) {
+                AnnotationUseSiteTarget.FIELD -> fieldPropertyAnnotations
+                AnnotationUseSiteTarget.PROPERTY_DELEGATE_FIELD -> delegatePropertyAnnotations
+                else -> regularPropertyAnnotations
+            }
+            destination += annotation
+        }
+
+        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)
+
+        constValueProvider?.buildValueProtoBufIfPropertyIsConst(property, annotationSerializer)?.let { constProtoBuf ->
+            proto.setExtension(protocol.compileTimeValue, constProtoBuf)
+        }
+    }
+
+    override fun serializeEnumEntry(enumEntry: FirEnumEntry, proto: ProtoBuf.EnumEntry.Builder) {
+        enumEntry.serializeAnnotations(proto, protocol.enumEntryAnnotation)
+    }
+
+    override fun serializeValueParameter(parameter: FirValueParameter, proto: ProtoBuf.ValueParameter.Builder) {
+        parameter.serializeAnnotations(proto, protocol.parameterAnnotation)
+    }
+
+    override fun serializeTypeAnnotations(annotations: List<FirAnnotation>, proto: ProtoBuf.Type.Builder) {
+        annotations.serializeAnnotations(proto, protocol.typeAnnotation)
+    }
+
+    override fun serializeTypeParameter(typeParameter: FirTypeParameter, proto: ProtoBuf.TypeParameter.Builder) {
+        typeParameter.serializeAnnotations(proto, protocol.typeParameterAnnotation)
+    }
+
+    override val customClassMembersProducer: ClassMembersProducer?
+        get() = super.customClassMembersProducer
+
+    @Suppress("Reformat")
+    private fun <
+        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
+        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+    > FirAnnotationContainer.serializeAnnotations(
+        proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+        extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?,
+    ) {
+        if (extension == null) return
+        this.nonSourceAnnotations(session).serializeAnnotations(proto, extension)
+    }
+
+    @Suppress("Reformat")
+    private fun <
+        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
+        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+    > List<FirAnnotation>.serializeAnnotations(
+        proto: GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+        extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<ProtoBuf.Annotation>>?,
+    ) {
+        if (extension == null) return
+        for (annotation in this) {
+            proto.addExtensionOrNull(extension, annotationSerializer.serializeAnnotation(annotation))
+        }
+    }
+
+    @Suppress("Reformat")
+    private fun <
+        MessageType : GeneratedMessageLite.ExtendableMessage<MessageType>,
+        BuilderType : GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>,
+        Type
+    > GeneratedMessageLite.ExtendableBuilder<MessageType, BuilderType>.addExtensionOrNull(
+        extension: GeneratedMessageLite.GeneratedExtension<MessageType, List<Type>>,
+        value: Type?
+    ) {
+        if (value != null) {
+            addExtension(extension, value)
+        }
+    }
+}

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

@@ -24,17 +24,26 @@ import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 import org.jetbrains.kotlin.types.ConstantValueKind
 
-internal fun FirExpression.toConstantValue(session: FirSession): ConstantValue<*>? = accept(FirToConstantValueTransformer, session)
+internal fun FirExpression.toConstantValue(session: FirSession): ConstantValue<*>? {
+    return accept(FirToConstantValueTransformerUnsafe, session)
+}
 
-internal object FirToConstantValueTransformer : FirDefaultVisitor<ConstantValue<*>?, FirSession>() {
+private object FirToConstantValueTransformerSafe : FirToConstantValueTransformer(failOnNonConst = false)
+private object FirToConstantValueTransformerUnsafe : FirToConstantValueTransformer(failOnNonConst = true)
+
+private abstract class FirToConstantValueTransformer(
+    private val failOnNonConst: Boolean
+) : FirDefaultVisitor<ConstantValue<*>?, FirSession>() {
     override fun visitElement(
         element: FirElement,
         data: FirSession
     ): ConstantValue<*>? {
-        error("Illegal element as annotation argument: ${element::class.qualifiedName} -> ${element.render()}")
+        if (failOnNonConst) {
+            error("Illegal element as annotation argument: ${element::class.qualifiedName} -> ${element.render()}")
+        }
+        return null
     }
 
-    @OptIn(ExperimentalUnsignedTypes::class)
     override fun <T> visitConstExpression(
         constExpression: FirConstExpression<T>,
         data: FirSession

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

@@ -13,10 +13,7 @@ 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
 
@@ -65,23 +62,6 @@ fun serializeSingleFirFile(
     )
 }
 
-open class FirKLibSerializerExtension(
+class FirElementAwareSerializableStringTable : FirElementAwareStringTable, SerializableStringTable() {
-    override val session: FirSession,
+    override fun getLocalClassIdReplacement(firClass: FirClass): ClassId = ClassId.topLevel(StandardNames.FqNames.any.toSafe())
-    override val metadataVersion: BinaryVersion,
-    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())
 }

compiler/fir/fir2ir/jvm-backend/src/org/jetbrains/kotlin/fir/backend/jvm/FirJvmSerializerExtension.kt

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.config.JvmDefaultMode
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.backend.ConstValueProviderImpl
 import org.jetbrains.kotlin.fir.backend.Fir2IrComponents
 import org.jetbrains.kotlin.fir.backend.FirMetadataSource
 import org.jetbrains.kotlin.fir.declarations.*
@@ -25,7 +24,6 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.toFirRegularClassSymbol
 import org.jetbrains.kotlin.fir.serialization.*
-import org.jetbrains.kotlin.fir.serialization.constant.ConstValueProvider
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.ir.declarations.MetadataSource
 import org.jetbrains.kotlin.load.kotlin.NON_EXISTENT_CLASS_NAME
@@ -60,7 +58,6 @@ class FirJvmSerializerExtension(
     override val metadataVersion: BinaryVersion,
     private val jvmDefaultMode: JvmDefaultMode,
     override val stringTable: FirElementAwareStringTable,
-    override val constValueProvider: ConstValueProvider?
 ) : FirSerializerExtension() {
 
     constructor(
@@ -76,7 +73,7 @@ class FirJvmSerializerExtension(
         session, bindings, metadata, localDelegatedProperties, approximator, components.scopeSession,
         state.globalSerializationBindings, state.useTypeTableInSerializer, state.moduleName, state.classBuilderMode,
         state.isParamAssertionsDisabled, state.unifiedNullChecks, state.metadataVersion, state.jvmDefaultMode,
-        FirJvmElementAwareStringTable(typeMapper, components), ConstValueProviderImpl(components),
+        FirJvmElementAwareStringTable(typeMapper, components),
     )
 
     override fun shouldUseTypeTable(): Boolean = useTypeTable
@@ -167,8 +164,10 @@ class FirJvmSerializerExtension(
         }
     }
 
-    override fun serializeTypeAnnotation(annotation: FirAnnotation, proto: ProtoBuf.Type.Builder) {
+    override fun serializeTypeAnnotations(annotations: List<FirAnnotation>, proto: ProtoBuf.Type.Builder) {
-        proto.addExtension(JvmProtoBuf.typeAnnotation, annotationSerializer.serializeAnnotation(annotation))
+        for (annotation in annotations) {
+            proto.addExtension(JvmProtoBuf.typeAnnotation, annotationSerializer.serializeAnnotation(annotation))
+        }
     }
 
 

compiler/fir/fir2ir/jvm-backend/src/org/jetbrains/kotlin/fir/backend/jvm/FirMetadataSerializer.kt

@@ -105,7 +105,6 @@ fun makeLocalFirMetadataSerializerForMetadataSource(
         configuration.metadataVersion(session.languageVersionSettings.languageVersion),
         session.languageVersionSettings.getFlag(JvmAnalysisFlags.jvmDefaultMode),
         stringTable,
-        constValueProvider = null,
     )
     return FirMetadataSerializer(
         globalSerializationBindings,

compiler/testData/codegen/box/innerNested/considerPropertyInitValueFromNestedClass.kt

@@ -1,4 +1,5 @@
-// IGNORE_BACKEND_K2: NATIVE
+// IGNORE_BACKEND_K2: NATIVE, JS_IR
+//   Ignore reason: KT-57611
 // MODULE: lib
 // FILE: Class.kt
 
@@ -18,4 +19,4 @@ import Class
 @Ann("${Class.Obj.Const}+")
 fun f(): String = "OK"
 
-fun box() = f()
+fun box() = f()

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/Fir2IrJsResultsConverter.kt

@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.fir.backend.jvm.FirJvmKotlinMangler
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.descriptors.FirModuleDescriptor
-import org.jetbrains.kotlin.fir.serialization.FirElementAwareSerializableStringTable
 import org.jetbrains.kotlin.fir.serialization.FirKLibSerializerExtension
 import org.jetbrains.kotlin.fir.serialization.serializeSingleFirFile
 import org.jetbrains.kotlin.incremental.components.LookupTracker
@@ -124,7 +123,9 @@ class Fir2IrJsResultsConverter(
                 components.scopeSession,
                 actualizedExpectDeclarations,
                 FirKLibSerializerExtension(
-                    components.session, metadataVersion, FirElementAwareSerializableStringTable(), ConstValueProviderImpl(components),
+                    components.session, metadataVersion,
+                    ConstValueProviderImpl(components),
+                    allowErrorTypes = false, exportKDoc = false
                 ),
                 configuration.languageVersionSettings,
             )

js/js.translator/testData/box/inlineMultiModule/extensionLambda.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_K2: JS_IR
 // EXPECTED_REACHABLE_NODES: 1285
 // MODULE: lib
 // FILE: lib.kt
@@ -25,4 +24,4 @@ fun box(): String {
     assertEquals(6, test(A(2), 3))
 
     return "OK"
-}
+}

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/FirSerializer.kt

@@ -12,26 +12,19 @@ import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.backend.ConstValueProviderImpl
 import org.jetbrains.kotlin.fir.backend.extractFirDeclarations
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.serialization.*
-import org.jetbrains.kotlin.fir.serialization.constant.ConstValueProvider
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.IrMessageLogger
 import org.jetbrains.kotlin.konan.library.KonanLibrary
 import org.jetbrains.kotlin.library.SerializedIrFile
-import org.jetbrains.kotlin.library.metadata.KlibMetadataProtoBuf
 import org.jetbrains.kotlin.library.metadata.resolver.TopologicalLibraryOrder
 import org.jetbrains.kotlin.metadata.ProtoBuf
-import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
-import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
-import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.utils.toMetadataVersion
 
 internal fun PhaseContext.firSerializer(
@@ -69,8 +62,10 @@ internal fun PhaseContext.firSerializer(
                 session,
                 scopeSession,
                 actualizedFirDeclarations,
-                FirNativeKLibSerializerExtension(
+                FirKLibSerializerExtension(
-                        session, metadataVersion, FirElementAwareSerializableStringTable(), ConstValueProviderImpl(input.components),
+                        session, metadataVersion,
+                        ConstValueProviderImpl(input.components),
+                        allowErrorTypes = false, exportKDoc = false
                 ),
                 configuration.languageVersionSettings,
         )
@@ -134,117 +129,3 @@ internal fun PhaseContext.serializeNativeModule(
 
     return SerializerOutput(serializedMetadata, serializedIr, null, dependencies)
 }
-
-class FirNativeKLibSerializerExtension(
-        override val session: FirSession,
-        override val metadataVersion: BinaryVersion,
-        override val stringTable: FirElementAwareSerializableStringTable,
-        override val constValueProvider: ConstValueProvider?
-) : FirKLibSerializerExtension(session, metadataVersion, stringTable, constValueProvider) {
-    private fun declarationFileId(declaration: FirMemberDeclaration): Int? {
-        val fileName = declaration.source.psi?.containingFile?.name ?: return null
-        return stringTable.getStringIndex(fileName)
-    }
-
-    override fun serializeFunction(
-            function: FirFunction,
-            proto: ProtoBuf.Function.Builder,
-            versionRequirementTable: MutableVersionRequirementTable?,
-            childSerializer: FirElementSerializer
-    ) {
-        // inspired by KlibMetadataSerializerExtension.serializeFunction
-        declarationFileId(function)?.let { proto.setExtension(KlibMetadataProtoBuf.functionFile, it) }
-        function.nonSourceAnnotations(session).forEach {
-            proto.addExtension(KlibMetadataProtoBuf.functionAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        function.receiverParameter?.nonSourceAnnotations(session)?.forEach {
-            proto.addExtension(KlibMetadataProtoBuf.functionExtensionReceiverAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        // TODO KT-56090 Serialize KDocString
-        super.serializeFunction(function, proto, versionRequirementTable, childSerializer)
-    }
-
-    override fun serializeValueParameter(parameter: FirValueParameter, proto: ProtoBuf.ValueParameter.Builder) {
-        parameter.nonSourceAnnotations(session).forEach {
-            proto.addExtension(KlibMetadataProtoBuf.parameterAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        super.serializeValueParameter(parameter, proto)
-    }
-
-    override fun serializeProperty(
-            property: FirProperty,
-            proto: ProtoBuf.Property.Builder,
-            versionRequirementTable: MutableVersionRequirementTable?,
-            childSerializer: FirElementSerializer
-    ) {
-        // inspired by KlibMetadataSerializerExtension.serializeProperty
-        declarationFileId(property)?.let { proto.setExtension(KlibMetadataProtoBuf.propertyFile, it) }
-        property.nonSourceAnnotations(session).forEach {
-            val extension = when (it.useSiteTarget) {  // Revise this code after KT-54385
-                AnnotationUseSiteTarget.FIELD -> KlibMetadataProtoBuf.propertyBackingFieldAnnotation
-                AnnotationUseSiteTarget.PROPERTY_DELEGATE_FIELD -> KlibMetadataProtoBuf.propertyDelegatedFieldAnnotation
-                else -> KlibMetadataProtoBuf.propertyAnnotation
-            }
-            proto.addExtension(extension, annotationSerializer.serializeAnnotation(it))
-        }
-        property.receiverParameter?.nonSourceAnnotations(session)?.forEach {
-            proto.addExtension(KlibMetadataProtoBuf.propertyExtensionReceiverAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        property.getter?.nonSourceAnnotations(session)?.forEach {
-            proto.addExtension(KlibMetadataProtoBuf.propertyGetterAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        property.setter?.nonSourceAnnotations(session)?.forEach {
-            proto.addExtension(KlibMetadataProtoBuf.propertySetterAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-
-        serializeConstant(property, proto)
-
-        // TODO KT-56090 Serialize KDocString
-        super.serializeProperty(property, proto, versionRequirementTable, childSerializer)
-    }
-
-    override fun serializeClass(
-            klass: FirClass,
-            proto: ProtoBuf.Class.Builder,
-            versionRequirementTable: MutableVersionRequirementTable,
-            childSerializer: FirElementSerializer
-    ) {
-        declarationFileId(klass)?.let { proto.setExtension(KlibMetadataProtoBuf.classFile, it) }
-        klass.nonSourceAnnotations(session).forEach {
-            proto.addExtension(KlibMetadataProtoBuf.classAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        // TODO KT-56090 Serialize KDocString
-        super.serializeClass(klass, proto, versionRequirementTable, childSerializer)
-    }
-
-    override fun serializeConstructor(
-            constructor: FirConstructor,
-            proto: ProtoBuf.Constructor.Builder,
-            childSerializer: FirElementSerializer
-    ) {
-        constructor.nonSourceAnnotations(session).forEach {
-            proto.addExtension(KlibMetadataProtoBuf.constructorAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        // TODO KT-56090 Serialize KDocString
-        super.serializeConstructor(constructor, proto, childSerializer)
-    }
-
-    override fun serializeEnumEntry(enumEntry: FirEnumEntry, proto: ProtoBuf.EnumEntry.Builder) {
-        enumEntry.nonSourceAnnotations(session).forEach {
-            proto.addExtension(KlibMetadataProtoBuf.enumEntryAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        super.serializeEnumEntry(enumEntry, proto)
-    }
-
-    override fun serializeTypeAnnotation(annotation: FirAnnotation, proto: ProtoBuf.Type.Builder) {
-        proto.addExtension(KlibMetadataProtoBuf.typeAnnotation, annotationSerializer.serializeAnnotation(annotation))
-        super.serializeTypeAnnotation(annotation, proto)
-    }
-
-    override fun serializeTypeParameter(typeParameter: FirTypeParameter, proto: ProtoBuf.TypeParameter.Builder) {
-        typeParameter.nonSourceAnnotations(session).forEach {
-            proto.addExtension(KlibMetadataProtoBuf.typeParameterAnnotation, annotationSerializer.serializeAnnotation(it))
-        }
-        super.serializeTypeParameter(typeParameter, proto)
-    }
-}

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/K2Tests.kt

@@ -53,7 +53,7 @@ class CustomK2Tests : KGPBaseTest() {
     }
 
     @GradleTest
-    @DisplayName("HMPP compilation without JS target. KT-57376, KT-57377, KT-57635")
+    @DisplayName("HMPP compilation without JS target. KT-57376, KT-57377, KT-57635, KT-57654")
     fun testHmppCompilationWithoutJsTarget(gradleVersion: GradleVersion) {
         with(project("k2-mpp-without-js", gradleVersion)) {
             val taskToExecute = ":compileIntermediateMainKotlinMetadata"

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/resources/testProject/k2-mpp-without-js/src/commonMain/kotlin/kt57654.common.kt

@@ -0,0 +1,8 @@
+// ISSUE: KT-57654
+package kt57654
+
+class A {
+    fun foo() {}
+}
+
+fun useA(block: A.() -> Unit) {}

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/resources/testProject/k2-mpp-without-js/src/intermediateMain/kotlin/kt57654.intermediate.kt

@@ -0,0 +1,8 @@
+// ISSUE: KT-57654
+package kt57654
+
+fun test() {
+    useA {
+        foo()
+    }
+}