Commonized FileRegistry with SourceFileMap for klib metadata

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

@@ -12,8 +12,13 @@ import org.jetbrains.kotlin.descriptors.MemberDescriptor
 import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
 import org.jetbrains.kotlin.resolve.checkers.ExpectedActualDeclarationChecker
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.library.KotlinLibrary
+import org.jetbrains.kotlin.library.metadata.KlibMetadataPackageFragment
+import org.jetbrains.kotlin.library.metadata.KlibMetadataProtoBuf
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.OverridingUtil
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPropertyDescriptor
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedSimpleFunctionDescriptor
 
 internal val DeclarationDescriptor.isExpectMember: Boolean
     get() = this is MemberDescriptor && this.isExpect
@@ -49,3 +54,23 @@ internal fun <T : CallableMemberDescriptor> T.resolveFakeOverrideMaybeAbstract()
 // This is Native specific. Try to eliminate.
 val ModuleDescriptor.isForwardDeclarationModule get() =
     name == Name.special("<forward declarations>")
+
+private fun sourceByIndex(descriptor: CallableMemberDescriptor, index: Int): SourceFile {
+    val fragment = descriptor.findPackage() as KlibMetadataPackageFragment
+    return fragment.fileRegistry.sourceFile(index)
+}
+
+fun CallableMemberDescriptor.findSourceFile(): SourceFile {
+    val source = this.source.containingFile
+    if (source != SourceFile.NO_SOURCE_FILE)
+        return source
+    return when {
+        this is DeserializedSimpleFunctionDescriptor && proto.hasExtension(KlibMetadataProtoBuf.functionFile) ->
+            sourceByIndex(
+                this, proto.getExtension(KlibMetadataProtoBuf.functionFile))
+        this is DeserializedPropertyDescriptor && proto.hasExtension(KlibMetadataProtoBuf.propertyFile) ->
+            sourceByIndex(
+                this, proto.getExtension(KlibMetadataProtoBuf.propertyFile))
+        else -> TODO()
+    }
+}

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

@@ -3,7 +3,7 @@
  * 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.lower.serialization.metadata
+package org.jetbrains.kotlin.backend.common.serialization.metadata
 
 import org.jetbrains.kotlin.library.metadata.KlibMetadataPackageFragment
 import org.jetbrains.kotlin.psi.KtFile

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

@@ -1,41 +0,0 @@
-/*
- * Copyright 2010-2019 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.lower.serialization.metadata
-
-import org.jetbrains.kotlin.backend.common.serialization.metadata.extractFileId
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithSource
-import org.jetbrains.kotlin.library.metadata.KlibMetadataPackageFragment
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.DescriptorUtils
-import org.jetbrains.kotlin.resolve.source.PsiSourceFile
-
-
-class KlibMetadataFileRegistry {
-    private val fileIdsImpl = mutableMapOf<KlibFileMetadata, Int>()
-
-    fun lookup(file: KlibFileMetadata) = fileIdsImpl.getOrPut(file) { fileIdsImpl.size }
-
-    val fileIds: Map<KlibFileMetadata, Int>
-        get() = fileIdsImpl
-
-    fun getFileId(descriptor: DeclarationDescriptor): Int? {
-        if (!DescriptorUtils.isTopLevelDeclaration(descriptor) || descriptor !is DeclarationDescriptorWithSource) return null
-
-        val fileId = descriptor.extractFileId()
-        if (fileId != null) {
-            (descriptor.containingDeclaration as? KlibMetadataPackageFragment)?.let { packageFragment ->
-                return this.lookup(KotlinDeserializedFileMetadata(packageFragment, fileId))
-            }
-        }
-
-        val file = descriptor.source.containingFile as? PsiSourceFile ?: return null
-
-        val psiFile = file.psiFile
-        return (psiFile as? KtFile)?.let { this.lookup(KotlinPsiFileMetadata(it)) }
-    }
-
-}

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

@@ -11,8 +11,7 @@ import org.jetbrains.kotlin.config.AnalysisFlags
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
-import org.jetbrains.kotlin.ir.backend.js.lower.serialization.metadata.KlibMetadataFileRegistry
-import org.jetbrains.kotlin.ir.backend.js.lower.serialization.metadata.KotlinPsiFileMetadata
+import org.jetbrains.kotlin.library.metadata.KlibMetadataFileRegistry
 import org.jetbrains.kotlin.library.metadata.KlibMetadataProtoBuf
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
@@ -35,7 +34,6 @@ abstract class KlibMetadataSerializer(
     val metadataVersion: BinaryVersion,
     val descriptorTable: DescriptorTable
 ) {
-
     val fileRegistry = KlibMetadataFileRegistry()
 
     lateinit var serializerContext: SerializerContext
@@ -57,7 +55,7 @@ abstract class KlibMetadataSerializer(
             languageVersionSettings,
             metadataVersion,
             ::declarationTableHandler,
-            {descriptor -> fileRegistry.getFileId(descriptor) } ,
+            {descriptor: DeclarationDescriptorWithSource -> fileRegistry.assign(descriptor.source.containingFile) } ,
             KlibMetadataStringTable()
         )
         return SerializerContext(
@@ -82,10 +80,6 @@ abstract class KlibMetadataSerializer(
 
         val (stringTableProto, nameTableProto) = serializerExtension.stringTable.buildProto()
 
-        // TODO: we place files table to each and every fragment.
-        // Need to refactor it out sonehow.
-        val files = serializeFiles(fileRegistry, bindingContext, AnnotationSerializer(serializerExtension.stringTable))
-
         return ProtoBuf.PackageFragment.newBuilder()
             .setPackage(packageProto)
             .addAllClass_(classesProto.map { it.first })
@@ -95,7 +89,6 @@ abstract class KlibMetadataSerializer(
                 classesProto.forEach {
                     packageFragment.addExtension(KlibMetadataProtoBuf.className, it.second )
                 }
-                packageFragment.setExtension(KlibMetadataProtoBuf.packageFragmentFiles, files)
                 packageFragment.setExtension(KlibMetadataProtoBuf.isEmpty, isEmpty)
                 packageFragment.setExtension(KlibMetadataProtoBuf.fqName, fqName.asString())
             }
@@ -224,32 +217,17 @@ abstract class KlibMetadataSerializer(
     }
 
     private fun serializeFiles(
-        fileRegistry: KlibMetadataFileRegistry,
-        bindingContext: BindingContext,
-        serializer: AnnotationSerializer
-    ): KlibMetadataProtoBuf.Files {
-        val filesProto = KlibMetadataProtoBuf.Files.newBuilder()
-        for ((file, id) in fileRegistry.fileIds.entries.sortedBy { it.value }) {
+        header: KlibMetadataProtoBuf.Header.Builder,
+        fileRegistry: KlibMetadataFileRegistry
+    ) {
+
+        fileRegistry.filesAndClear().map { it.name ?: "" }.forEach {
             val fileProto = KlibMetadataProtoBuf.File.newBuilder()
-            if (id != filesProto.fileCount) {
-                fileProto.id = id
-            }
-            val annotations = when (file) {
-                is KotlinPsiFileMetadata -> file.ktFile.annotationEntries.map { bindingContext[BindingContext.ANNOTATION, it]!! }
-                //is KotlinDeserializedFileMetadata -> file.packageFragment.fileMap[file.fileId]!!.annotations
-                else -> TODO("support other file types")
-            }
-            for (annotation in annotations.filterOutSourceAnnotations()) {
-                fileProto.addAnnotation(serializer.serializeAnnotation(annotation))
-            }
-            val name = when (file) {
-                is KotlinPsiFileMetadata -> file.ktFile.getName()
-                else -> TODO("support other file types")
-            }
-            fileProto.name = name
-            filesProto.addFile(fileProto)
+                .setName(it)
+                .build()
+
+            header.addFile(fileProto)
         }
-        return filesProto.build()
     }
 
     protected fun getPackagesFqNames(module: ModuleDescriptor): Set<FqName> {
@@ -296,6 +274,9 @@ abstract class KlibMetadataSerializer(
         emptyPackages.forEach {
             header.addEmptyPackage(it)
         }
+
+        serializeFiles(header, fileRegistry)
+
         return header.build()
     }
 

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

@@ -22,7 +22,7 @@ class KlibMetadataSerializerExtension(
     private val languageVersionSettings: LanguageVersionSettings,
     override val metadataVersion: BinaryVersion,
     val declarationTableHandler: (DeclarationDescriptor) -> KlibMetadataProtoBuf.DescriptorUniqId?,
-    val descriptorFileId: (DeclarationDescriptor) -> Int?,
+    val descriptorFileId: (DeclarationDescriptorWithSource) -> Int?,
     override val stringTable: StringTableImpl
 ) : KotlinSerializerExtensionBase(KlibMetadataSerializerProtocol) {
     override fun shouldUseTypeTable(): Boolean = true