FIR JS: Refactor Klib serialization for FIR support

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2022 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.diagnostics.ConeIntermediateDiagnostic
+import org.jetbrains.kotlin.fir.languageVersionSettings
+import org.jetbrains.kotlin.fir.types.ConeErrorType
+import org.jetbrains.kotlin.fir.types.typeContext
+import org.jetbrains.kotlin.types.AbstractTypeApproximator
+import org.jetbrains.kotlin.types.model.SimpleTypeMarker
+
+class TypeApproximatorForMetadataSerializer(session: FirSession) :
+    AbstractTypeApproximator(session.typeContext, session.languageVersionSettings) {
+
+    override fun createErrorType(debugName: String): SimpleTypeMarker {
+        return ConeErrorType(ConeIntermediateDiagnostic(debugName))
+    }
+}
+

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

@@ -21,13 +21,13 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.builder.buildProperty
 import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameterCopy
-import org.jetbrains.kotlin.fir.diagnostics.ConeIntermediateDiagnostic
 import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
 import org.jetbrains.kotlin.fir.serialization.FirElementAwareStringTable
 import org.jetbrains.kotlin.fir.serialization.FirElementSerializer
+import org.jetbrains.kotlin.fir.serialization.TypeApproximatorForMetadataSerializer
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.ir.declarations.IrClass
@@ -39,7 +39,6 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.protobuf.MessageLite
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
-import org.jetbrains.kotlin.types.model.SimpleTypeMarker
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.Method
 
@@ -189,14 +188,6 @@ internal fun makeElementSerializer(
         else -> null
     }
 
-internal class TypeApproximatorForMetadataSerializer(session: FirSession) :
-    AbstractTypeApproximator(session.typeContext, session.languageVersionSettings) {
-
-    override fun createErrorType(debugName: String): SimpleTypeMarker {
-        return ConeErrorType(ConeIntermediateDiagnostic(debugName))
-    }
-}
-
 private fun FirFunction.copyToFreeAnonymousFunction(approximator: AbstractTypeApproximator): FirAnonymousFunction {
     val function = this
     return buildAnonymousFunction {

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/metadata/KlibMetadataSerializer.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.serialization.ApproximatingStringTable
 import org.jetbrains.kotlin.serialization.DescriptorSerializer
+import org.jetbrains.kotlin.serialization.SerializableStringTable
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPropertyDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedSimpleFunctionDescriptor
@@ -66,30 +67,6 @@ abstract class KlibMetadataSerializer(
     }
 
 
-    private fun SerializerContext.buildFragment(
-        packageProto: ProtoBuf.Package,
-        classesProto: List<Pair<ProtoBuf.Class, Int>>,
-        fqName: FqName,
-        isEmpty: Boolean
-    ): ProtoBuf.PackageFragment {
-
-        val (stringTableProto, nameTableProto) = serializerExtension.stringTable.buildProto()
-
-        return ProtoBuf.PackageFragment.newBuilder()
-            .setPackage(packageProto)
-            .addAllClass_(classesProto.map { it.first })
-            .setStrings(stringTableProto)
-            .setQualifiedNames(nameTableProto)
-            .also { packageFragment ->
-                classesProto.forEach {
-                    packageFragment.addExtension(KlibMetadataProtoBuf.className, it.second )
-                }
-                packageFragment.setExtension(KlibMetadataProtoBuf.isEmpty, isEmpty)
-                packageFragment.setExtension(KlibMetadataProtoBuf.fqName, fqName.asString())
-            }
-            .build()
-    }
-
     private fun serializeClass(packageName: FqName,
                                classDescriptor: ClassDescriptor): List<Pair<ProtoBuf.Class, Int>> {
         with(serializerContext) {
@@ -169,11 +146,12 @@ abstract class KlibMetadataSerializer(
                     else
                         buildPackageProto(fqName, nonCassDescriptors)
 
-                    buildFragment(
+                    buildKlibPackageFragment(
                         packageProto,
                         serializeClasses(fqName, classifierDescriptors),
                         fqName,
-                        topLevelDescriptors.isEmpty() && classifierDescriptors.isEmpty()
+                        topLevelDescriptors.isEmpty() && classifierDescriptors.isEmpty(),
+                        serializerExtension.stringTable
                     )
                 }
             )
@@ -195,33 +173,36 @@ abstract class KlibMetadataSerializer(
                     if (typeAliases.isNotEmpty()) buildPackageProto(fqName, typeAliases)
                     else emptyPackageProto()
 
-                buildFragment(
+                buildKlibPackageFragment(
                     packageProto,
                     classesProto,
                     fqName,
-                    descriptors.isEmpty()
+                    descriptors.isEmpty(),
+                    serializerExtension.stringTable
                 )
             }
         }
 
         result += topLevelDescriptors.maybeChunked(TOP_LEVEL_DECLARATION_COUNT_PER_FILE) { descriptors ->
             withNewContext {
-                buildFragment(
+                buildKlibPackageFragment(
                     buildPackageProto(fqName, descriptors),
                     emptyList(),
                     fqName,
-                    descriptors.isEmpty()
+                    descriptors.isEmpty(),
+                    serializerExtension.stringTable
                 )
             }
         }
 
         if (result.isEmpty()) {
             result += withNewContext {
-                buildFragment(
+                buildKlibPackageFragment(
                     emptyPackageProto(),
                     emptyList(),
                     fqName,
-                    true
+                    true,
+                    serializerExtension.stringTable
                 )
             }
         }
@@ -297,3 +278,29 @@ fun DeclarationDescriptor.extractFileId(): Int? = when (this) {
 internal val ModuleDescriptor.packageFragmentProviderForModuleContentWithoutDependencies: PackageFragmentProvider
     get() = (this as? ModuleDescriptorImpl)?.packageFragmentProviderForModuleContentWithoutDependencies
         ?: error("Can't get a module content package fragments, it's not a ${ModuleDescriptorImpl::class.simpleName}.")
+
+fun buildKlibPackageFragment(
+    packageProto: ProtoBuf.Package,
+    classesProto: List<Pair<ProtoBuf.Class, Int>>,
+    fqName: FqName,
+    isEmpty: Boolean,
+    stringTable: SerializableStringTable,
+): ProtoBuf.PackageFragment {
+
+    val (stringTableProto, nameTableProto) = stringTable.buildProto()
+
+    return ProtoBuf.PackageFragment.newBuilder()
+        .setPackage(packageProto)
+        .addAllClass_(classesProto.map { it.first })
+        .setStrings(stringTableProto)
+        .setQualifiedNames(nameTableProto)
+        .also { packageFragment ->
+            classesProto.forEach {
+                packageFragment.addExtension(KlibMetadataProtoBuf.className, it.second )
+            }
+            packageFragment.setExtension(KlibMetadataProtoBuf.isEmpty, isEmpty)
+            packageFragment.setExtension(KlibMetadataProtoBuf.fqName, fqName.asString())
+        }
+        .build()
+}
+

compiler/ir/serialization.js/build.gradle.kts

@@ -9,6 +9,7 @@ dependencies {
     api(project(":js:js.frontend"))
 
     implementation(project(":compiler:ir.backend.common"))
+    implementation(project(":compiler:fir:fir-serialization"))
 
     compileOnly(intellijCore())
 }

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/firKlibSerialization.kt

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2010-2022 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.ir.backend.js
+
+import org.jetbrains.kotlin.backend.common.serialization.metadata.buildKlibPackageFragment
+import org.jetbrains.kotlin.config.CommonConfigurationKeys
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.fir.FirSession
+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.FirElementAwareStringTable
+import org.jetbrains.kotlin.fir.serialization.FirElementSerializer
+import org.jetbrains.kotlin.fir.serialization.FirSerializerExtension
+import org.jetbrains.kotlin.fir.serialization.TypeApproximatorForMetadataSerializer
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
+import org.jetbrains.kotlin.serialization.SerializableStringTable
+
+fun serializeSingleFirFile(file: FirFile, session: FirSession, scopeSession: ScopeSession, configuration: CompilerConfiguration): ProtoBuf.PackageFragment {
+    val serializerExtension = FirKLibSerializerExtension(
+        session,
+        configuration.get(CommonConfigurationKeys.METADATA_VERSION)!!,
+        FirElementAwareSerializableStringTable()
+    )
+    val approximator = TypeApproximatorForMetadataSerializer(session)
+    val packageSerializer = FirElementSerializer.createTopLevel(session, scopeSession, serializerExtension, approximator)
+
+    // TODO: typealiases (see klib serializer)
+    // TODO: split package fragment (see klib serializer)
+    // TODO: handle incremental/monolothic (see klib serializer) - maybe externally
+
+    val packageProto = packageSerializer.packagePartProto(file.packageFqName, file).build()
+
+    // TODO: filter out expects
+    val classifiers = file.declarations.filterIsInstance<FirClass>().sortedBy { it.classId.asFqNameString() }
+
+    val classesProto = classifiers.map {
+        val classSerializer = FirElementSerializer.create(session, scopeSession, it, serializerExtension, null, approximator)
+        val index = classSerializer.stringTable.getFqNameIndex(it)
+        classSerializer.classProto(it).build() to index
+    }
+
+    val hasTopLevelDeclarations = file.declarations.any {
+        it is FirProperty || it is FirSimpleFunction || it is FirTypeAlias
+    }
+
+    return buildKlibPackageFragment(
+        packageProto,
+        classesProto,
+        file.packageFqName,
+        hasTopLevelDeclarations && classesProto.isEmpty(),
+        serializerExtension.stringTable as SerializableStringTable
+    )
+}
+
+class FirKLibSerializerExtension(
+    override val session: FirSession,
+    override val metadataVersion: BinaryVersion,
+    override val stringTable: FirElementAwareSerializableStringTable
+) : FirSerializerExtension()
+
+class FirElementAwareSerializableStringTable() : FirElementAwareStringTable, SerializableStringTable()

compiler/serialization/src/org/jetbrains/kotlin/serialization/StringTableImpl.kt

@@ -17,13 +17,14 @@
 package org.jetbrains.kotlin.serialization
 
 import org.jetbrains.kotlin.metadata.ProtoBuf
-import org.jetbrains.kotlin.metadata.ProtoBuf.QualifiedNameTable.QualifiedName
 import org.jetbrains.kotlin.metadata.serialization.Interner
+import org.jetbrains.kotlin.metadata.serialization.StringTable
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 
-open class StringTableImpl : DescriptorAwareStringTable {
+open class SerializableStringTable : StringTable {
-    private class FqNameProto(val fqName: QualifiedName.Builder) {
+
+    private class FqNameProto(val fqName: ProtoBuf.QualifiedNameTable.QualifiedName.Builder) {
         override fun hashCode(): Int {
             var result = 13
             result = 31 * result + fqName.parentQualifiedName
@@ -48,14 +49,14 @@ open class StringTableImpl : DescriptorAwareStringTable {
     override fun getStringIndex(string: String): Int = strings.intern(string)
 
     override fun getQualifiedClassNameIndex(className: String, isLocal: Boolean): Int =
-        getQualifiedClassNameIndex(ClassId.fromString(className, isLocal))
+        getClassIdIndex(ClassId.fromString(className, isLocal))
 
-    override fun getQualifiedClassNameIndex(classId: ClassId): Int {
+    fun getClassIdIndex(classId: ClassId): Int {
-        val builder = QualifiedName.newBuilder()
+        val builder = ProtoBuf.QualifiedNameTable.QualifiedName.newBuilder()
-        builder.kind = QualifiedName.Kind.CLASS
+        builder.kind = ProtoBuf.QualifiedNameTable.QualifiedName.Kind.CLASS
 
         builder.parentQualifiedName =
-            classId.outerClassId?.let(this::getQualifiedClassNameIndex)
+            classId.outerClassId?.let(this::getClassIdIndex)
                 ?: getPackageFqNameIndex(classId.packageFqName)
 
         builder.shortName = getStringIndex(classId.shortClassName.asString())
@@ -66,7 +67,7 @@ open class StringTableImpl : DescriptorAwareStringTable {
     fun getPackageFqNameIndex(fqName: FqName): Int {
         var result = -1
         for (segment in fqName.pathSegments()) {
-            val builder = QualifiedName.newBuilder()
+            val builder = ProtoBuf.QualifiedNameTable.QualifiedName.newBuilder()
             builder.shortName = getStringIndex(segment.asString())
             if (result != -1) {
                 builder.parentQualifiedName = result
@@ -89,6 +90,11 @@ open class StringTableImpl : DescriptorAwareStringTable {
 
         return Pair(strings.build(), qualifiedNames.build())
     }
+}
+
+open class StringTableImpl : DescriptorAwareStringTable, SerializableStringTable() {
+
+    override fun getQualifiedClassNameIndex(classId: ClassId): Int = getClassIdIndex(classId)
 
     override val isLocalClassIdReplacementKeptGeneric: Boolean
         get() = false