Eliminated global state in IR deserializer.

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

@@ -21,7 +21,6 @@ import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.*
-import org.jetbrains.kotlin.ir.util.IrDeserializer
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.backend.common.serialization.KotlinIr.IrStatement.*
@@ -44,7 +43,7 @@ abstract class IrModuleDeserializer(
     val logger: LoggingContext,
     val builtIns: IrBuiltIns,
     val symbolTable: SymbolTable
-) : IrDeserializer {
+) {
 
     abstract fun deserializeIrSymbol(proto: KotlinIr.IrSymbol): IrSymbol
     abstract fun deserializeIrType(proto: KotlinIr.IrTypeIndex): IrType
@@ -1102,10 +1101,9 @@ abstract class IrModuleDeserializer(
     private val allKnownOrigins =
         IrDeclarationOrigin::class.nestedClasses.toList() + DeclarationFactory.FIELD_FOR_OUTER_THIS::class
     val originIndex = allKnownOrigins.map { it.objectInstance as IrDeclarationOriginImpl }.associateBy { it.name }
-    val irrelevantOrigin = object : IrDeclarationOriginImpl("irrelevant") {}
     fun deserializeIrDeclarationOrigin(proto: KotlinIr.IrDeclarationOrigin) = originIndex[deserializeString(proto.custom)]!!
 
-    protected fun deserializeDeclaration(proto: KotlinIr.IrDeclaration, parent: IrDeclarationParent?): IrDeclaration {
+    public open fun deserializeDeclaration(proto: KotlinIr.IrDeclaration, parent: IrDeclarationParent?): IrDeclaration {
 
         val start = proto.coordinates.startOffset
         val end = proto.coordinates.endOffset
@@ -1168,3 +1166,5 @@ abstract class IrModuleDeserializer(
         return declaration
     }
 }
+
+val irrelevantOrigin = object : IrDeclarationOriginImpl("irrelevant") {}

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

@@ -11,59 +11,57 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.EmptyPackageFragmentDescriptor
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.impl.*
+import org.jetbrains.kotlin.ir.declarations.impl.IrClassImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrModuleFragmentImpl
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.expressions.impl.IrLoopBase
-import org.jetbrains.kotlin.ir.symbols.*
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
-import org.jetbrains.kotlin.ir.symbols.impl.*
+import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
+import org.jetbrains.kotlin.ir.symbols.IrSymbol
+import org.jetbrains.kotlin.ir.symbols.impl.IrAnonymousInitializerSymbolImpl
+import org.jetbrains.kotlin.ir.symbols.impl.IrFileSymbolImpl
+import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
+import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
+import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.IrDeserializer
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite.*
+import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite.newInstance
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedCallableMemberDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
-import java.io.File
 
 abstract class KotlinIrLinker(
-    logger: LoggingContext,
+    val logger: LoggingContext,
-    builtIns: IrBuiltIns,
+    val builtIns: IrBuiltIns,
-    symbolTable: SymbolTable,
+    val symbolTable: SymbolTable,
     private val forwardModuleDescriptor: ModuleDescriptor?,
-    val firstKnownBuiltinsIndex: Long)
+    private val firstKnownBuiltinsIndex: Long
-        : IrModuleDeserializer(logger, builtIns, symbolTable),
+) : DescriptorUniqIdAware, IrDeserializer {
-          DescriptorUniqIdAware {
 
     protected val deserializedSymbols = mutableMapOf<UniqIdKey, IrSymbol>()
     private val reachableTopLevels = mutableSetOf<UniqIdKey>()
     private val deserializedTopLevels = mutableSetOf<UniqIdKey>()
     private val forwardDeclarations = mutableSetOf<IrSymbol>()
 
-    private var deserializedModuleDescriptor: ModuleDescriptor? = null
+    private val deserializersForModules = mutableMapOf<ModuleDescriptor, IrDeserializerForModule>()
-    private var deserializedModuleProtoSymbolTables = mutableMapOf<ModuleDescriptor, KotlinIr.IrSymbolTable>()
-    private var deserializedModuleProtoStringTables = mutableMapOf<ModuleDescriptor, KotlinIr.StringTable>()
-    private var deserializedModuleProtoTypeTables = mutableMapOf<ModuleDescriptor, KotlinIr.IrTypeTable>()
-    private var deserializedModuleLoops = mutableMapOf<Pair<ModuleDescriptor, Int>, IrLoopBase>()
 
     val resolvedForwardDeclarations = mutableMapOf<UniqIdKey, UniqIdKey>()
 
-    abstract protected val descriptorReferenceDeserializer: DescriptorReferenceDeserializer
+    inner class IrDeserializerForModule(
+        private val moduleDescriptor: ModuleDescriptor,
+        private val moduleProto: KotlinIr.IrModule
+    ) : IrModuleDeserializer(logger, builtIns, symbolTable) {
 
-    protected val indexAfterKnownBuiltins = loadKnownBuiltinSymbols()
+        private var moduleLoops = mutableMapOf<Int, IrLoopBase>()
 
-    fun loadKnownBuiltinSymbols(): Long {
+        private fun referenceDeserializedSymbol(
-        var currentIndex = firstKnownBuiltinsIndex
+            proto: KotlinIr.IrSymbolData,
-        builtIns.knownBuiltins.forEach {
+            descriptor: DeclarationDescriptor?
-            require(it is IrFunction)
+        ): IrSymbol = when (proto.kind) {
-            deserializedSymbols.put(UniqIdKey(null, UniqId(currentIndex, isLocal = false)), it.symbol)
-            assert(symbolTable.referenceSimpleFunction(it.descriptor) == it.symbol)
-            currentIndex++
-        }
-        return currentIndex
-    }
-
-    private fun referenceDeserializedSymbol(proto: KotlinIr.IrSymbolData, descriptor: DeclarationDescriptor?): IrSymbol = when (proto.kind) {
             KotlinIr.IrSymbolKind.ANONYMOUS_INIT_SYMBOL ->
                 IrAnonymousInitializerSymbolImpl(
                     descriptor as ClassDescriptor?
@@ -120,26 +118,24 @@ abstract class KotlinIrLinker(
         }
 
         override fun deserializeIrSymbol(proto: KotlinIr.IrSymbol): IrSymbol {
-        val symbolData =
+            val symbolData = moduleProto.symbolTable.getSymbols(proto.index)
-            deserializedModuleProtoSymbolTables[deserializedModuleDescriptor]!!.getSymbols(proto.index)
             return deserializeIrSymbolData(symbolData)
         }
 
         override fun deserializeIrType(proto: KotlinIr.IrTypeIndex): IrType {
-        val typeData =
+            val typeData = moduleProto.typeTable.getTypes(proto.index)
-            deserializedModuleProtoTypeTables[deserializedModuleDescriptor]!!.getTypes(proto.index)
             return deserializeIrTypeData(typeData)
         }
 
-    override fun deserializeString(proto: KotlinIr.String) =
+        override fun deserializeString(proto: KotlinIr.String): String =
-        deserializedModuleProtoStringTables[deserializedModuleDescriptor]!!.getStrings(proto.index)
+            moduleProto.stringTable.getStrings(proto.index)
 
         override fun deserializeLoopHeader(loopIndex: Int, loopBuilder: () -> IrLoopBase) =
-            deserializedModuleLoops.getOrPut(deserializedModuleDescriptor!! to loopIndex, loopBuilder)
+            moduleLoops.getOrPut(loopIndex, loopBuilder)
 
-    fun deserializeIrSymbolData(proto: KotlinIr.IrSymbolData): IrSymbol {
+        private fun deserializeIrSymbolData(proto: KotlinIr.IrSymbolData): IrSymbol {
-        val key = proto.uniqId.uniqIdKey(deserializedModuleDescriptor!!)
+            val key = proto.uniqId.uniqIdKey(moduleDescriptor)
-        val topLevelKey = proto.topLevelUniqId.uniqIdKey(deserializedModuleDescriptor!!)
+            val topLevelKey = proto.topLevelUniqId.uniqIdKey(moduleDescriptor)
 
             if (!deserializedTopLevels.contains(topLevelKey)) reachableTopLevels.add(topLevelKey)
 
@@ -181,6 +177,28 @@ abstract class KotlinIrLinker(
                 isTypeParameter = proto.isTypeParameter
             )
 
+
+        override fun getPrimitiveTypeOrNull(symbol: IrClassifierSymbol, hasQuestionMark: Boolean) =
+            this@KotlinIrLinker.getPrimitiveTypeOrNull(symbol, hasQuestionMark)
+    }
+
+    protected open fun getPrimitiveTypeOrNull(symbol: IrClassifierSymbol, hasQuestionMark: Boolean): IrSimpleType? = null
+
+    protected abstract val descriptorReferenceDeserializer: DescriptorReferenceDeserializer
+
+    protected val indexAfterKnownBuiltins = loadKnownBuiltinSymbols()
+
+    private fun loadKnownBuiltinSymbols(): Long {
+        var currentIndex = firstKnownBuiltinsIndex
+        builtIns.knownBuiltins.forEach {
+            require(it is IrFunction)
+            deserializedSymbols[UniqIdKey(null, UniqId(currentIndex, isLocal = false))] = it.symbol
+            assert(symbolTable.referenceSimpleFunction(it.descriptor) == it.symbol)
+            currentIndex++
+        }
+        return currentIndex
+    }
+
     private val ByteArray.codedInputStream: org.jetbrains.kotlin.protobuf.CodedInputStream
         get() {
             val codedInputStream = org.jetbrains.kotlin.protobuf.CodedInputStream.newInstance(this)
@@ -190,13 +208,18 @@ abstract class KotlinIrLinker(
 
     private val reversedFileIndex = mutableMapOf<UniqIdKey, IrFile>()
 
-    private val UniqIdKey.moduleOfOrigin get() =
+    private val UniqIdKey.moduleOfOrigin
+        get() =
             this.moduleDescriptor ?: reversedFileIndex[this]?.packageFragmentDescriptor?.containingDeclaration
 
     private fun deserializeTopLevelDeclaration(uniqIdKey: UniqIdKey): IrDeclaration {
         val proto = loadTopLevelDeclarationProto(uniqIdKey)
-        return deserializeDeclaration(proto, reversedFileIndex[uniqIdKey]!!)
+        return deserializersForModules[uniqIdKey.moduleOfOrigin]!!.deserializeDeclaration(
+            proto,
+            reversedFileIndex[uniqIdKey]!!
+        )
     }
+
     protected abstract fun reader(moduleDescriptor: ModuleDescriptor, uniqId: UniqId): ByteArray
 
     private fun loadTopLevelDeclarationProto(uniqIdKey: UniqIdKey): KotlinIr.IrDeclaration {
@@ -230,14 +253,14 @@ abstract class KotlinIrLinker(
             if (deserializedSymbols[key]?.isBound == true ||
                 // The key.moduleOrigin is null for uniqIds that we haven't seen in any of the library headers.
                 // Just skip it for now and handle it elsewhere.
-                key.moduleOfOrigin == null) {
+                key.moduleOfOrigin == null
+            ) {
 
                 reachableTopLevels.remove(key)
                 deserializedTopLevels.add(key)
                 continue
             }
 
-            deserializedModuleDescriptor = key.moduleOfOrigin
             val reachable = deserializeTopLevelDeclaration(key)
             val file = reversedFileIndex[key]!!
             file.declarations.add(reachable)
@@ -253,8 +276,7 @@ abstract class KotlinIrLinker(
     override fun findDeserializedDeclaration(symbol: IrSymbol): IrDeclaration? {
 
         if (!symbol.isBound) {
-            val topLevelDesecriptor = findDeserializedDeclarationForDescriptor(symbol.descriptor)
+            findDeserializedDeclarationForDescriptor(symbol.descriptor) ?: return null
-            if (topLevelDesecriptor == null) return null
         }
 
         assert(symbol.isBound) {
@@ -291,7 +313,8 @@ abstract class KotlinIrLinker(
             val declarations = symbols.map {
 
                 val classDescriptor = it.descriptor as ClassDescriptor
-                val declaration = symbolTable.declareClass(UNDEFINED_OFFSET, UNDEFINED_OFFSET, irrelevantOrigin,
+                val declaration = symbolTable.declareClass(
+                    UNDEFINED_OFFSET, UNDEFINED_OFFSET, irrelevantOrigin,
                     classDescriptor,
                     classDescriptor.modality
                 ) { symbol: IrClassSymbol -> IrClassImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, irrelevantOrigin, symbol) }
@@ -305,14 +328,22 @@ abstract class KotlinIrLinker(
         }
     }
 
-    fun deserializeIrFile(fileProto: KotlinIr.IrFile, moduleDescriptor: ModuleDescriptor, deseralizationStrategy: DeserializationStrategy): IrFile {
+    fun deserializeIrFile(
+        fileProto: KotlinIr.IrFile,
+        moduleDescriptor: ModuleDescriptor,
+        deseralizationStrategy: DeserializationStrategy
+    ): IrFile {
+
+        val moduleDeserializer = deserializersForModules[moduleDescriptor]!!
+
         val fileEntry = NaiveSourceBasedFileEntryImpl(
-            deserializeString(fileProto.fileEntry.name),
+            moduleDeserializer.deserializeString(fileProto.fileEntry.name),
             fileProto.fileEntry.lineStartOffsetsList.toIntArray()
         )
 
         // TODO: we need to store "" in protobuf, I suppose. Or better yet, reuse fqname storage from metadata.
-        val fqName = deserializeString(fileProto.fqName).let { if (it == "<root>") FqName.ROOT else FqName(it) }
+        val fqName = moduleDeserializer.deserializeString(fileProto.fqName)
+            .let { if (it == "<root>") FqName.ROOT else FqName(it) }
 
         val packageFragmentDescriptor = EmptyPackageFragmentDescriptor(moduleDescriptor, fqName)
 
@@ -328,22 +359,23 @@ abstract class KotlinIrLinker(
             }
         }
 
-        val annotations = deserializeAnnotations(fileProto.annotations)
+        val annotations = moduleDeserializer.deserializeAnnotations(fileProto.annotations)
         file.annotations.addAll(annotations)
 
 
         if (deseralizationStrategy == DeserializationStrategy.EXPLICITLY_EXPORTED)
-            fileProto.explicitlyExportedToCompilerList.forEach { deserializeIrSymbol(it) }
+            fileProto.explicitlyExportedToCompilerList.forEach { moduleDeserializer.deserializeIrSymbol(it) }
 
         return file
     }
 
-    fun deserializeIrModuleHeader(proto: KotlinIr.IrModule, moduleDescriptor: ModuleDescriptor, deserializationStrategy: DeserializationStrategy): IrModuleFragment {
+    fun deserializeIrModuleHeader(
+        proto: KotlinIr.IrModule,
+        moduleDescriptor: ModuleDescriptor,
+        deserializationStrategy: DeserializationStrategy
+    ): IrModuleFragment {
 
-        deserializedModuleDescriptor = moduleDescriptor
+        deserializersForModules[moduleDescriptor] = IrDeserializerForModule(moduleDescriptor, proto)
-        deserializedModuleProtoSymbolTables.put(moduleDescriptor, proto.symbolTable)
-        deserializedModuleProtoStringTables.put(moduleDescriptor, proto.stringTable)
-        deserializedModuleProtoTypeTables.put(moduleDescriptor, proto.typeTable)
 
         val files = proto.fileList.map {
             deserializeIrFile(it, moduleDescriptor, deserializationStrategy)
@@ -353,7 +385,11 @@ abstract class KotlinIrLinker(
         return module
     }
 
-    fun deserializeIrModuleHeader(moduleDescriptor: ModuleDescriptor, byteArray: ByteArray, deserializationStrategy: DeserializationStrategy = DeserializationStrategy.ONLY_REFERENCED): IrModuleFragment {
+    fun deserializeIrModuleHeader(
+        moduleDescriptor: ModuleDescriptor,
+        byteArray: ByteArray,
+        deserializationStrategy: DeserializationStrategy = DeserializationStrategy.ONLY_REFERENCED
+    ): IrModuleFragment {
         val proto = KotlinIr.IrModule.parseFrom(byteArray.codedInputStream, newInstance())
         return deserializeIrModuleHeader(proto, moduleDescriptor, deserializationStrategy)
     }