[IR SERIALIZATION] Refactored uniqId allocation

- split DeclarationTable into local and global one
 - minimize usages of kotlin mangler
 - clean up interfaces

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

@@ -7,9 +7,7 @@ package org.jetbrains.kotlin.backend.common.serialization
 
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.impl.IrAnonymousInitializerImpl
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
-import org.jetbrains.kotlin.ir.symbols.IrSymbol
 
 class DescriptorTable {
     private val descriptors = mutableMapOf<DeclarationDescriptor, Long>()
@@ -20,39 +18,52 @@ class DescriptorTable {
 }
 
 // TODO: We don't manage id clashes anyhow now.
-abstract class DeclarationTable(val builtIns: IrBuiltIns, val descriptorTable: DescriptorTable, mangler: KotlinMangler) : KotlinMangler by mangler {
-
-    private val builtInsTable = mutableMapOf<IrSymbol, UniqId>()
+abstract class GlobalDeclarationTable(private val mangler: KotlinMangler) {
     private val table = mutableMapOf<IrDeclaration, UniqId>()
-    val descriptors = descriptorTable
-    protected abstract var currentIndex: Long
 
-    open fun loadKnownBuiltins(): Long {
+    protected open fun loadKnownBuiltins(builtIns: IrBuiltIns, startIndex: Long): Long {
+        var index = startIndex
         builtIns.knownBuiltins.forEach {
-            builtInsTable[it] = UniqId(currentIndex++, false)
+            table[it.owner] = UniqId(index++, false)
         }
-        return currentIndex
-    }
-
-    fun uniqIdByDeclaration(value: IrDeclaration) = (value as? IrSymbolOwner)?.let { builtInsTable[it.symbol] } ?: table.getOrPut(value) {
-        computeUniqIdByDeclaration(value)
-    }
-
-    open protected fun computeUniqIdByDeclaration(value: IrDeclaration): UniqId {
-        val index = if (value.origin == IrDeclarationOrigin.FAKE_OVERRIDE ||
-            !value.isExported()
-            || value is IrVariable
-            || value is IrValueParameter
-            || value is IrAnonymousInitializer
-            || value is IrLocalDelegatedProperty
-        ) {
-            UniqId(currentIndex++, true)
-        } else {
-            UniqId(value.hashedMangle, false)
-        }
-
         return index
     }
+
+    open fun computeUniqIdByDeclaration(declaration: IrDeclaration): UniqId {
+        return table.getOrPut(declaration) {
+            UniqId(mangler.hashedMangleImpl(declaration), false)
+        }
+    }
+
+    fun isExportedDeclaration(declaration: IrDeclaration): Boolean = mangler.isExportedImpl(declaration)
+}
+
+class DeclarationTable(
+    private val descriptorTable: DescriptorTable,
+    private val globalDeclarationTable: GlobalDeclarationTable,
+    startIndex: Long
+) {
+    private val table = mutableMapOf<IrDeclaration, UniqId>()
+
+    private fun IrDeclaration.isLocalDeclaration(): Boolean {
+        return origin == IrDeclarationOrigin.FAKE_OVERRIDE || !isExportedDeclaration(this) || this is IrVariable || this is IrValueParameter || this is IrAnonymousInitializer || this is IrLocalDelegatedProperty
+    }
+
+    private var localIndex = startIndex
+
+    fun isExportedDeclaration(declaration: IrDeclaration) = globalDeclarationTable.isExportedDeclaration(declaration)
+
+    private fun computeUniqIdByDeclaration(declaration: IrDeclaration): UniqId {
+        return if (declaration.isLocalDeclaration()) {
+            table.getOrPut(declaration) { UniqId(localIndex++, true) }
+        } else globalDeclarationTable.computeUniqIdByDeclaration(declaration)
+    }
+
+    fun uniqIdByDeclaration(declaration: IrDeclaration): UniqId {
+        val uniqId = computeUniqIdByDeclaration(declaration)
+        descriptorTable.put(declaration.descriptor, uniqId)
+        return uniqId
+    }
 }
 
 // This is what we pre-populate tables with

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

@@ -125,7 +125,6 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.Visibility as Pro
 open class IrFileSerializer(
     val logger: LoggingContext,
     private val declarationTable: DeclarationTable,
-    mangler: KotlinMangler,
     private val bodiesOnlyForInlines: Boolean = false
 ) {
 
@@ -148,7 +147,7 @@ open class IrFileSerializer(
     private val protoBodyArray = mutableListOf<XStatementOrExpression>()
 
     private val descriptorReferenceSerializer =
-        DescriptorReferenceSerializer(declarationTable, { serializeString(it) }, { serializeFqName(it) }, mangler)
+        DescriptorReferenceSerializer(declarationTable, { serializeString(it) }, { serializeFqName(it) })
 
     sealed class XStatementOrExpression {
         abstract fun toByteArray(): ByteArray

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

@@ -11,12 +11,7 @@ import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.library.IrIrSerializedIrFile
 import org.jetbrains.kotlin.library.SerializedIrModule
 
-abstract class IrModuleSerializer<F : IrFileSerializer>(
-    protected val logger: LoggingContext,
-    protected val mangler: KotlinMangler
-) {
-
-
+abstract class IrModuleSerializer<F : IrFileSerializer>(protected val logger: LoggingContext) {
     abstract fun createSerializerForFile(file: IrFile): F
 
     private fun serializeIrFile(file: IrFile): IrIrSerializedIrFile {

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

@@ -16,7 +16,9 @@ import org.jetbrains.kotlin.name.Name
 interface KotlinMangler {
     val String.hashMangle: Long
     val IrDeclaration.hashedMangle: Long
+    fun hashedMangleImpl(declaration: IrDeclaration): Long
     fun IrDeclaration.isExported(): Boolean
+    fun isExportedImpl(declaration: IrDeclaration): Boolean
     val IrFunction.functionName: String
     val IrType.isInlined: Boolean
     val Long.isSpecial: Boolean
@@ -31,14 +33,20 @@ interface KotlinMangler {
 abstract class KotlinManglerImpl : KotlinMangler {
     override val String.hashMangle get() = this.cityHash64()
 
+    override fun hashedMangleImpl(declaration: IrDeclaration): Long {
+        return declaration.uniqSymbolName().hashMangle
+    }
+
     override val IrDeclaration.hashedMangle: Long
-        get() = this.uniqSymbolName().hashMangle
+        get() = hashedMangleImpl(this)
 
 
     // We can't call "with (super) { this.isExported() }" in children.
     // So provide a hook.
     protected open fun IrDeclaration.isPlatformSpecificExported(): Boolean = false
 
+    override fun IrDeclaration.isExported(): Boolean = isExportedImpl(this)
+
     /**
      * Defines whether the declaration is exported, i.e. visible from other modules.
      *
@@ -46,27 +54,27 @@ abstract class KotlinManglerImpl : KotlinMangler {
      * that doesn't depend on any internal transformations (e.g. IR lowering),
      * and so should be computable from the descriptor itself without checking a backend state.
      */
-    override tailrec fun IrDeclaration.isExported(): Boolean {
+    tailrec override fun isExportedImpl(declaration: IrDeclaration): Boolean {
         // TODO: revise
-        val descriptorAnnotations = this.descriptor.annotations
+        val descriptorAnnotations = declaration.descriptor.annotations
 
-        if (this.isPlatformSpecificExported()) return true
+        if (declaration.isPlatformSpecificExported()) return true
 
         if (descriptorAnnotations.hasAnnotation(publishedApiAnnotation)) {
             return true
         }
 
-        if (this.isAnonymousObject)
+        if (declaration.isAnonymousObject)
             return false
 
-        if (this is IrConstructor && constructedClass.kind.isSingleton) {
+        if (declaration is IrConstructor && declaration.constructedClass.kind.isSingleton) {
             // Currently code generator can access the constructor of the singleton,
             // so ignore visibility of the constructor itself.
-            return constructedClass.isExported()
+            return isExportedImpl(declaration.constructedClass)
         }
 
-        if (this is IrFunction) {
-            val descriptor = this.descriptor
+        if (declaration is IrFunction) {
+            val descriptor = declaration.descriptor
             // TODO: this code is required because accessor doesn't have a reference to property.
             if (descriptor is PropertyAccessorDescriptor) {
                 val property = descriptor.correspondingProperty
@@ -74,12 +82,12 @@ abstract class KotlinManglerImpl : KotlinMangler {
             }
         }
 
-        val visibility = when (this) {
-            is IrClass -> this.visibility
-            is IrFunction -> this.visibility
-            is IrProperty -> this.visibility
-            is IrField -> this.visibility
-            is IrTypeAlias -> this.visibility
+        val visibility = when (declaration) {
+            is IrClass -> declaration.visibility
+            is IrFunction -> declaration.visibility
+            is IrProperty -> declaration.visibility
+            is IrField -> declaration.visibility
+            is IrTypeAlias -> declaration.visibility
             else -> null
         }
 
@@ -92,9 +100,9 @@ abstract class KotlinManglerImpl : KotlinMangler {
             return false
         }
 
-        val parent = this.parent
+        val parent = declaration.parent
         if (parent is IrDeclaration) {
-            return parent.isExported()
+            return isExportedImpl(parent)
         }
 
         return true
@@ -205,7 +213,7 @@ abstract class KotlinManglerImpl : KotlinMangler {
 
     val IrClass.typeInfoSymbolName: String
         get() {
-            assert(this.isExported())
+            assert(isExportedImpl(this))
             if (isBuiltInFunction(this))
                 return KotlinMangler.functionClassSymbolName(name)
             return "ktype:" + this.fqNameForIrSerialization.toString()

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

@@ -22,8 +22,8 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.classId
 open class DescriptorReferenceSerializer(
     val declarationTable: DeclarationTable,
     val serializeString: (String) -> ProtoString,
-    val serializeFqName: (FqName) -> ProtoFqName,
-    mangler: KotlinMangler): KotlinMangler by mangler {
+    val serializeFqName: (FqName) -> ProtoFqName
+) {
 
     private fun isEnumSpecialMember(descriptor: DeclarationDescriptor): Boolean {
         if (descriptor !is SimpleFunctionDescriptor) return false
@@ -54,7 +54,7 @@ open class DescriptorReferenceSerializer(
 
         val descriptor = declaration.descriptor
 
-        if (!declaration.isExported() &&
+        if (!declarationTable.isExportedDeclaration(declaration) &&
             !((declaration as? IrDeclarationWithVisibility)?.visibility == Visibilities.INVISIBLE_FAKE)) {
             return null
         }
@@ -111,7 +111,6 @@ open class DescriptorReferenceSerializer(
         } else descriptor.name.toString()
 
         val uniqId = discoverableDescriptorsDeclaration?.let { declarationTable.uniqIdByDeclaration(it) }
-        uniqId?.let { declarationTable.descriptors.put(discoverableDescriptorsDeclaration.descriptor, it) }
 
         val proto = ProtoDescriptorReference.newBuilder()
             .setPackageFqName(serializeFqName(packageFqName))

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/klib.kt

@@ -19,6 +19,9 @@ import org.jetbrains.kotlin.descriptors.impl.CompositePackageFragmentProvider
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.*
+import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsIrLinker
+import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsIrModuleSerializer
+import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.newJsDescriptorUniqId
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.metadata.*
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
@@ -291,8 +294,6 @@ fun serializeModuleIntoKlib(
     moduleFragment: IrModuleFragment,
     nopack: Boolean
 ) {
-    val declarationTable = JsDeclarationTable(moduleFragment.irBuiltins, DescriptorTable())
-
     val descriptorTable = DescriptorTable()
 
     val serializedIr = JsIrModuleSerializer(emptyLoggingContext, moduleFragment.irBuiltins, descriptorTable).serializedIrModule(moduleFragment)

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/lower/serialization/ir/JsDeclarationTable.kt

@@ -17,4 +17,4 @@ class JsDeclarationTable(builtIns: IrBuiltIns, descriptorTable: DescriptorTable)
         loadKnownBuiltins()
         currentIndex = 0
     }
-}
+}

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/lower/serialization/ir/JsIrFileSerializer.kt

@@ -13,4 +13,4 @@ class JsIrFileSerializer(
     logger: LoggingContext,
     declarationTable: DeclarationTable,
     bodiesOnlyForInlines: Boolean = false
-) : IrFileSerializer(logger, declarationTable, JsMangler, bodiesOnlyForInlines)
+) : IrFileSerializer(logger, declarationTable, bodiesOnlyForInlines)

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/lower/serialization/ir/JsIrModuleSerializer.kt

@@ -6,14 +6,17 @@
 package org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir
 
 import org.jetbrains.kotlin.backend.common.LoggingContext
+import org.jetbrains.kotlin.backend.common.serialization.DeclarationTable
 import org.jetbrains.kotlin.backend.common.serialization.DescriptorTable
 import org.jetbrains.kotlin.backend.common.serialization.IrModuleSerializer
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 
-class JsIrModuleSerializer(logger: LoggingContext, private val irBuiltIns: IrBuiltIns, private val descriptorTable: DescriptorTable) :
-    IrModuleSerializer<JsIrFileSerializer>(logger, JsMangler) {
+class JsIrModuleSerializer(logger: LoggingContext, irBuiltIns: IrBuiltIns, private val descriptorTable: DescriptorTable) :
+    IrModuleSerializer<JsIrFileSerializer>(logger) {
+
+    private val globalDeclarationTable = JsGlobalDeclarationTable(irBuiltIns)
 
     override fun createSerializerForFile(file: IrFile): JsIrFileSerializer =
-        JsIrFileSerializer(logger, JsDeclarationTable(irBuiltIns, descriptorTable))
+        JsIrFileSerializer(logger, DeclarationTable(descriptorTable, globalDeclarationTable, 0))
 }