[Wasm] Support lazy associated object initialisation

Fix #KT-63939

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

@@ -77,8 +77,7 @@ class WasmSymbols(
         getProperty(FqName.fromSegments(listOf("kotlin", "wasm", "internal", "isNotFirstWasmExportCall")))
     )
 
-    internal val initAssociatedObjects = getInternalFunction("initAssociatedObjects")
+    internal val tryGetAssociatedObject = getInternalFunction("tryGetAssociatedObject")
-    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,11 +7,12 @@ 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.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.backend.js.utils.associatedObject
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.IrClass
@@ -45,7 +46,7 @@ class AssociatedObjectsLowering(val context: WasmBackendContext) : FileLoweringP
     }
 
     private val visitor = object : IrElementVisitorVoid {
-        val initFunctionStatements = (context.wasmSymbols.initAssociatedObjects.owner.body as IrBlockBody).statements
+        val tryGetAssociatedObject = (context.wasmSymbols.tryGetAssociatedObject.owner.body as IrBlockBody).statements
 
         override fun visitElement(element: IrElement) {
             if (element is IrClass) {
@@ -62,49 +63,44 @@ class AssociatedObjectsLowering(val context: WasmBackendContext) : FileLoweringP
                 if (klassAnnotation.valueArgumentsCount != 1) continue
                 val associatedObject = klassAnnotation.associatedObject() ?: continue
 
-                val builder = cachedBuilder ?: context.createIrBuilder(context.wasmSymbols.initAssociatedObjects)
+                val builder = cachedBuilder ?: context.createIrBuilder(context.wasmSymbols.tryGetAssociatedObject)
                 cachedBuilder = builder
 
-                val addCall = builder.createAssociatedObjectAdd(
+                val selector = builder.createAssociatedObjectSelector(
                     wasmSymbols = context.wasmSymbols,
                     irBuiltIns = context.irBuiltIns,
                     targetClass = declaration.symbol,
                     keyAnnotation = annotationClass.symbol,
                     associatedObject = associatedObject.symbol
                 )
-                initFunctionStatements.add(addCall)
+                tryGetAssociatedObject.add(0, selector)
             }
         }
     }
 }
 
-private fun IrBuilderWithScope.createAssociatedObjectAdd(
+private fun IrBuilderWithScope.createAssociatedObjectSelector(
     wasmSymbols: WasmSymbols,
     irBuiltIns: IrBuiltIns,
     targetClass: IrClassSymbol,
     keyAnnotation: IrClassSymbol,
     associatedObject: IrClassSymbol
-): IrCall = buildStatement(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+): IrStatement {
-    irCall(wasmSymbols.addAssociatedObject, irBuiltIns.unitType).also { addCall ->
+    val classIdParam = irGet(wasmSymbols.tryGetAssociatedObject.owner.valueParameters[0])
-        addCall.putValueArgument(
+    val keyIdParam = irGet(wasmSymbols.tryGetAssociatedObject.owner.valueParameters[1])
-            0,
+
-            irGet(wasmSymbols.initAssociatedObjects.owner.valueParameters[0])
+    val classId = irCall(wasmSymbols.wasmTypeId, irBuiltIns.intType).also {
-        )
+        it.putTypeArgument(0, targetClass.defaultType)
-        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)
-        )
     }
+    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))
+        )
+    )
 }

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

@@ -10314,6 +10314,12 @@ 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,6 +10208,12 @@ 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,6 +10208,12 @@ 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,6 +10314,12 @@ 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,6 +10208,12 @@ 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

@@ -0,0 +1,64 @@
+@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,35 +8,16 @@ 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 packIntoULong(a: Int, b: Int): ULong =
+internal fun tryGetAssociatedObject(
-    (a.toUInt().toULong() shl Int.SIZE_BITS) or b.toUInt().toULong()
+    @Suppress("UNUSED_PARAMETER") klassId: Int,
-
+    @Suppress("UNUSED_PARAMETER") keyId: Int,
-internal fun addAssociatedObject(
+): Any? {
-    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
-    // addAssociatedObject(mapToInit, ...)
+    return null
 }

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)?)\\.kt$"
+    val jsTranslatorReflectionPattern = "^(findAssociatedObject(InSeparatedFile)?(Lazyness)?)\\.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)?)\\.kt$"), null, TargetBackend.WASM, true);
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("js/js.translator/testData/box/reflection"), Pattern.compile("^(findAssociatedObject(InSeparatedFile)?(Lazyness)?)\\.kt$"), null, TargetBackend.WASM, true);
         }
 
         @Test
@@ -528,6 +528,12 @@ 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)?)\\.kt$"), null, TargetBackend.WASM, true);
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("js/js.translator/testData/box/reflection"), Pattern.compile("^(findAssociatedObject(InSeparatedFile)?(Lazyness)?)\\.kt$"), null, TargetBackend.WASM, true);
         }
 
         @Test
@@ -528,6 +528,12 @@ 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