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

2020-04-16 17:33:44 +03:00
parent f4b6e1bff6
commit 442331acc9
14 changed files with 354 additions and 5 deletions
@@ -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) }
+}
@@ -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
@@ -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
+}
@@ -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?
 }
@@ -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
+    }
 }
@@ -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"))
    }
@@ -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");
@@ -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");
@@ -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"
+}
@@ -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"
+}
@@ -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
+    }
+}
@@ -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
+}
@@ -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)
@@ -75,4 +75,4 @@ internal fun <T : Any> getKClass1(jClass: JsClass<T>): KClass<T> {
    } else {
        SimpleKClassImpl(jClass)
    }
-}
+}