IR JS: support findAssociatedObject feature (KT-37418 fixed)

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"
+}