Revert "[Wasm] Support lazy associated object initialisation"

This reverts commit 552ee1ee38.

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmSymbols.kt

@@ -77,7 +77,8 @@ class WasmSymbols(
         getProperty(FqName.fromSegments(listOf("kotlin", "wasm", "internal", "isNotFirstWasmExportCall")))
     )
 
-    internal val tryGetAssociatedObject = getInternalFunction("tryGetAssociatedObject")
+    internal val initAssociatedObjects = getInternalFunction("initAssociatedObjects")
+    internal val addAssociatedObject = getInternalFunction("addAssociatedObject")
 
     override val throwNullPointerException = getInternalFunction("THROW_NPE")
     override val throwISE = getInternalFunction("THROW_ISE")

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/AssociatedObjectsLowering.kt

@@ -7,12 +7,11 @@ package org.jetbrains.kotlin.backend.wasm.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
-import org.jetbrains.kotlin.backend.common.lower.irIfThen
 import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.backend.wasm.WasmSymbols
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.backend.js.utils.associatedObject
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.IrClass
@@ -46,7 +45,7 @@ class AssociatedObjectsLowering(val context: WasmBackendContext) : FileLoweringP
     }
 
     private val visitor = object : IrElementVisitorVoid {
-        val tryGetAssociatedObject = (context.wasmSymbols.tryGetAssociatedObject.owner.body as IrBlockBody).statements
+        val initFunctionStatements = (context.wasmSymbols.initAssociatedObjects.owner.body as IrBlockBody).statements
 
         override fun visitElement(element: IrElement) {
             if (element is IrClass) {
@@ -63,44 +62,49 @@ class AssociatedObjectsLowering(val context: WasmBackendContext) : FileLoweringP
                 if (klassAnnotation.valueArgumentsCount != 1) continue
                 val associatedObject = klassAnnotation.associatedObject() ?: continue
 
-                val builder = cachedBuilder ?: context.createIrBuilder(context.wasmSymbols.tryGetAssociatedObject)
+                val builder = cachedBuilder ?: context.createIrBuilder(context.wasmSymbols.initAssociatedObjects)
                 cachedBuilder = builder
 
-                val selector = builder.createAssociatedObjectSelector(
+                val addCall = builder.createAssociatedObjectAdd(
                     wasmSymbols = context.wasmSymbols,
                     irBuiltIns = context.irBuiltIns,
                     targetClass = declaration.symbol,
                     keyAnnotation = annotationClass.symbol,
                     associatedObject = associatedObject.symbol
                 )
-                tryGetAssociatedObject.add(0, selector)
+                initFunctionStatements.add(addCall)
             }
         }
     }
 }
 
-private fun IrBuilderWithScope.createAssociatedObjectSelector(
+private fun IrBuilderWithScope.createAssociatedObjectAdd(
     wasmSymbols: WasmSymbols,
     irBuiltIns: IrBuiltIns,
     targetClass: IrClassSymbol,
     keyAnnotation: IrClassSymbol,
     associatedObject: IrClassSymbol
-): IrStatement {
+): IrCall = buildStatement(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
-    val classIdParam = irGet(wasmSymbols.tryGetAssociatedObject.owner.valueParameters[0])
+    irCall(wasmSymbols.addAssociatedObject, irBuiltIns.unitType).also { addCall ->
-    val keyIdParam = irGet(wasmSymbols.tryGetAssociatedObject.owner.valueParameters[1])
+        addCall.putValueArgument(
-
+            0,
-    val classId = irCall(wasmSymbols.wasmTypeId, irBuiltIns.intType).also {
+            irGet(wasmSymbols.initAssociatedObjects.owner.valueParameters[0])
-        it.putTypeArgument(0, targetClass.defaultType)
-    }
-    val keyId = irCall(wasmSymbols.wasmTypeId, irBuiltIns.intType).also {
-        it.putTypeArgument(0, keyAnnotation.defaultType)
-    }
-
-    return irIfThen(
-        condition = irEquals(classIdParam, classId),
-        thenPart = irIfThen(
-            condition = irEquals(keyIdParam, keyId),
-            thenPart = irReturn(irGetObjectValue(irBuiltIns.anyType, associatedObject))
         )
-    )
+        addCall.putValueArgument(
+            1,
+            irCall(wasmSymbols.wasmTypeId, irBuiltIns.intType).also {
+                it.putTypeArgument(0, targetClass.defaultType)
+            }
+        )
+        addCall.putValueArgument(
+            2,
+            irCall(wasmSymbols.wasmTypeId, irBuiltIns.intType).also {
+                it.putTypeArgument(0, keyAnnotation.defaultType)
+            }
+        )
+        addCall.putValueArgument(
+            3,
+            irGetObjectValue(irBuiltIns.anyType, associatedObject)
+        )
+    }
 }

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/fir/FirJsES6BoxTestGenerated.java

@@ -10314,12 +10314,6 @@ public class FirJsES6BoxTestGenerated extends AbstractFirJsES6BoxTest {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
         }
 
-        @Test
-        @TestMetadata("findAssociatedObjectLazyness.kt")
-        public void testFindAssociatedObjectLazyness() throws Exception {
-            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectLazyness.kt");
-        }
-
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/fir/FirLightTreeJsBoxTestGenerated.java

@@ -10208,12 +10208,6 @@ public class FirLightTreeJsBoxTestGenerated extends AbstractFirLightTreeJsBoxTes
             runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
         }
 
-        @Test
-        @TestMetadata("findAssociatedObjectLazyness.kt")
-        public void testFindAssociatedObjectLazyness() throws Exception {
-            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectLazyness.kt");
-        }
-
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/fir/FirPsiJsBoxTestGenerated.java

@@ -10208,12 +10208,6 @@ public class FirPsiJsBoxTestGenerated extends AbstractFirPsiJsBoxTest {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
         }
 
-        @Test
-        @TestMetadata("findAssociatedObjectLazyness.kt")
-        public void testFindAssociatedObjectLazyness() throws Exception {
-            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectLazyness.kt");
-        }
-
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrBoxJsES6TestGenerated.java

@@ -10314,12 +10314,6 @@ public class IrBoxJsES6TestGenerated extends AbstractIrBoxJsES6Test {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
         }
 
-        @Test
-        @TestMetadata("findAssociatedObjectLazyness.kt")
-        public void testFindAssociatedObjectLazyness() throws Exception {
-            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectLazyness.kt");
-        }
-
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrBoxJsTestGenerated.java

@@ -10208,12 +10208,6 @@ public class IrBoxJsTestGenerated extends AbstractIrBoxJsTest {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
         }
 
-        @Test
-        @TestMetadata("findAssociatedObjectLazyness.kt")
-        public void testFindAssociatedObjectLazyness() throws Exception {
-            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectLazyness.kt");
-        }
-
         @Test
         @TestMetadata("findAssociatedObject_oldBE.kt")
         public void testFindAssociatedObject_oldBE() throws Exception {

js/js.translator/testData/box/reflection/findAssociatedObjectLazyness.kt

@@ -1,64 +0,0 @@
-@file:OptIn(ExperimentalAssociatedObjects::class)
-
-import kotlin.reflect.AssociatedObjectKey
-import kotlin.reflect.ExperimentalAssociatedObjects
-import kotlin.reflect.findAssociatedObject
-import kotlin.reflect.KClass
-
-private var obj1Init = false
-private var obj2Init = false
-private var obj3Init = false
-
-@AssociatedObjectKey
-annotation class KEY1(val kClass: KClass<*>)
-
-@AssociatedObjectKey
-annotation class KEY2(val kClass: KClass<*>)
-
-private object OBJ1 {
-    init { obj1Init = true }
-}
-
-private object OBJ2 {
-    init { obj2Init = true }
-}
-
-private object OBJ3 {
-    init { obj3Init = true }
-}
-
-@KEY1(OBJ1::class)
-class CLS1
-
-@KEY1(OBJ2::class)
-@KEY2(OBJ3::class)
-class CLS2
-
-fun box(): String {
-    // No objects initialised
-    if (obj1Init) return "FAIL1"
-    if (obj2Init) return "FAIL2"
-    if (obj3Init) return "FAIL3"
-
-    CLS1::class.findAssociatedObject<KEY2>()
-    if (obj1Init) return "FAIL4"
-    if (obj2Init) return "FAIL5"
-    if (obj3Init) return "FAIL6"
-
-    CLS1::class.findAssociatedObject<KEY1>()
-    if (!obj1Init) return "FAIL7"
-    if (obj2Init) return "FAIL8"
-    if (obj3Init) return "FAIL9"
-
-    CLS2::class.findAssociatedObject<KEY1>()
-    if (!obj1Init) return "FAIL10"
-    if (!obj2Init) return "FAIL11"
-    if (obj3Init) return "FAIL12"
-
-    CLS2::class.findAssociatedObject<KEY2>()
-    if (!obj1Init) return "FAIL13"
-    if (!obj2Init) return "FAIL14"
-    if (!obj3Init) return "FAIL15"
-
-    return "OK"
-}

libraries/stdlib/wasm/src/kotlin/reflect/associatedObjectsImpl.kt

@@ -8,16 +8,35 @@ package kotlin.wasm.internal
 import kotlin.reflect.wasm.internal.KClassImpl
 import kotlin.reflect.KClass
 
+internal var associatedObjects: Map<ULong, Any>? = null
+
 @PublishedApi
 internal fun findAssociatedObject(klass: KClass<*>, key: Int): Any? {
     val klassId = (klass as? KClassImpl<*>)?.typeData?.typeId ?: return null
-    return tryGetAssociatedObject(klassId, key)
+
+    val map = associatedObjects ?: run {
+        val newMap: MutableMap<ULong, Any> = mutableMapOf()
+        initAssociatedObjects(newMap)
+        associatedObjects = newMap
+        newMap
+    }
+
+    return map[packIntoULong(klassId, key)]
 }
 
-internal fun tryGetAssociatedObject(
+internal fun packIntoULong(a: Int, b: Int): ULong =
-    @Suppress("UNUSED_PARAMETER") klassId: Int,
+    (a.toUInt().toULong() shl Int.SIZE_BITS) or b.toUInt().toULong()
-    @Suppress("UNUSED_PARAMETER") keyId: Int,
+
-): Any? {
+internal fun addAssociatedObject(
+    mapToInit: MutableMap<ULong, Any>,
+    klassId: Int,
+    keyId: Int,
+    instance: Any
+) {
+    mapToInit[packIntoULong(klassId, keyId)] = instance
+}
+
+internal fun initAssociatedObjects(@Suppress("UNUSED_PARAMETER") mapToInit: MutableMap<ULong, Any>) {
     // Init implicitly with AssociatedObjectsLowering
-    return null
+    // addAssociatedObject(mapToInit, ...)
 }

wasm/wasm.tests/test/org/jetbrains/kotlin/generators/tests/GenerateWasmTests.kt

@@ -21,7 +21,7 @@ fun main(args: Array<String>) {
     val k2BoxTestDir = "multiplatform/k2"
 
     val jsTranslatorTestPattern = "^([^_](.+))\\.kt$"
-    val jsTranslatorReflectionPattern = "^(findAssociatedObject(InSeparatedFile)?(Lazyness)?)\\.kt$"
+    val jsTranslatorReflectionPattern = "^(findAssociatedObject(InSeparatedFile)?)\\.kt$"
     val jsTranslatorEsModulesExcludedDirs = listOf(
         // JsExport is not supported for classes
         "jsExport", "native", "export",

wasm/wasm.tests/tests-gen/org/jetbrains/kotlin/wasm/test/FirWasmJsTranslatorTestGenerated.java

@@ -514,7 +514,7 @@ public class FirWasmJsTranslatorTestGenerated extends AbstractFirWasmJsTranslato
     public class Reflection {
         @Test
         public void testAllFilesPresentInReflection() throws Exception {
-            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("js/js.translator/testData/box/reflection"), Pattern.compile("^(findAssociatedObject(InSeparatedFile)?(Lazyness)?)\\.kt$"), null, TargetBackend.WASM, true);
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("js/js.translator/testData/box/reflection"), Pattern.compile("^(findAssociatedObject(InSeparatedFile)?)\\.kt$"), null, TargetBackend.WASM, true);
         }
 
         @Test
@@ -528,12 +528,6 @@ public class FirWasmJsTranslatorTestGenerated extends AbstractFirWasmJsTranslato
         public void testFindAssociatedObjectInSeparatedFile() throws Exception {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
         }
-
-        @Test
-        @TestMetadata("findAssociatedObjectLazyness.kt")
-        public void testFindAssociatedObjectLazyness() throws Exception {
-            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectLazyness.kt");
-        }
     }
 
     @Nested

wasm/wasm.tests/tests-gen/org/jetbrains/kotlin/wasm/test/K1WasmJsTranslatorTestGenerated.java

@@ -514,7 +514,7 @@ public class K1WasmJsTranslatorTestGenerated extends AbstractK1WasmJsTranslatorT
     public class Reflection {
         @Test
         public void testAllFilesPresentInReflection() throws Exception {
-            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("js/js.translator/testData/box/reflection"), Pattern.compile("^(findAssociatedObject(InSeparatedFile)?(Lazyness)?)\\.kt$"), null, TargetBackend.WASM, true);
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("js/js.translator/testData/box/reflection"), Pattern.compile("^(findAssociatedObject(InSeparatedFile)?)\\.kt$"), null, TargetBackend.WASM, true);
         }
 
         @Test
@@ -528,12 +528,6 @@ public class K1WasmJsTranslatorTestGenerated extends AbstractK1WasmJsTranslatorT
         public void testFindAssociatedObjectInSeparatedFile() throws Exception {
             runTest("js/js.translator/testData/box/reflection/findAssociatedObjectInSeparatedFile.kt");
         }
-
-        @Test
-        @TestMetadata("findAssociatedObjectLazyness.kt")
-        public void testFindAssociatedObjectLazyness() throws Exception {
-            runTest("js/js.translator/testData/box/reflection/findAssociatedObjectLazyness.kt");
-        }
     }
 
     @Nested