JS: fix serialization of fileId in case of incremental compilation

js/js.serializer/src/org/jetbrains/kotlin/serialization/js/KotlinFileMetadata.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.serialization.js
+
+import org.jetbrains.kotlin.psi.KtFile
+
+sealed class KotlinFileMetadata
+
+data class KotlinPsiFileMetadata(val ktFile: KtFile) : KotlinFileMetadata()
+
+data class KotlinDeserializedFileMetadata(
+        val packageFragment: KotlinJavascriptPackageFragment,
+        val fileId: Int
+) : KotlinFileMetadata()

js/js.serializer/src/org/jetbrains/kotlin/serialization/js/KotlinFileRegistry.kt

@@ -16,13 +16,11 @@
 
 package org.jetbrains.kotlin.serialization.js
 
-import org.jetbrains.kotlin.psi.KtFile
-
 class KotlinFileRegistry {
-    private val fileIdsImpl = mutableMapOf<KtFile, Int>()
+    private val fileIdsImpl = mutableMapOf<KotlinFileMetadata, Int>()
 
-    fun lookup(file: KtFile) = fileIdsImpl.getOrPut(file) { fileIdsImpl.size }
+    fun lookup(file: KotlinFileMetadata) = fileIdsImpl.getOrPut(file) { fileIdsImpl.size }
 
-    val fileIds: Map<KtFile, Int>
+    val fileIds: Map<KotlinFileMetadata, Int>
         get() = fileIdsImpl
 }

js/js.serializer/src/org/jetbrains/kotlin/serialization/js/KotlinJavascriptPackageFragment.kt

@@ -29,10 +29,7 @@ import org.jetbrains.kotlin.serialization.deserialization.AnnotationDeserializer
 import org.jetbrains.kotlin.serialization.deserialization.DeserializationConfiguration
 import org.jetbrains.kotlin.serialization.deserialization.DeserializedPackageFragmentImpl
 import org.jetbrains.kotlin.serialization.deserialization.IncompatibleVersionErrorData
-import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
-import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPropertyDescriptor
-import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedSimpleFunctionDescriptor
 import org.jetbrains.kotlin.storage.StorageManager
 import org.jetbrains.kotlin.storage.getValue
 
@@ -44,7 +41,7 @@ class KotlinJavascriptPackageFragment(
         header: JsProtoBuf.Header,
         configuration: DeserializationConfiguration
 ) : DeserializedPackageFragmentImpl(fqName, storageManager, module, proto, JsContainerSource(fqName, header, configuration)) {
-    private val fileMap: Map<Int, FileHolder> by storageManager.createLazyValue {
+    val fileMap: Map<Int, FileHolder> by storageManager.createLazyValue {
         this.proto.getExtension(JsProtoBuf.packageFragmentFiles).fileList.withIndex().associate { (index, file) ->
             (if (file.hasId()) file.id else index) to FileHolder(file.annotationList)
         }
@@ -58,17 +55,12 @@ class KotlinJavascriptPackageFragment(
         if (DescriptorUtils.getParentOfType(descriptor, PackageFragmentDescriptor::class.java) != this) {
             throw IllegalArgumentException("Provided descriptor $descriptor does not belong to this package $this")
         }
-        val fileId = when (descriptor) {
-            is DeserializedClassDescriptor -> descriptor.classProto.getExtension(JsProtoBuf.classContainingFileId)
-            is DeserializedSimpleFunctionDescriptor -> descriptor.proto.getExtension(JsProtoBuf.functionContainingFileId)
-            is DeserializedPropertyDescriptor -> descriptor.proto.getExtension(JsProtoBuf.propertyContainingFileId)
-            else -> null
-        }
+        val fileId = descriptor.extractFileId()
 
         return fileId?.let { fileMap[it] }?.annotations.orEmpty()
     }
 
-    private inner class FileHolder(val annotationsProto: List<ProtoBuf.Annotation>) {
+    inner class FileHolder(private val annotationsProto: List<ProtoBuf.Annotation>) {
         val annotations: List<AnnotationDescriptor> by storageManager.createLazyValue {
             annotationsProto.map { annotationDeserializer.deserializeAnnotation(it, nameResolver) }
         }

js/js.serializer/src/org/jetbrains/kotlin/serialization/js/KotlinJavascriptSerializationUtil.kt

@@ -29,6 +29,9 @@ import org.jetbrains.kotlin.serialization.AnnotationSerializer
 import org.jetbrains.kotlin.serialization.DescriptorSerializer
 import org.jetbrains.kotlin.serialization.ProtoBuf
 import org.jetbrains.kotlin.serialization.deserialization.DeserializationConfiguration
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPropertyDescriptor
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedSimpleFunctionDescriptor
 import org.jetbrains.kotlin.storage.StorageManager
 import org.jetbrains.kotlin.utils.JsMetadataVersion
 import org.jetbrains.kotlin.utils.KotlinJavascriptMetadataUtils
@@ -130,7 +133,9 @@ object KotlinJavascriptSerializationUtil {
         val builder = ProtoBuf.PackageFragment.newBuilder()
 
         // TODO: ModuleDescriptor should be able to return the package only with the contents of that module, without dependencies
-        val skip: (DeclarationDescriptor) -> Boolean = { DescriptorUtils.getContainingModule(it) != module || (it is MemberDescriptor && it.isExpect) }
+        val skip: (DeclarationDescriptor) -> Boolean = {
+            DescriptorUtils.getContainingModule(it) != module || (it is MemberDescriptor && it.isExpect)
+        }
 
         val fileRegistry = KotlinFileRegistry()
         val serializerExtension = KotlinJavascriptSerializerExtension(fileRegistry)
@@ -183,8 +188,11 @@ object KotlinJavascriptSerializationUtil {
             if (id != filesProto.fileCount) {
                 fileProto.id = id
             }
-            for (annotationPsi in file.annotationEntries) {
-                val annotation = bindingContext[BindingContext.ANNOTATION, annotationPsi]!!
+            val annotations = when (file) {
+                is KotlinPsiFileMetadata -> file.ktFile.annotationEntries.map { bindingContext[BindingContext.ANNOTATION, it]!! }
+                is KotlinDeserializedFileMetadata -> file.packageFragment.fileMap[file.fileId]!!.annotations
+            }
+            for (annotation in annotations) {
                 fileProto.addAnnotation(serializer.serializeAnnotation(annotation))
             }
             filesProto.addFile(fileProto)
@@ -282,4 +290,11 @@ object KotlinJavascriptSerializationUtil {
     }
 }
 
-data class KotlinJavaScriptLibraryParts(val header: JsProtoBuf.Header, val body: List<ProtoBuf.PackageFragment>)
+data class KotlinJavaScriptLibraryParts(val header: JsProtoBuf.Header, val body: List<ProtoBuf.PackageFragment>)
+
+internal fun DeclarationDescriptor.extractFileId(): Int? = when (this) {
+    is DeserializedClassDescriptor -> classProto.getExtension(JsProtoBuf.classContainingFileId)
+    is DeserializedSimpleFunctionDescriptor -> proto.getExtension(JsProtoBuf.functionContainingFileId)
+    is DeserializedPropertyDescriptor -> proto.getExtension(JsProtoBuf.propertyContainingFileId)
+    else -> null
+}

js/js.serializer/src/org/jetbrains/kotlin/serialization/js/KotlinJavascriptSerializerExtension.kt

@@ -64,11 +64,17 @@ class KotlinJavascriptSerializerExtension(
     private fun getFileId(descriptor: DeclarationDescriptor): Int? {
         if (!DescriptorUtils.isTopLevelDeclaration(descriptor) || descriptor !is DeclarationDescriptorWithSource) return null
 
-        val file = descriptor.source.containingFile
-        if (file !is PsiSourceFile) return null
+        val fileId = descriptor.extractFileId()
+        if (fileId != null) {
+            (descriptor.containingDeclaration as? KotlinJavascriptPackageFragment)?.let { packageFragment ->
+                return fileRegistry.lookup(KotlinDeserializedFileMetadata(packageFragment, fileId))
+            }
+        }
+
+        val file = descriptor.source.containingFile as? PsiSourceFile ?: return null
 
         val psiFile = file.psiFile
-        return (psiFile as? KtFile)?.let { fileRegistry.lookup(it) }
+        return (psiFile as? KtFile)?.let { fileRegistry.lookup(KotlinPsiFileMetadata(it)) }
     }
 }