Add initial experimental support for @AssociatedObjectKey (#2947)

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/Reporting.kt

@@ -24,14 +24,16 @@ internal fun CommonBackendContext.reportCompilationWarning(message: String) {
     report(null, null, message, false)
 }
 
-internal fun error(irFile: IrFile, element: IrElement?, message: String): Nothing {
+internal fun error(irFile: IrFile?, element: IrElement?, message: String): Nothing {
     error(buildString {
         append("Internal compiler error: $message\n")
         if (element == null) {
             append("(IR element is null)")
         } else {
-            val location = element.getCompilerMessageLocation(irFile)
-            append("at $location\n")
+            if (irFile != null) {
+                val location = element.getCompilerMessageLocation(irFile)
+                append("at $location\n")
+            }
 
             val renderedElement = try {
                 element.render()

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/RuntimeNames.kt

@@ -11,4 +11,5 @@ object RuntimeNames {
     val independent = FqName("kotlin.native.internal.Independent")
     val filterExceptions = FqName("kotlin.native.internal.FilterExceptions")
     val kotlinNativeInternalPackageName = FqName.fromSegments(listOf("kotlin", "native", "internal"))
+    val associatedObjectKey = FqName("kotlin.reflect.AssociatedObjectKey")
 }

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/RTTIGenerator.kt

@@ -6,19 +6,17 @@
 package org.jetbrains.kotlin.backend.konan.llvm
 
 import llvm.*
+import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.backend.konan.Context
-import org.jetbrains.kotlin.backend.konan.computePrimitiveBinaryTypeOrNull
 import org.jetbrains.kotlin.backend.konan.descriptors.*
 import org.jetbrains.kotlin.backend.konan.ir.*
+import org.jetbrains.kotlin.backend.konan.ir.isAnonymousObject
+import org.jetbrains.kotlin.backend.konan.ir.isLocal
 import org.jetbrains.kotlin.backend.konan.isExternalObjCClassMethod
-import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrField
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrProperty
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrClassReference
 import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.util.fqNameSafe
-import org.jetbrains.kotlin.ir.util.isAnnotationClass
-import org.jetbrains.kotlin.ir.util.isInterface
+import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.konan.KonanAbiVersion
 import org.jetbrains.kotlin.name.FqName
 
@@ -85,7 +83,8 @@ internal class RTTIGenerator(override val context: Context) : ContextUtils {
             packageName: String?,
             relativeName: String?,
             flags: Int,
-            writableTypeInfo: ConstPointer?) :
+            writableTypeInfo: ConstPointer?,
+            associatedObjects: ConstPointer?) :
 
             Struct(
                     runtime.typeInfoType,
@@ -114,7 +113,9 @@ internal class RTTIGenerator(override val context: Context) : ContextUtils {
 
                     Int32(flags),
 
-                    *listOfNotNull(writableTypeInfo).toTypedArray()
+                    *listOfNotNull(writableTypeInfo).toTypedArray(),
+
+                    associatedObjects
             )
 
     private fun kotlinStringLiteral(string: String?): ConstPointer = if (string == null) {
@@ -219,7 +220,8 @@ internal class RTTIGenerator(override val context: Context) : ContextUtils {
                 reflectionInfo.packageName,
                 reflectionInfo.relativeName,
                 flagsFromClass(irClass),
-                llvmDeclarations.writableTypeInfoGlobal?.pointer
+                llvmDeclarations.writableTypeInfoGlobal?.pointer,
+                associatedObjects = genAssociatedObjects(irClass)
         )
 
         val typeInfoGlobalValue = if (!irClass.typeInfoHasVtableAttached) {
@@ -312,6 +314,31 @@ internal class RTTIGenerator(override val context: Context) : ContextUtils {
         return result.pointer
     }
 
+    private fun genAssociatedObjects(irClass: IrClass): ConstPointer? {
+        val associatedObjects = getAssociatedObjects(irClass)
+        if (associatedObjects.isEmpty()) {
+            return null
+        }
+
+        val associatedObjectTableRecords = associatedObjects.map { (key, value) ->
+            val associatedObjectGetter = generateFunction(
+                    CodeGenerator(context),
+                    functionType(kObjHeaderPtr, false, kObjHeaderPtrPtr),
+                    ""
+            ) {
+                ret(getObjectValue(value, ExceptionHandler.Caller, locationInfo = null))
+            }
+
+            Struct(runtime.associatedObjectTableRecordType, key.typeInfoPtr, constPointer(associatedObjectGetter))
+        }
+
+        return staticData.placeGlobalConstArray(
+                name = "kassociatedobjects:${irClass.fqNameSafe}",
+                elemType = runtime.associatedObjectTableRecordType,
+                elements = associatedObjectTableRecords + Struct(runtime.associatedObjectTableRecordType, null, null)
+        )
+    }
+
     // TODO: extract more code common with generate().
     fun generateSyntheticInterfaceImpl(
             irClass: IrClass,
@@ -370,7 +397,8 @@ internal class RTTIGenerator(override val context: Context) : ContextUtils {
                 packageName = reflectionInfo.packageName,
                 relativeName = reflectionInfo.relativeName,
                 flags = flagsFromClass(irClass) or (if (immutable) TF_IMMUTABLE else 0),
-                writableTypeInfo = writableTypeInfo
+                writableTypeInfo = writableTypeInfo,
+                associatedObjects = null
               ), vtable)
 
         typeInfoWithVtableGlobal.setInitializer(typeInfoWithVtable)
@@ -397,6 +425,42 @@ internal class RTTIGenerator(override val context: Context) : ContextUtils {
     }
 }
 
+private fun getAssociatedObjects(irClass: IrClass): Map<IrClass, IrClass> {
+    val result = mutableMapOf<IrClass, IrClass>()
+
+    irClass.annotations.forEach {
+        val irFile = irClass.getContainingFile()
+
+        val annotationClass = (it.symbol.owner as? IrConstructor)?.constructedClass
+                ?: error(irFile, it, "unexpected annotation")
+
+        if (annotationClass.hasAnnotation(RuntimeNames.associatedObjectKey)) {
+            val argument = it.getValueArgument(0)
+
+            val irClassReference = argument as? IrClassReference
+                    ?: error(irFile, argument, "unexpected annotation argument")
+
+            val associatedObject = irClassReference.symbol.owner
+
+            if (associatedObject !is IrClass || !associatedObject.isObject) {
+                error(irFile, irClassReference, "argument is not a singleton")
+            }
+
+            if (annotationClass in result) {
+                error(
+                        irFile,
+                        it,
+                        "duplicate value for ${annotationClass.name}, previous was ${result[annotationClass]?.name}"
+                )
+            }
+
+            result[annotationClass] = associatedObject
+        }
+    }
+
+    return result
+}
+
 // Keep in sync with Konan_TypeFlags in TypeInfo.h.
 private const val TF_IMMUTABLE = 1
 private const val TF_ACYCLIC   = 2

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/Runtime.kt

@@ -24,6 +24,7 @@ class Runtime(bitcodeFile: String) {
     val writableTypeInfoType = getStructTypeOrNull("WritableTypeInfo")
     val methodTableRecordType = getStructType("MethodTableRecord")
     val globalHashType = getStructType("GlobalHash")
+    val associatedObjectTableRecordType = getStructType("AssociatedObjectTableRecord")
 
     val objHeaderType = getStructType("ObjHeader")
     val objHeaderPtrType = pointerType(objHeaderType)

backend.native/tests/build.gradle

@@ -1438,6 +1438,10 @@ task kclassEnumArgument(type: RunKonanTest) {
     source = "codegen/kclass/kClassEnumArgument.kt"
 }
 
+task associatedObjects1(type: RunKonanTest) {
+    source = "codegen/associatedObjects/associatedObjects1.kt"
+}
+
 task coroutines_simple(type: RunKonanTest) {
     goldValue = "42\n"
     source = "codegen/coroutines/simple.kt"

backend.native/tests/codegen/associatedObjects/associatedObjects1.kt

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+package codegen.associatedObjects.associatedObjects1
+
+import kotlin.test.*
+import kotlin.reflect.*
+
+@Test
+@UseExperimental(ExperimentalAssociatedObjects::class)
+fun testBasics1() {
+    assertSame(Bar, Foo::class.findAssociatedObject<Associated1>())
+    assertSame(Baz, Foo::class.findAssociatedObject<Associated2>())
+    assertSame(null, Foo::class.findAssociatedObject<Associated3>())
+
+    assertSame(null, Bar::class.findAssociatedObject<Associated1>())
+}
+
+@UseExperimental(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated1(val kClass: KClass<*>)
+
+@UseExperimental(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated2(val kClass: KClass<*>)
+
+@UseExperimental(ExperimentalAssociatedObjects::class)
+@AssociatedObjectKey
+@Retention(AnnotationRetention.BINARY)
+annotation class Associated3(val kClass: KClass<*>)
+
+@Associated1(Bar::class)
+@Associated2(Baz::class)
+class Foo
+
+object Bar
+object Baz
+
+@Test
+@UseExperimental(ExperimentalAssociatedObjects::class)
+fun testGlobalOptimizations1() {
+    val i1 = I1ImplHolder::class.findAssociatedObject<Associated1>()!! as I1
+    assertEquals(42, i1.foo())
+}
+
+private interface I1 {
+    fun foo(): Int
+}
+
+private object I1Impl : I1 {
+    override fun foo() = 42
+}
+
+@Associated1(I1Impl::class)
+private class I1ImplHolder
+
+@Test
+@UseExperimental(ExperimentalAssociatedObjects::class)
+fun testGlobalOptimizations2() {
+    val i2 = I2ImplHolder()::class.findAssociatedObject<Associated1>()!! as I2
+    assertEquals(17, i2.foo())
+}
+
+private interface I2 {
+    fun foo(): Int
+}
+
+private object I2Impl : I2 {
+    override fun foo() = 17
+}
+
+@Associated1(I2Impl::class)
+private class I2ImplHolder

runtime/src/main/cpp/TypeInfo.h

@@ -27,6 +27,7 @@ struct WritableTypeInfo;
 #endif
 
 struct ObjHeader;
+struct AssociatedObjectTableRecord;
 
 // An element of sorted by hash in-place array representing methods.
 // For systems where introspection is not needed - only open methods are in
@@ -114,6 +115,9 @@ struct TypeInfo {
     WritableTypeInfo* writableInfo_;
 #endif
 
+    // Null-terminated array.
+    const AssociatedObjectTableRecord* associatedObjects;
+
     // vtable starts just after declared contents of the TypeInfo:
     // void* const vtable_[];
 #ifdef __cplusplus

runtime/src/main/cpp/Types.cpp

@@ -62,6 +62,26 @@ OBJ_GETTER(Kotlin_TypeInfo_getRelativeName, KNativePtr typeInfo) {
   RETURN_OBJ(reinterpret_cast<const TypeInfo*>(typeInfo)->relativeName_);
 }
 
+struct AssociatedObjectTableRecord {
+  const TypeInfo* key;
+  OBJ_GETTER0((*getAssociatedObjectInstance));
+};
+
+OBJ_GETTER(Kotlin_TypeInfo_findAssociatedObject, KNativePtr typeInfo, KNativePtr key) {
+  const AssociatedObjectTableRecord* associatedObjects = reinterpret_cast<const TypeInfo*>(typeInfo)->associatedObjects;
+  if (associatedObjects == nullptr) {
+    RETURN_OBJ(nullptr);
+  }
+
+  for (int index = 0; associatedObjects[index].key != nullptr; ++index) {
+    if (associatedObjects[index].key == key) {
+      RETURN_RESULT_OF0(associatedObjects[index].getAssociatedObjectInstance);
+    }
+  }
+
+  RETURN_OBJ(nullptr);
+}
+
 bool IsSubInterface(const TypeInfo* thiz, const TypeInfo* other) {
   for (int i = 0; i < thiz->implementedInterfacesCount_; ++i) {
     if (thiz->implementedInterfaces_[i] == other) {

runtime/src/main/kotlin/kotlin/native/internal/KClassImpl.kt

@@ -40,8 +40,19 @@ internal class KClassImpl<T : Any>(private val typeInfo: NativePtr) : KClass<T>
     override fun toString(): String {
         return "class " + (qualifiedName ?: simpleName ?: "<anonymous>")
     }
+
+    internal fun findAssociatedObjectImpl(key: KClassImpl<*>): Any? =
+            findAssociatedObjectImpl(this.typeInfo, key.typeInfo)
 }
 
+@PublishedApi
+internal fun KClass<*>.findAssociatedObject(key: KClass<*>): Any? =
+        if (this is KClassImpl<*> && key is KClassImpl<*>) {
+            this.findAssociatedObjectImpl(key)
+        } else {
+            null
+        }
+
 @PublishedApi
 internal class KClassUnsupportedImpl(private val message: String) : KClass<Any> {
     override val simpleName: String? get() = error(message)
@@ -57,6 +68,9 @@ internal class KClassUnsupportedImpl(private val message: String) : KClass<Any>
     override fun toString(): String = "unreflected class ($message)"
 }
 
+@SymbolName("Kotlin_TypeInfo_findAssociatedObject")
+private external fun findAssociatedObjectImpl(typeInfo: NativePtr, key: NativePtr): Any?
+
 @ExportForCompiler
 @SymbolName("Kotlin_Any_getTypeInfo")
 internal external fun getObjectTypeInfo(obj: Any): NativePtr

runtime/src/main/kotlin/kotlin/reflect/AssociatedObjects.kt

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+package kotlin.reflect
+
+import kotlin.native.internal.*
+
+/**
+ * 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 [UseExperimental] annotation, e.g. `@UseExperimental(ExperimentalAssociatedObjects::class)`,
+ * or by using the compiler argument `-Xuse-experimental=kotlin.reflect.ExperimentalAssociatedObjects`.
+ */
+@Experimental(level = Experimental.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)