[JS IR] Change global map in context to mapping

^KT-46204 fixed

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt

@@ -65,8 +65,6 @@ class JsIrBackendContext(
     val fileToInitializationFuns: MutableMap<IrFile, IrSimpleFunction?> = mutableMapOf()
     val fileToInitializerPureness: MutableMap<IrFile, Boolean> = mutableMapOf()
 
-    val suspendArityStore: MutableMap<IrDeclaration, Set<Int>> = mutableMapOf()
-
     val extractedLocalClasses: MutableSet<IrClass> = hashSetOf()
 
     override val builtIns = module.builtIns

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/JsSuspendArityStoreLowering.kt

@@ -10,18 +10,19 @@ import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
-import kotlin.collections.set
 
-class JsSuspendArityStoreLowering(private val context: JsIrBackendContext) : DeclarationTransformer {
+class JsSuspendArityStoreLowering(context: JsIrBackendContext) : DeclarationTransformer {
+
+    private var IrClass.suspendArityStore by context.mapping.suspendArityStore
+
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
         if (declaration !is IrClass) return null
 
         declaration.declarations
             .filterIsInstance<IrSimpleFunction>()
             .filter { it.isSuspend }
-            .map { it.valueParameters.size }
             .let {
-                context.suspendArityStore[declaration] = it.toSet()
+                declaration.suspendArityStore = it
             }
 
         return null

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/JsClassGenerator.kt

@@ -234,9 +234,10 @@ class JsClassGenerator(private val irClass: IrClass, val context: JsGenerationCo
 
     private fun isCoroutineClass(): Boolean = irClass.superTypes.any { it.isSuspendFunctionTypeOrSubtype() }
 
-    private fun generateSuspendArity(): JsPropertyInitializer? {
+    private fun generateSuspendArity(): JsPropertyInitializer {
-        val arity = context.staticContext.backendContext.suspendArityStore
+        val arity = context.staticContext.backendContext.mapping.suspendArityStore[irClass]!!
-            .getValue(irClass)
+            .map { it.valueParameters.size }
+            .distinct()
             .map { JsIntLiteral(it) }
 
         return JsPropertyInitializer(JsNameRef(Namer.METADATA_SUSPEND_ARITY), JsArrayLiteral(arity))

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

@@ -35,6 +35,8 @@ class JsMapping(val state: JsMappingState) : DefaultMapping(state) {
     val enumConstructorOldToNewValueParameters = state.newDeclarationToDeclarationMapping<IrValueDeclaration, IrValueParameter>()
     val enumEntryToCorrespondingField = state.newDeclarationToDeclarationMapping<IrEnumEntry, IrField>()
     val enumClassToInitEntryInstancesFun = state.newDeclarationToDeclarationMapping<IrClass, IrSimpleFunction>()
+
+    val suspendArityStore = state.newDeclarationToDeclarationCollectionMapping<IrClass, Collection<IrSimpleFunction>>()
 }