From 2f0f88062a9bcdab00c3fe5f1f3d73c27e885279 Mon Sep 17 00:00:00 2001
From: Ilya Goncharov <ilya.goncharov@jetbrains.com>
Date: Wed, 4 Aug 2021 14:25:39 +0300
Subject: [PATCH] [JS IR] Change global map in context to mapping

^KT-46204 fixed
---
 .../jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt | 2 --
 .../js/lower/coroutines/JsSuspendArityStoreLowering.kt   | 9 +++++----
 .../backend/js/transformers/irToJs/JsClassGenerator.kt   | 7 ++++---
 .../src/org/jetbrains/kotlin/ir/backend/js/JsMapping.kt  | 2 ++
 4 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt b/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt
index 373dd71c921..ad528e344ee 100644
--- a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt
+++ b/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
diff --git a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/JsSuspendArityStoreLowering.kt b/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/JsSuspendArityStoreLowering.kt
index 256659a46ce..62868206b4e 100644
--- a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/JsSuspendArityStoreLowering.kt
+++ b/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
diff --git a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/JsClassGenerator.kt b/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/JsClassGenerator.kt
index d3de84ab5a0..f750ef956bb 100644
--- a/compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/JsClassGenerator.kt
+++ b/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? {
-        val arity = context.staticContext.backendContext.suspendArityStore
-            .getValue(irClass)
+    private fun generateSuspendArity(): JsPropertyInitializer {
+        val arity = context.staticContext.backendContext.mapping.suspendArityStore[irClass]!!
+            .map { it.valueParameters.size }
+            .distinct()
             .map { JsIntLiteral(it) }
 
         return JsPropertyInitializer(JsNameRef(Namer.METADATA_SUSPEND_ARITY), JsArrayLiteral(arity))
diff --git a/compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/JsMapping.kt b/compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/JsMapping.kt
index 9e5eceeec15..267d28dfeb9 100644
--- a/compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/JsMapping.kt
+++ b/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>>()
 }