~ add a test for the old BE

The `findAssociatedObject` function always returns `null` in old BE

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

@@ -8,8 +8,10 @@ package org.jetbrains.kotlin.ir.backend.js
 import org.jetbrains.kotlin.backend.common.ir.isMemberOfOpenClass
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.backend.js.export.isExported
+import org.jetbrains.kotlin.ir.backend.js.utils.associatedObject
 import org.jetbrains.kotlin.ir.backend.js.utils.getJsName
 import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
+import org.jetbrains.kotlin.ir.backend.js.utils.isAssociatedObjectAnnotatedAnnotation
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
@@ -78,8 +80,19 @@ private fun removeUselessDeclarations(module: IrModuleFragment, usefulDeclaratio
                 process(declaration)
             }
 
+            private fun IrConstructorCall.shouldKeepAnnotation(): Boolean {
+                associatedObject()?.let { obj ->
+                    if (obj !in usefulDeclarations) return false
+                }
+                return true
+            }
+
             override fun visitClass(declaration: IrClass) {
                 process(declaration)
+                // Remove annotations for `findAssociatedObject` feature, which reference objects eliminated by the DCE.
+                // Otherwise `JsClassGenerator.generateAssociatedKeyProperties` will try to reference the object factory (which is removed).
+                // That will result in an error from the Namer. It cannot generate a name for an absent declaration.
+                declaration.annotations = declaration.annotations.filter { it.shouldKeepAnnotation() }
             }
 
             // TODO bring back the primary constructor fix
@@ -118,6 +131,10 @@ fun usefulDeclarations(roots: Iterable<IrDeclaration>, context: JsIrBackendConte
     val contagiousReachableDeclarations = hashSetOf<IrOverridableDeclaration<*>>()
     val constructedClasses = hashSetOf<IrClass>()
 
+    val classesWithObjectAssociations = hashSetOf<IrClass>()
+    val referencedJsClasses = hashSetOf<IrDeclaration>()
+    val referencedJsClassesFromExpressions = hashSetOf<IrClass>()
+
     fun IrDeclaration.enqueue(
         from: IrDeclaration?,
         description: String?,
@@ -211,6 +228,14 @@ fun usefulDeclarations(roots: Iterable<IrDeclaration>, context: JsIrBackendConte
                         }
                     }
                 }
+
+                declaration.annotations.forEach {
+                    val annotationClass = it.symbol.owner.constructedClass
+                    if (annotationClass.isAssociatedObjectAnnotatedAnnotation) {
+                        classesWithObjectAssociations += declaration
+                        annotationClass.enqueue("@AssociatedObject annotated annotation class")
+                    }
+                }
             }
 
             if (declaration is IrSimpleFunction && declaration.isFakeOverride) {
@@ -268,7 +293,13 @@ fun usefulDeclarations(roots: Iterable<IrDeclaration>, context: JsIrBackendConte
                             constructor.enqueue("intrinsic: jsBoxIntrinsic")
                         }
                         context.intrinsics.jsClass -> {
-                            (expression.getTypeArgument(0)!!.classifierOrFail.owner as IrDeclaration).enqueue("intrinsic: jsClass")
+                            val ref = expression.getTypeArgument(0)!!.classifierOrFail.owner as IrDeclaration
+                            ref.enqueue("intrinsic: jsClass")
+                            referencedJsClasses += ref
+                        }
+                        context.intrinsics.jsGetKClassFromExpression -> {
+                            val ref = expression.getTypeArgument(0)?.classOrNull ?: context.irBuiltIns.anyClass
+                            referencedJsClassesFromExpressions += ref.owner
                         }
                         context.intrinsics.jsObjectCreate.symbol -> {
                             val classToCreate = expression.getTypeArgument(0)!!.classifierOrFail.owner as IrClass
@@ -314,6 +345,21 @@ fun usefulDeclarations(roots: Iterable<IrDeclaration>, context: JsIrBackendConte
             return null
         }
 
+        // Handle objects, constructed via `findAssociatedObject` annotation
+        referencedJsClassesFromExpressions += constructedClasses.filterDescendantsOf(referencedJsClassesFromExpressions) // Grow the set of possible results of instance::class expression
+        for (klass in classesWithObjectAssociations) {
+            if (klass !in referencedJsClasses && klass !in referencedJsClassesFromExpressions) continue
+
+            for (annotation in klass.annotations) {
+                val annotationClass = annotation.symbol.owner.constructedClass
+                if (annotationClass !in referencedJsClasses) continue
+
+                annotation.associatedObject()?.let { obj ->
+                    context.mapping.objectToGetInstanceFunction[obj]?.enqueue(klass, "associated object factory")
+                }
+            }
+        }
+
         for (klass in constructedClasses) {
             // TODO a better way to support inverse overrides.
             for (declaration in ArrayList(klass.declarations)) {
@@ -369,3 +415,29 @@ fun usefulDeclarations(roots: Iterable<IrDeclaration>, context: JsIrBackendConte
 
     return result
 }
+
+private fun Collection<IrClass>.filterDescendantsOf(bases: Collection<IrClass>): Collection<IrClass> {
+    val visited = hashSetOf<IrClass>()
+    val baseDescendants = hashSetOf<IrClass>()
+    baseDescendants += bases
+
+    fun overridesAnyBase(klass: IrClass): Boolean {
+        if (klass in baseDescendants) return true
+        if (klass in visited) return false
+
+        visited += klass
+
+        klass.superTypes.forEach {
+            (it.classifierOrNull as? IrClassSymbol)?.owner?.let {
+                if (overridesAnyBase(it)) {
+                    baseDescendants += klass
+                    return true
+                }
+            }
+        }
+
+        return false
+    }
+
+    return this.filter { overridesAnyBase(it) }
+}

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.ir.backend.js.export.isExported
 import org.jetbrains.kotlin.ir.backend.js.utils.*
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrClassReference
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
 import org.jetbrains.kotlin.ir.types.IrType
@@ -189,6 +190,8 @@ class JsClassGenerator(private val irClass: IrClass, val context: JsGenerationCo
 
         metadataLiteral.propertyInitializers += generateSuperClasses()
 
+        metadataLiteral.propertyInitializers += generateAssociatedKeyProperties()
+
         if (isCoroutineClass()) {
             metadataLiteral.propertyInitializers += generateSuspendArity()
         }
@@ -218,6 +221,31 @@ class JsClassGenerator(private val irClass: IrClass, val context: JsGenerationCo
             )
         )
     }
+
+    private fun generateAssociatedKeyProperties(): List<JsPropertyInitializer> {
+        var result = emptyList<JsPropertyInitializer>()
+
+        context.getAssociatedObjectKey(irClass)?.let { key ->
+            result = result + JsPropertyInitializer(JsStringLiteral("associatedObjectKey"), JsIntLiteral(key))
+        }
+
+        val associatedObjects = irClass.annotations.mapNotNull { annotation ->
+            val annotationClass = annotation.symbol.owner.constructedClass
+            context.getAssociatedObjectKey(annotationClass)?.let { key ->
+                annotation.associatedObject()?.let { obj ->
+                    context.staticContext.backendContext.mapping.objectToGetInstanceFunction[obj]?.let { factory ->
+                        JsPropertyInitializer(JsIntLiteral(key), context.staticContext.getNameForStaticFunction(factory).makeRef())
+                    }
+                }
+            }
+        }
+
+        if (associatedObjects.isNotEmpty()) {
+            result = result + JsPropertyInitializer(JsStringLiteral("associatedObjects"), JsObjectLiteral(associatedObjects))
+        }
+
+        return result
+    }
 }
 
 private val IrClassifierSymbol.isInterface get() = (owner as? IrClass)?.isInterface == true

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/utils/AnnotationUtils.kt

@@ -7,10 +7,11 @@ package org.jetbrains.kotlin.ir.backend.js.utils
 
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrClassReference
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
-import org.jetbrains.kotlin.ir.util.getAnnotation
-import org.jetbrains.kotlin.ir.util.hasAnnotation
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
@@ -45,4 +46,15 @@ fun IrDeclarationWithName.getJsNameOrKotlinName(): Name =
     when (val jsName = getJsName()) {
         null -> name
         else -> Name.identifier(jsName)
-    }
+    }
+
+private val associatedObjectKeyAnnotationFqName = FqName("kotlin.reflect.AssociatedObjectKey")
+
+val IrClass.isAssociatedObjectAnnotatedAnnotation: Boolean
+    get() = isAnnotationClass && annotations.any { it.symbol.owner.constructedClass.fqNameWhenAvailable == associatedObjectKeyAnnotationFqName }
+
+fun IrConstructorCall.associatedObject(): IrClass? {
+    if (!symbol.owner.constructedClass.isAssociatedObjectAnnotatedAnnotation) return null
+    val klass = ((getValueArgument(0) as? IrClassReference)?.symbol as? IrClassSymbol)?.owner ?: return null
+    return if (klass.isObject) klass else null
+}

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/utils/IrNamer.kt

@@ -22,4 +22,5 @@ interface IrNamer {
     fun getNameForProperty(property: IrProperty): JsName
     fun getRefForExternalClass(klass: IrClass): JsNameRef
     fun getNameForLoop(loop: IrLoop): JsName?
+    fun getAssociatedObjectKey(irClass: IrClass): Int?
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/utils/IrNamerImpl.kt

@@ -73,4 +73,14 @@ class IrNamerImpl(private val newNameTables: NameTables) : IrNamer {
                 error("Unsupported external class parent $parent")
         }
     }
+
+    private val associatedObjectKeyMap = mutableMapOf<IrClass, Int>()
+
+    override fun getAssociatedObjectKey(irClass: IrClass): Int? {
+        if (irClass.isAssociatedObjectAnnotatedAnnotation) {
+
+            return associatedObjectKeyMap.getOrPut(irClass) { associatedObjectKeyMap.size }
+        }
+        return null
+    }
 }

js/js.frontend/src/org/jetbrains/kotlin/js/resolve/diagnostics/JsReflectionAPICallChecker.kt

@@ -18,11 +18,20 @@ package org.jetbrains.kotlin.js.resolve.diagnostics
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.builtins.ReflectionTypes
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.diagnostics.Errors.UNSUPPORTED
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.calls.checkers.AbstractReflectionApiCallChecker
 import org.jetbrains.kotlin.resolve.calls.checkers.CallCheckerContext
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.storage.StorageManager
 
+private val ADDITIONAL_ALLOWED_CLASSES = setOf(
+    FqName("kotlin.reflect.AssociatedObjectKey"),
+    FqName("kotlin.reflect.ExperimentalAssociatedObjects")
+)
+
 class JsReflectionAPICallChecker(
     reflectionTypes: ReflectionTypes,
     storageManager: StorageManager
@@ -30,6 +39,12 @@ class JsReflectionAPICallChecker(
     override val isWholeReflectionApiAvailable: Boolean
         get() = false
 
+    override fun isAllowedReflectionApi(descriptor: CallableDescriptor, containingClass: ClassDescriptor): Boolean {
+        return super.isAllowedReflectionApi(descriptor, containingClass) ||
+                containingClass.fqNameSafe in ADDITIONAL_ALLOWED_CLASSES ||
+                descriptor.name.asString() == "findAssociatedObject"
+    }
+
     override fun report(element: PsiElement, context: CallCheckerContext) {
         context.trace.report(UNSUPPORTED.on(element, "This reflection API is not supported yet in JavaScript"))
     }

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrBoxJsTestGenerated.java

@@ -6751,6 +6751,16 @@ public class IrBoxJsTestGenerated extends AbstractIrBoxJsTest {
             runTest("js/js.translator/testData/box/reflection/external.kt");
         }
 
+        @TestMetadata("findAssociatedObject.kt")
+        public void testFindAssociatedObject() throws Exception {
+            runTest("js/js.translator/testData/box/reflection/findAssociatedObject.kt");
+        }
+
+        @TestMetadata("findAssociatedObject_oldBE.kt")
+        public void testFindAssociatedObject_oldBE() throws Exception {
+            runTest("js/js.translator/testData/box/reflection/findAssociatedObject_oldBE.kt");
+        }
+
         @TestMetadata("kClass.kt")
         public void testKClass() throws Exception {
             runTest("js/js.translator/testData/box/reflection/kClass.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/BoxJsTestGenerated.java

@@ -6781,6 +6781,16 @@ public class BoxJsTestGenerated extends AbstractBoxJsTest {
             runTest("js/js.translator/testData/box/reflection/external.kt");
         }
 
+        @TestMetadata("findAssociatedObject.kt")
+        public void testFindAssociatedObject() throws Exception {
+            runTest("js/js.translator/testData/box/reflection/findAssociatedObject.kt");
+        }
+
+        @TestMetadata("findAssociatedObject_oldBE.kt")
+        public void testFindAssociatedObject_oldBE() throws Exception {
+            runTest("js/js.translator/testData/box/reflection/findAssociatedObject_oldBE.kt");
+        }
+
         @TestMetadata("jsClass.kt")
         public void testJsClass() throws Exception {
             runTest("js/js.translator/testData/box/reflection/jsClass.kt");

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

@@ -0,0 +1,97 @@
+// IGNORE_BACKEND: JS
+// KJS_WITH_FULL_RUNTIME
+
+import kotlin.reflect.*
+
+@OptIn(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated1(val kClass: KClass<*>)
+
+@OptIn(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated2(val kClass: KClass<*>)
+
+@OptIn(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated3(val kClass: KClass<*>)
+
+@Associated1(Bar::class)
+@Associated2(Baz::class)
+class Foo
+
+object Bar
+object Baz
+
+private class C(var list: List<String>?)
+
+private interface I1 {
+    fun foo(): Int
+    fun bar(c: C)
+}
+
+private object I1Impl : I1 {
+    override fun foo() = 42
+    override fun bar(c: C) {
+        c.list = mutableListOf("zzz")
+    }
+}
+
+@Associated1(I1Impl::class)
+private class I1ImplHolder
+
+private interface I2 {
+    fun foo(): Int
+}
+
+private object I2Impl : I2 {
+    override fun foo() = 17
+}
+
+@Associated1(I2Impl::class)
+private class I2ImplHolder
+
+@Associated2(A.Companion::class)
+class A {
+    companion object : I2 {
+        override fun foo() = 20
+    }
+}
+
+@OptIn(ExperimentalAssociatedObjects::class)
+fun KClass<*>.getAssociatedObjectByAssociated2(): Any? {
+    return this.findAssociatedObject<Associated2>()
+}
+
+@OptIn(ExperimentalAssociatedObjects::class)
+fun box(): String {
+
+    if (Foo::class.findAssociatedObject<Associated1>() != Bar) return "fail 1"
+
+    if (Foo::class.findAssociatedObject<Associated2>() != Baz) return "fail 2"
+
+    if (Foo::class.findAssociatedObject<Associated3>() != null) return "fail 3"
+
+    if (Bar::class.findAssociatedObject<Associated1>() != null) return "fail 4"
+
+    val i1 = I1ImplHolder::class.findAssociatedObject<Associated1>() as I1
+    if (i1.foo() != 42) return "fail 5"
+
+    val c = C(null)
+    i1.bar(c)
+    if (c.list!![0] != "zzz") return "fail 6"
+
+    val i2 = I2ImplHolder()::class.findAssociatedObject<Associated1>() as I2
+    if (i2.foo() != 17) return "fail 7"
+
+    val a = A::class.findAssociatedObject<Associated2>() as I2
+    if (a.foo() != 20) return "fail 8"
+
+    if (Foo::class.getAssociatedObjectByAssociated2() != Baz) return "fail 9"
+
+    if ((A::class.getAssociatedObjectByAssociated2() as I2).foo() != 20) return "fail 10"
+
+    return "OK"
+}

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

@@ -0,0 +1,22 @@
+// EXPECTED_REACHABLE_NODES: 1321
+// IGNORE_BACKEND: JS_IR
+
+import kotlin.reflect.*
+
+@OptIn(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated1(val kClass: KClass<*>)
+
+@Associated1(Bar::class)
+class Foo
+
+object Bar
+
+@OptIn(ExperimentalAssociatedObjects::class)
+fun box(): String {
+    // This API is not implented in the old backend.
+    if (Foo::class.findAssociatedObject<Associated1>() != null) return "fail 1"
+
+    return "OK"
+}

libraries/stdlib/js-ir/src/kotlin/reflection_js-ir.kt

@@ -0,0 +1,19 @@
+/*
+ * 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.
+ */
+
+import kotlin.reflect.*
+import kotlin.reflect.js.internal.*
+
+@PublishedApi
+internal fun <T : Annotation> KClass<*>.findAssociatedObject(annotationClass: KClass<T>): Any? {
+    return if (this is KClassImpl<*> && annotationClass is KClassImpl<T>) {
+        val key = annotationClass.jClass.asDynamic().`$metadata$`.associatedObjectKey?.unsafeCast<Int>() ?: return null
+        val map = this.jClass.asDynamic().`$metadata$`.associatedObjects ?: return null
+        val factory = map[key] ?: return null
+        return factory()
+    } else {
+        null
+    }
+}

libraries/stdlib/js-v1/src/kotlin/reflection_js-v1.kt

@@ -0,0 +1,12 @@
+/*
+ * 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.
+ */
+
+import kotlin.reflect.KClass
+
+@PublishedApi
+internal fun <T : Annotation> KClass<*>.findAssociatedObject(annotationClass: KClass<T>): Any? {
+    // This API is not supported in js-v1. Return `null` to be source-compatible with js-ir.
+    return null
+}

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

@@ -0,0 +1,41 @@
+/*
+ * 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 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 `-Xopt-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
+public inline fun <reified T : Annotation> KClass<*>.findAssociatedObject(): Any? =
+    this.findAssociatedObject(T::class)

libraries/stdlib/js/src/kotlin/reflect/reflection.kt

@@ -75,4 +75,4 @@ internal fun <T : Any> getKClass1(jClass: JsClass<T>): KClass<T> {
     } else {
         SimpleKClassImpl(jClass)
     }
-}
+}