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

2023-05-08 16:39:47 +02:00
parent bb72a8bcb7
commit aef8ece123
2 changed files with 138 additions and 115 deletions
@@ -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) {
-    DECLARED_SIGNATURE(0),
-    COMMON_SIGNATURE(1),
-    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
+internal class IdSignatureSerialization(private val library: KotlinLibraryHeader) {
+    private enum class IdSignatureProtoType(val id: Int) {
+        DECLARED_SIGNATURE(0),
+        COMMON_SIGNATURE(1),
+        COMPOSITE_SIGNATURE(2),
+        ACCESSOR_SIGNATURE(3);
    }

-    when (signature) {
-        is IdSignature.CommonSignature -> {
-            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")
-        }
+    interface IdSignatureICSerializer {
+        fun serializeIdSignature(out: CodedOutputStream, signature: IdSignature)
    }
-}

-internal fun CodedInputStream.readIdSignature(indexToSignatureMapper: (Int) -> IdSignature): IdSignature {
-    when (val signatureType = readInt32()) {
-        IdSignatureProtoType.DECLARED_SIGNATURE.id -> {
-            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")
-        }
+    interface IdSignatureICDeserializer {
+        fun deserializeIdSignature(input: CodedInputStream): IdSignature
+        fun skipIdSignature(input: CodedInputStream)
    }
-}

-internal fun CodedInputStream.skipIdSignature() {
-    when (val signatureType = readInt32()) {
-        IdSignatureProtoType.DECLARED_SIGNATURE.id -> {
-            readInt32()
+    inner class FileIdSignatureSerialization(srcFile: KotlinSourceFile) : IdSignatureICSerializer, IdSignatureICDeserializer {
+        private val deserializer by lazy {
+            library.sourceFileDeserializers[srcFile] ?: notFoundIcError("signature deserializer", library.libraryFile, srcFile)
        }
-        IdSignatureProtoType.COMMON_SIGNATURE.id -> {
-            readString()
-            readString()
-            if (readBool()) {
-                readFixed64()
+
+        internal val signatureToIndexMapping = hashMapOf<IdSignature, Int>()
+
+        override fun serializeIdSignature(out: CodedOutputStream, signature: IdSignature) {
+            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()
-            skipIdSignature()
+
+        override fun deserializeIdSignature(input: CodedInputStream): IdSignature {
+            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()
-            skipIdSignature()
-        }
-        else -> {
-            icError("can not skip signature type $signatureType")
+
+        override fun skipIdSignature(input: CodedInputStream) {
+            when (val signatureType = input.readInt32()) {
+                IdSignatureProtoType.DECLARED_SIGNATURE.id -> {
+                    input.readInt32()
+                }
+                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) }
+    }
 }
@@ -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: 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
-            }
+            val deserializer = idSignatureSerialization.getIdSignatureDeserializer(srcFile)

            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()
-        fun CodedOutputStream.serializeIdSignature(signature: IdSignature) =
-            writeIdSignature(signature) { signatureToIndexMapping[it] ?: signatureToIndexMappingSaved[it] }
+        val serializer = idSignatureSerialization.getIdSignatureSerializer(srcFile, signatureToIndexMapping)

        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)
            }
        }