[Wasm] Support for Associated Objects feature

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

@@ -509,6 +509,13 @@ private val builtInsLoweringPhase = makeWasmModulePhase(
     description = "Lower IR builtins"
 )
 
+private val associatedObjectsLowering = makeWasmModulePhase(
+    ::AssociatedObjectsLowering,
+    name = "AssociatedObjectsLowering",
+    description = "Load associated object init body",
+    prerequisite = setOf(localClassExtractionPhase)
+)
+
 private val objectDeclarationLoweringPhase = makeWasmModulePhase(
     ::ObjectDeclarationLowering,
     name = "ObjectDeclarationLowering",
@@ -687,6 +694,9 @@ val wasmPhases = SameTypeNamedCompilerPhase(
             expressionBodyTransformer then
             eraseVirtualDispatchReceiverParametersTypes then
             bridgesConstructionPhase then
+
+            associatedObjectsLowering then
+
             objectDeclarationLoweringPhase then
             fieldInitializersLoweringPhase then
             genericReturnTypeLowering then

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

@@ -73,6 +73,10 @@ class WasmSymbols(
     internal val isNotFirstWasmExportCall: IrPropertySymbol = symbolTable.referenceProperty(
         getProperty(FqName.fromSegments(listOf("kotlin", "wasm", "internal", "isNotFirstWasmExportCall")))
     )
+
+    internal val initAssociatedObjects = getInternalFunction("initAssociatedObjects")
+    internal val addAssociatedObject = getInternalFunction("addAssociatedObject")
+
     internal val throwAsJsException: IrSimpleFunctionSymbol = getInternalFunction("throwAsJsException")
 
     override val throwNullPointerException = getInternalFunction("THROW_NPE")

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

@@ -0,0 +1,114 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+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.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.backend.js.utils.associatedObject
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+
+/**
+ * Adding associated objects into objects dictionary
+ *
+ * For code like this:
+ * annotation class Key(klass: KClass<*>)
+ * object OBJ
+ *
+ * @Key(OBJ::class)
+ * class C
+ *
+ * add initializer expression into initAssociatedObjects() body:
+ * internal fun initAssociatedObjects(mapToInit: MutableMap<Int, MutableMap<Int, Any>>) {
+ *   ...
+ *   addAssociatedObject(mapToInit, C.klassId, Key.klassId, OBJ)
+ * }
+ */
+class AssociatedObjectsLowering(val context: WasmBackendContext) : FileLoweringPass {
+    override fun lower(irFile: IrFile) {
+        irFile.acceptChildrenVoid(visitor)
+    }
+
+    private val visitor = object : IrElementVisitorVoid {
+        val initFunctionStatements = (context.wasmSymbols.initAssociatedObjects.owner.body as IrBlockBody).statements
+
+        override fun visitElement(element: IrElement) {
+            if (element is IrClass) {
+                element.acceptChildrenVoid(this)
+            }
+        }
+
+        override fun visitClass(declaration: IrClass) {
+            super.visitClass(declaration)
+
+            var cachedBuilder: IrBuilderWithScope? = null
+            for (klassAnnotation in declaration.annotations) {
+                val annotationClass = klassAnnotation.symbol.owner.parentClassOrNull ?: continue
+                if (klassAnnotation.valueArgumentsCount != 1) continue
+                val associatedObject = klassAnnotation.associatedObject() ?: continue
+
+                val builder = cachedBuilder ?: context.createIrBuilder(context.wasmSymbols.initAssociatedObjects)
+                cachedBuilder = builder
+
+                val addCall = builder.createAssociatedObjectAdd(
+                    wasmSymbols = context.wasmSymbols,
+                    irBuiltIns = context.irBuiltIns,
+                    targetClass = declaration.symbol,
+                    keyAnnotation = annotationClass.symbol,
+                    associatedObject = associatedObject.symbol
+                )
+                initFunctionStatements.add(addCall)
+            }
+        }
+    }
+}
+
+private fun IrBuilderWithScope.createAssociatedObjectAdd(
+    wasmSymbols: WasmSymbols,
+    irBuiltIns: IrBuiltIns,
+    targetClass: IrClassSymbol,
+    keyAnnotation: IrClassSymbol,
+    associatedObject: IrClassSymbol
+): IrCall = buildStatement(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+    irCall(wasmSymbols.addAssociatedObject, irBuiltIns.unitType).also { addCall ->
+        addCall.putValueArgument(
+            0,
+            irGet(wasmSymbols.initAssociatedObjects.owner.valueParameters[0])
+        )
+        val classIdGetter = when (targetClass.owner.isInterface) {
+            true -> wasmSymbols.wasmInterfaceId
+            false -> wasmSymbols.wasmClassId
+        }
+        addCall.putValueArgument(
+            1,
+            irCall(classIdGetter, irBuiltIns.intType).also {
+                it.putTypeArgument(0, targetClass.defaultType)
+            }
+        )
+        addCall.putValueArgument(
+            2,
+            irCall(wasmSymbols.wasmClassId, irBuiltIns.intType).also {
+                it.putTypeArgument(0, keyAnnotation.defaultType)
+            }
+        )
+        addCall.putValueArgument(
+            3,
+            irGetObjectValue(irBuiltIns.anyType, associatedObject)
+        )
+    }
+}

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

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package kotlin.reflect
+
+import kotlin.wasm.internal.wasmClassId
+import kotlin.wasm.internal.findAssociatedObject
+
+/**
+ * The experimental marker for associated objects API.
+ *
+ * Any usage of a declaration annotated with `@ExperimentalAssociatedObjects` must be accepted either by
+ * annotating that usage with the [OptIn] annotation, e.g. `@OptIn(ExperimentalAssociatedObjects::class)`,
+ * or by using the compiler argument `-opt-in=kotlin.reflect.ExperimentalAssociatedObjects`.
+ */
+@RequiresOptIn(level = RequiresOptIn.Level.ERROR)
+@Retention(value = AnnotationRetention.BINARY)
+public annotation class ExperimentalAssociatedObjects
+
+/**
+ * Makes the annotated annotation class an associated object key.
+ *
+ * An associated object key annotation should have single [KClass] parameter.
+ * When applied to a class with reference to an object declaration as an argument, it binds
+ * the object to the class, making this binding discoverable at runtime using [findAssociatedObject].
+ */
+@ExperimentalAssociatedObjects
+@Retention(AnnotationRetention.BINARY)
+@Target(AnnotationTarget.ANNOTATION_CLASS)
+public annotation class AssociatedObjectKey
+
+/**
+ * If [T] is an @[AssociatedObjectKey]-annotated annotation class and [this] class is annotated with @[T] (`S::class`),
+ * returns object `S`.
+ *
+ * Otherwise returns `null`.
+ */
+@ExperimentalAssociatedObjects
+@Suppress("NON_PUBLIC_CALL_FROM_PUBLIC_INLINE")
+public inline fun <reified T : Annotation> KClass<*>.findAssociatedObject(): Any? =
+    findAssociatedObject(this, wasmClassId<T>())

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

@@ -23,7 +23,7 @@ internal object ErrorKClass : KClass<Nothing> {
     override fun isInstance(value: Any?): Boolean = error("Can's check isInstance on ErrorKClass")
 }
 
-internal class KClassImpl<T : Any>(private val typeData: TypeInfoData) : KClass<T> {
+internal class KClassImpl<T : Any>(internal val typeData: TypeInfoData) : KClass<T> {
     override val simpleName: String get() = typeData.typeName
     override val qualifiedName: String =
         if (typeData.packageName.isEmpty()) typeData.typeName else "${typeData.packageName}.${typeData.typeName}"

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+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
+
+    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 =
+    (a.toUInt().toULong() shl Int.SIZE_BITS) or b.toUInt().toULong()
+
+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
+    // addAssociatedObject(mapToInit, ...)
+}