[JS IR] Refactor signature serializer in JS IR IC infra

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/ic/IdSignatureSerialization.kt

@@ -9,109 +9,143 @@ import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.protobuf.CodedInputStream
 import org.jetbrains.kotlin.protobuf.CodedOutputStream
 
-private enum class IdSignatureProtoType(val id: Int) {
+internal class IdSignatureSerialization(private val library: KotlinLibraryHeader) {
-    DECLARED_SIGNATURE(0),
+    private enum class IdSignatureProtoType(val id: Int) {
-    COMMON_SIGNATURE(1),
+        DECLARED_SIGNATURE(0),
-    COMPOSITE_SIGNATURE(2),
+        COMMON_SIGNATURE(1),
-    ACCESSOR_SIGNATURE(3);
+        COMPOSITE_SIGNATURE(2),
-}
+        ACCESSOR_SIGNATURE(3);
-
-internal fun CodedOutputStream.writeIdSignature(signature: IdSignature, signatureToIndexMapper: (IdSignature) -> Int?) {
-    val index = signatureToIndexMapper(signature)
-    if (index != null) {
-        writeInt32NoTag(IdSignatureProtoType.DECLARED_SIGNATURE.id)
-        writeInt32NoTag(index)
-        return
     }
 
-    when (signature) {
+    interface IdSignatureICSerializer {
-        is IdSignature.CommonSignature -> {
+        fun serializeIdSignature(out: CodedOutputStream, signature: IdSignature)
-            writeInt32NoTag(IdSignatureProtoType.COMMON_SIGNATURE.id)
-            writeStringNoTag(signature.packageFqName)
-            writeStringNoTag(signature.declarationFqName)
-            val id = signature.id
-            if (id != null) {
-                writeBoolNoTag(true)
-                writeFixed64NoTag(id)
-            } else {
-                writeBoolNoTag(false)
-            }
-            writeInt64NoTag(signature.mask)
-        }
-        is IdSignature.CompositeSignature -> {
-            writeInt32NoTag(IdSignatureProtoType.COMPOSITE_SIGNATURE.id)
-            writeIdSignature(signature.container, signatureToIndexMapper)
-            writeIdSignature(signature.inner, signatureToIndexMapper)
-        }
-        is IdSignature.AccessorSignature -> {
-            writeInt32NoTag(IdSignatureProtoType.ACCESSOR_SIGNATURE.id)
-            writeIdSignature(signature.propertySignature, signatureToIndexMapper)
-            writeIdSignature(signature.accessorSignature, signatureToIndexMapper)
-        }
-        else -> {
-            icError("can not write $signature signature")
-        }
     }
-}
 
-internal fun CodedInputStream.readIdSignature(indexToSignatureMapper: (Int) -> IdSignature): IdSignature {
+    interface IdSignatureICDeserializer {
-    when (val signatureType = readInt32()) {
+        fun deserializeIdSignature(input: CodedInputStream): IdSignature
-        IdSignatureProtoType.DECLARED_SIGNATURE.id -> {
+        fun skipIdSignature(input: CodedInputStream)
-            return indexToSignatureMapper(readInt32())
-        }
-        IdSignatureProtoType.COMMON_SIGNATURE.id -> {
-            val packageFqName = readString()
-            val declarationFqName = readString()
-            val id = if (readBool()) {
-                readFixed64()
-            } else {
-                null
-            }
-            val mask = readInt64()
-            return IdSignature.CommonSignature(packageFqName, declarationFqName, id, mask)
-        }
-        IdSignatureProtoType.COMPOSITE_SIGNATURE.id -> {
-            val containerSignature = readIdSignature(indexToSignatureMapper)
-            val innerSignature = readIdSignature(indexToSignatureMapper)
-            return IdSignature.CompositeSignature(containerSignature, innerSignature)
-        }
-        IdSignatureProtoType.ACCESSOR_SIGNATURE.id -> {
-            val propertySignature = readIdSignature(indexToSignatureMapper)
-            val accessorSignature = readIdSignature(indexToSignatureMapper)
-            if (accessorSignature !is IdSignature.CommonSignature) {
-                icError("can not read accessor signature")
-            }
-            return IdSignature.AccessorSignature(propertySignature, accessorSignature)
-        }
-        else -> {
-            icError("can not read signature type $signatureType")
-        }
     }
-}
 
-internal fun CodedInputStream.skipIdSignature() {
+    inner class FileIdSignatureSerialization(srcFile: KotlinSourceFile) : IdSignatureICSerializer, IdSignatureICDeserializer {
-    when (val signatureType = readInt32()) {
+        private val deserializer by lazy {
-        IdSignatureProtoType.DECLARED_SIGNATURE.id -> {
+            library.sourceFileDeserializers[srcFile] ?: notFoundIcError("signature deserializer", library.libraryFile, srcFile)
-            readInt32()
         }
-        IdSignatureProtoType.COMMON_SIGNATURE.id -> {
+
-            readString()
+        internal val signatureToIndexMapping = hashMapOf<IdSignature, Int>()
-            readString()
+
-            if (readBool()) {
+        override fun serializeIdSignature(out: CodedOutputStream, signature: IdSignature) {
-                readFixed64()
+            val index = signatureToIndexMapping[signature]
+            if (index != null) {
+                out.writeInt32NoTag(IdSignatureProtoType.DECLARED_SIGNATURE.id)
+                out.writeInt32NoTag(index)
+                return
+            }
+
+            when (signature) {
+                is IdSignature.CommonSignature -> {
+                    out.writeInt32NoTag(IdSignatureProtoType.COMMON_SIGNATURE.id)
+                    out.writeStringNoTag(signature.packageFqName)
+                    out.writeStringNoTag(signature.declarationFqName)
+                    val id = signature.id
+                    if (id != null) {
+                        out.writeBoolNoTag(true)
+                        out.writeFixed64NoTag(id)
+                    } else {
+                        out.writeBoolNoTag(false)
+                    }
+                    out.writeInt64NoTag(signature.mask)
+                }
+                is IdSignature.CompositeSignature -> {
+                    out.writeInt32NoTag(IdSignatureProtoType.COMPOSITE_SIGNATURE.id)
+                    serializeIdSignature(out, signature.container)
+                    serializeIdSignature(out, signature.inner)
+                }
+                is IdSignature.AccessorSignature -> {
+                    out.writeInt32NoTag(IdSignatureProtoType.ACCESSOR_SIGNATURE.id)
+                    serializeIdSignature(out, signature.propertySignature)
+                    serializeIdSignature(out, signature.accessorSignature)
+                }
+                else -> {
+                    icError("can not write $signature signature")
+                }
             }
-            readInt64()
         }
-        IdSignatureProtoType.COMPOSITE_SIGNATURE.id -> {
+
-            skipIdSignature()
+        override fun deserializeIdSignature(input: CodedInputStream): IdSignature {
-            skipIdSignature()
+            when (val signatureType = input.readInt32()) {
+                IdSignatureProtoType.DECLARED_SIGNATURE.id -> {
+                    val index = input.readInt32()
+                    val signature = deserializer.deserializeIdSignature(index)
+                    signatureToIndexMapping[signature] = index
+                    return signature
+                }
+                IdSignatureProtoType.COMMON_SIGNATURE.id -> {
+                    val packageFqName = input.readString()
+                    val declarationFqName = input.readString()
+                    val id = if (input.readBool()) {
+                        input.readFixed64()
+                    } else {
+                        null
+                    }
+                    val mask = input.readInt64()
+                    return IdSignature.CommonSignature(packageFqName, declarationFqName, id, mask)
+                }
+                IdSignatureProtoType.COMPOSITE_SIGNATURE.id -> {
+                    val containerSignature = deserializeIdSignature(input)
+                    val innerSignature = deserializeIdSignature(input)
+                    return IdSignature.CompositeSignature(containerSignature, innerSignature)
+                }
+                IdSignatureProtoType.ACCESSOR_SIGNATURE.id -> {
+                    val propertySignature = deserializeIdSignature(input)
+                    val accessorSignature = deserializeIdSignature(input)
+                    if (accessorSignature !is IdSignature.CommonSignature) {
+                        icError("can not read accessor signature")
+                    }
+                    return IdSignature.AccessorSignature(propertySignature, accessorSignature)
+                }
+                else -> {
+                    icError("can not read signature type $signatureType")
+                }
+            }
         }
-        IdSignatureProtoType.ACCESSOR_SIGNATURE.id -> {
+
-            skipIdSignature()
+        override fun skipIdSignature(input: CodedInputStream) {
-            skipIdSignature()
+            when (val signatureType = input.readInt32()) {
-        }
+                IdSignatureProtoType.DECLARED_SIGNATURE.id -> {
-        else -> {
+                    input.readInt32()
-            icError("can not skip signature type $signatureType")
+                }
+                IdSignatureProtoType.COMMON_SIGNATURE.id -> {
+                    input.readString()
+                    input.readString()
+                    if (input.readBool()) {
+                        input.readFixed64()
+                    }
+                    input.readInt64()
+                }
+                IdSignatureProtoType.COMPOSITE_SIGNATURE.id -> {
+                    skipIdSignature(input)
+                    skipIdSignature(input)
+                }
+                IdSignatureProtoType.ACCESSOR_SIGNATURE.id -> {
+                    skipIdSignature(input)
+                    skipIdSignature(input)
+                }
+                else -> {
+                    icError("can not skip signature type $signatureType")
+                }
+            }
         }
     }
+
+    private val fileSerializers = hashMapOf<KotlinSourceFile, FileIdSignatureSerialization>()
+
+    fun getIdSignatureDeserializer(srcFile: KotlinSourceFile): IdSignatureICDeserializer {
+        return fileSerializers.getOrPut(srcFile) { FileIdSignatureSerialization(srcFile) }
+    }
+
+    fun getIdSignatureSerializer(srcFile: KotlinSourceFile, signatureToIndexMapping: Map<IdSignature, Int>): IdSignatureICSerializer {
+        return fileSerializers.getOrPut(srcFile) {
+            FileIdSignatureSerialization(srcFile)
+        }.also { it.signatureToIndexMapping.putAll(signatureToIndexMapping) }
+    }
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/ic/IncrementalCache.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.ir.backend.js.ic
 
-import org.jetbrains.kotlin.backend.common.serialization.IdSignatureDeserializer
 import org.jetbrains.kotlin.backend.common.serialization.cityHash64
 import org.jetbrains.kotlin.backend.common.serialization.FingerprintHash
 import org.jetbrains.kotlin.ir.util.IdSignature
@@ -30,7 +29,7 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
 
     private val kotlinLibrarySourceFileMetadata = hashMapOf<KotlinSourceFile, KotlinSourceFileMetadata>()
 
-    private val signatureToIndexMappingFromMetadata = hashMapOf<KotlinSourceFile, MutableMap<IdSignature, Int>>()
+    private val idSignatureSerialization = IdSignatureSerialization(library)
 
     private val cacheHeaderFromDisk by lazy(LazyThreadSafetyMode.NONE) {
         cacheHeaderFile.useCodedInputIfExists {
@@ -148,16 +147,7 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
 
     private fun fetchSourceFileMetadata(srcFile: KotlinSourceFile, loadSignatures: Boolean) =
         kotlinLibrarySourceFileMetadata.getOrPut(srcFile) {
-            val signatureToIndexMapping = signatureToIndexMappingFromMetadata.getOrPut(srcFile) { hashMapOf() }
+            val deserializer = idSignatureSerialization.getIdSignatureDeserializer(srcFile)
-            val deserializer: IdSignatureDeserializer by lazy(LazyThreadSafetyMode.NONE) {
-                library.sourceFileDeserializers[srcFile] ?: notFoundIcError("signature deserializer", library.libraryFile, srcFile)
-            }
-
-            fun CodedInputStream.deserializeIdSignatureAndSave() = readIdSignature { index ->
-                val signature = deserializer.deserializeIdSignature(index)
-                signatureToIndexMapping[signature] = index
-                signature
-            }
 
             fun <T> CodedInputStream.readDependencies(signaturesReader: () -> T) = buildMapUntil(readInt32()) {
                 val libFile = KotlinLibraryFile.fromProtoStream(this@readDependencies)
@@ -171,12 +161,12 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
             fun CodedInputStream.readDirectDependencies() = readDependencies {
                 if (loadSignatures) {
                     buildMapUntil(readInt32()) {
-                        val signature = deserializeIdSignatureAndSave()
+                        val signature = deserializer.deserializeIdSignature(this@readDirectDependencies)
                         put(signature, ICHash.fromProtoStream(this@readDirectDependencies))
                     }
                 } else {
                     repeat(readInt32()) {
-                        skipIdSignature()
+                        deserializer.skipIdSignature(this@readDirectDependencies)
                         ICHash.fromProtoStream(this@readDirectDependencies)
                     }
                     emptyMap()
@@ -185,9 +175,9 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
 
             fun CodedInputStream.readInverseDependencies() = readDependencies {
                 if (loadSignatures) {
-                    buildSetUntil(readInt32()) { add(deserializeIdSignatureAndSave()) }
+                    buildSetUntil(readInt32()) { add(deserializer.deserializeIdSignature(this@readInverseDependencies)) }
                 } else {
-                    repeat(readInt32()) { skipIdSignature() }
+                    repeat(readInt32()) { deserializer.skipIdSignature(this@readInverseDependencies) }
                     emptySet()
                 }
             }
@@ -204,9 +194,10 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
         signatureToIndexMapping: Map<IdSignature, Int>
     ): SourceFileCacheArtifact {
         val binaryAstFile = srcFile.getCacheFile(BINARY_AST_SUFFIX)
-        val headerCacheFile = srcFile.getCacheFile(METADATA_SUFFIX)
         val sourceFileMetadata = kotlinLibrarySourceFileMetadata[srcFile]
             ?: return SourceFileCacheArtifact.DoNotChangeMetadata(srcFile, binaryAstFile)
+
+        val headerCacheFile = srcFile.getCacheFile(METADATA_SUFFIX)
         if (sourceFileMetadata.isEmpty()) {
             return SourceFileCacheArtifact.RemoveMetadata(srcFile, binaryAstFile, headerCacheFile)
         }
@@ -214,9 +205,7 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
             return SourceFileCacheArtifact.DoNotChangeMetadata(srcFile, binaryAstFile)
         }
 
-        val signatureToIndexMappingSaved = signatureToIndexMappingFromMetadata[srcFile] ?: emptyMap()
+        val serializer = idSignatureSerialization.getIdSignatureSerializer(srcFile, signatureToIndexMapping)
-        fun CodedOutputStream.serializeIdSignature(signature: IdSignature) =
-            writeIdSignature(signature) { signatureToIndexMapping[it] ?: signatureToIndexMappingSaved[it] }
 
         fun <T> CodedOutputStream.writeDependencies(depends: KotlinSourceFileMap<T>, signaturesWriter: (T) -> Unit) {
             writeInt32NoTag(depends.size)
@@ -233,7 +222,7 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
         fun CodedOutputStream.writeDirectDependencies(depends: KotlinSourceFileMap<Map<IdSignature, ICHash>>) = writeDependencies(depends) {
             writeInt32NoTag(it.size)
             for ((signature, hash) in it) {
-                serializeIdSignature(signature)
+                serializer.serializeIdSignature(this@writeDirectDependencies, signature)
                 hash.toProtoStream(this)
             }
         }
@@ -241,7 +230,7 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
         fun CodedOutputStream.writeInverseDependencies(depends: KotlinSourceFileMap<Set<IdSignature>>) = writeDependencies(depends) {
             writeInt32NoTag(it.size)
             for (signature in it) {
-                serializeIdSignature(signature)
+                serializer.serializeIdSignature(this@writeInverseDependencies, signature)
             }
         }