KT-36118 Introduce AbstractCoroutineContextKey and corresponding extensions required by its implementors

* It enables subtyping relationship for keys and elements, allowing type-safe and concise manipulations on the context when complex hierarchies of elements (e.g. implementors of ContinuationInterceptor) are present
    * It unblocks possibility to implement suspendCancellableCoroutine in stdlib
    * Implementing proper polymorphic get and minusKey in AbstractCoroutineContextElement is impossible because it will introduce source-incompatible change
    * ContinuationInterceptor get and minusKey are adjusted in a completely backwards-compatible way

libraries/stdlib/src/kotlin/coroutines/ContinuationInterceptor.kt

@@ -9,6 +9,12 @@ package kotlin.coroutines
  * Marks coroutine context element that intercepts coroutine continuations.
  * The coroutines framework uses [ContinuationInterceptor.Key] to retrieve the interceptor and
  * intercepts all coroutine continuations with [interceptContinuation] invocations.
+ *
+ * [ContinuationInterceptor] behaves like a [polymorphic element][AbstractCoroutineContextKey], meaning that
+ * its implementation delegates [get][CoroutineContext.Element.get] and [minusKey][CoroutineContext.Element.minusKey]
+ * to [getPolymorphicElement] and [minusPolymorphicKey] respectively.
+ * [ContinuationInterceptor] subtypes can be extracted from the coroutine context using either [ContinuationInterceptor.Key]
+ * or subtype key if it extends [AbstractCoroutineContextKey].
  */
 @SinceKotlin("1.3")
 public interface ContinuationInterceptor : CoroutineContext.Element {
@@ -43,12 +49,24 @@ public interface ContinuationInterceptor : CoroutineContext.Element {
         /* do nothing by default */
     }
 
-    // Performance optimization for a singleton Key
+    public override operator fun <E : CoroutineContext.Element> get(key: CoroutineContext.Key<E>): E? {
-    public override operator fun <E : CoroutineContext.Element> get(key: CoroutineContext.Key<E>): E? =
+        // getPolymorphicKey specialized for ContinuationInterceptor key
+        @OptIn(ExperimentalStdlibApi::class)
+        if (key is AbstractCoroutineContextKey<*, *>) {
+            @Suppress("UNCHECKED_CAST")
+            return if (key.isSubKey(this.key)) key.tryCast(this) as? E else null
+        }
         @Suppress("UNCHECKED_CAST")
-        if (key === Key) this as E else null
+        return if (ContinuationInterceptor === key) this as E else null
+    }
 
-    // Performance optimization to a singleton Key
+
-    public override fun minusKey(key: CoroutineContext.Key<*>): CoroutineContext =
+    public override fun minusKey(key: CoroutineContext.Key<*>): CoroutineContext {
-        if (key === Key) EmptyCoroutineContext else this
+        // minusPolymorphicKey specialized for ContinuationInterceptor key
+        @OptIn(ExperimentalStdlibApi::class)
+        if (key is AbstractCoroutineContextKey<*, *>) {
+            return if (key.isSubKey(this.key) && key.tryCast(this) != null) EmptyCoroutineContext else this
+        }
+        return if (ContinuationInterceptor === key) EmptyCoroutineContext else this
+    }
 }

libraries/stdlib/src/kotlin/coroutines/CoroutineContextImpl.kt

@@ -1,12 +1,12 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.coroutines
 
-import kotlin.coroutines.CoroutineContext.*
+import kotlin.coroutines.CoroutineContext.Element
-import kotlin.io.Serializable
+import kotlin.coroutines.CoroutineContext.Key
 
 /**
  * Base class for [CoroutineContext.Element] implementations.
@@ -14,6 +14,80 @@ import kotlin.io.Serializable
 @SinceKotlin("1.3")
 public abstract class AbstractCoroutineContextElement(public override val key: Key<*>) : Element
 
+/**
+ * Base class for [CoroutineContext.Key] associated with polymorphic [CoroutineContext.Element] implementation.
+ * Polymorphic element implementation implies delegating its [get][Element.get] and [minusKey][Element.minusKey]
+ * to [getPolymorphicElement] and [minusPolymorphicKey] respectively.
+ *
+ * Polymorphic elements can be extracted from the coroutine context using both element key and its supertype key.
+ * Example of polymorphic elements:
+ * ```
+ * open class BaseElement : CoroutineContext.Element {
+ *     companion object Key : CoroutineContext.Key<BaseElement>
+ *     override val key: CoroutineContext.Key<*> get() = Key
+ *     // It is important to use getPolymorphicKey and minusPolymorphicKey
+ *     override fun <E : CoroutineContext.Element> get(key: CoroutineContext.Key<E>): E? = getPolymorphicElement(key)
+ *     override fun minusKey(key: CoroutineContext.Key<*>): CoroutineContext = minusPolymorphicKey(key)
+ * }
+ *
+ * class DerivedElement : BaseElement() {
+ *     companion object Key : AbstractCoroutineContextKey<BaseElement, DerivedElement>(BaseElement, { it as? DerivedElement })
+ * }
+ * // Now it is possible to query both `BaseElement` and `DerivedElement`
+ * someContext[BaseElement] // Returns BaseElement?, non-null both for BaseElement and DerivedElement instances
+ * someContext[DerivedElement] // Returns DerivedElement?, non-null only for DerivedElement instance
+ * ```
+ * @param B base class of a polymorphic element
+ * @param baseKey an instance of base key
+ * @param E element type associated with the current key
+ * @param safeCast a function that can safely cast abstract [CoroutineContext.Element] to the concrete [E] type
+ *                 and return the element if it is a subtype of [E] or `null` otherwise.
+ */
+@SinceKotlin("1.3")
+@ExperimentalStdlibApi
+public abstract class AbstractCoroutineContextKey<B : Element, E : B>(
+    baseKey: Key<B>,
+    private val safeCast: (element: Element) -> E?
+) : Key<E> {
+    private val topmostKey: Key<*> = if (baseKey is AbstractCoroutineContextKey<*, *>) baseKey.topmostKey else baseKey
+
+    internal fun tryCast(element: Element): E? = safeCast(element)
+    internal fun isSubKey(key: Key<*>): Boolean = key === this || topmostKey === key
+}
+
+/**
+ * Returns the current element is it is associated with the given [key] in a polymorphic manner or `null` otherwise.
+ * This method returns non-null value if either [Element.key] is equal to the given [key] or if the [key] is associated
+ * with [Element.key] via [AbstractCoroutineContextKey].
+ * See [AbstractCoroutineContextKey] for the example of usage.
+ */
+@SinceKotlin("1.3")
+@ExperimentalStdlibApi
+public fun <E : Element> Element.getPolymorphicElement(key: Key<E>): E? {
+    if (key is AbstractCoroutineContextKey<*, *>) {
+        @Suppress("UNCHECKED_CAST")
+        return if (key.isSubKey(this.key)) key.tryCast(this) as? E else null
+    }
+    @Suppress("UNCHECKED_CAST")
+    return if (this.key === key) this as E else null
+}
+
+/**
+ * Returns empty coroutine context if the element is associated with the given [key] in a polymorphic manner
+ * or `null` otherwise.
+ * This method returns empty context if either [Element.key] is equal to the given [key] or if the [key] is associated
+ * with [Element.key] via [AbstractCoroutineContextKey].
+ * See [AbstractCoroutineContextKey] for the example of usage.
+ */
+@SinceKotlin("1.3")
+@ExperimentalStdlibApi
+public fun Element.minusPolymorphicKey(key: Key<*>): CoroutineContext {
+    if (key is AbstractCoroutineContextKey<*, *>) {
+        return if (key.isSubKey(this.key) && key.tryCast(this) != null) EmptyCoroutineContext else this
+    }
+    return if (this.key === key) EmptyCoroutineContext else this
+}
+
 /**
  * An empty coroutine context.
  */

libraries/stdlib/test/coroutines/AbstractCoroutineContextElementTest.kt

@@ -0,0 +1,167 @@
+/*
+ * 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 test.coroutines
+
+import kotlin.coroutines.*
+import kotlin.test.*
+
+class AbstractCoroutineContextElementTest {
+
+    private val CoroutineContext.size get() = fold(0) { acc, _ -> acc + 1 }
+
+    abstract class Base : AbstractCoroutineContextElement(Key) {
+        companion object Key : CoroutineContext.Key<Base>
+
+        override fun <E : CoroutineContext.Element> get(key: CoroutineContext.Key<E>): E? = getPolymorphicElement(key)
+        override fun minusKey(key: CoroutineContext.Key<*>): CoroutineContext = minusPolymorphicKey(key)
+    }
+
+    class DerivedWithoutKey : Base() {
+        // No custom key
+    }
+
+    open class DerivedWithKey : Base() {
+        companion object Key : AbstractCoroutineContextKey<Base, DerivedWithKey>(Base, { it as? DerivedWithKey })
+    }
+
+    class SubDerivedWithKey : DerivedWithKey() {
+        companion object Key : AbstractCoroutineContextKey<Base, SubDerivedWithKey>(Base, { it as? SubDerivedWithKey })
+    }
+
+    class SubDerivedWithKeyAndDifferentBase : DerivedWithKey() {
+        // Note how different base class is used
+        companion object Key :
+            AbstractCoroutineContextKey<DerivedWithKey, SubDerivedWithKeyAndDifferentBase>(
+                DerivedWithKey,
+                { it as? SubDerivedWithKeyAndDifferentBase })
+    }
+
+    object IrrelevantElement : AbstractCoroutineContextElement(Key) {
+        object Key : CoroutineContext.Key<IrrelevantElement>
+    }
+
+    @Test
+    fun testDerivedWithoutKey() {
+        val derivedWithoutKey = DerivedWithoutKey()
+        assertSame(Base.Key, derivedWithoutKey.key)
+        testDerivedWithoutKey(EmptyCoroutineContext, derivedWithoutKey) // Single element
+        testDerivedWithoutKey(IrrelevantElement, derivedWithoutKey) // Combined context
+    }
+
+    @Test
+    fun testDerivedWithoutKeyOverridesDerived() {
+        val context = DerivedWithKey() + DerivedWithoutKey()
+        assertEquals(1, context.size)
+        assertTrue(context[Base] is DerivedWithoutKey)
+        assertNull(context[DerivedWithKey])
+        assertEquals(EmptyCoroutineContext, context.minusKey(Base))
+        assertSame(context, context.minusKey(DerivedWithKey))
+    }
+
+    private fun testDerivedWithoutKey(context: CoroutineContext, element: CoroutineContext.Element) {
+        val ctx = context + element
+        assertEquals(context.size + 1, ctx.size)
+        assertSame(element, ctx[Base]!!)
+        assertNull(ctx[DerivedWithKey])
+        assertEquals(context, ctx.minusKey(Base))
+        assertSame(ctx, ctx.minusKey(DerivedWithKey))
+    }
+
+    @Test
+    fun testDerivedWithKey() {
+        val derivedWithKey = DerivedWithKey()
+        assertSame(Base.Key, derivedWithKey.key)
+        testDerivedWithKey(EmptyCoroutineContext, derivedWithKey) // Single element
+        testDerivedWithKey(IrrelevantElement, derivedWithKey) // Combined context
+    }
+
+    private fun testDerivedWithKey(context: CoroutineContext, element: CoroutineContext.Element) {
+        val ctx = context + element
+        assertEquals(context.size + 1, ctx.size)
+        assertSame(element, ctx[Base]!!)
+        assertSame(element, ctx[DerivedWithKey]!!)
+        assertEquals(context, ctx.minusKey(Base))
+        assertEquals(context, ctx.minusKey(DerivedWithKey))
+    }
+
+    @Test
+    fun testSubDerivedWithKey() {
+        val subDerivedWithKey = SubDerivedWithKey()
+        assertSame(Base.Key, subDerivedWithKey.key)
+        testSubDerivedWithKey(EmptyCoroutineContext, subDerivedWithKey)
+        testSubDerivedWithKey(IrrelevantElement, subDerivedWithKey)
+    }
+
+    private fun testSubDerivedWithKey(context: CoroutineContext, element: CoroutineContext.Element) {
+        val ctx = context + element
+        assertEquals(context.size + 1, ctx.size)
+        assertSame(element, ctx[Base]!!)
+        assertSame(element, ctx[DerivedWithKey]!!)
+        assertSame(element, ctx[SubDerivedWithKey]!!)
+        assertNull(ctx[SubDerivedWithKeyAndDifferentBase])
+        assertEquals(context, ctx.minusKey(Base))
+        assertEquals(context, ctx.minusKey(DerivedWithKey))
+        assertEquals(context, ctx.minusKey(SubDerivedWithKey))
+        assertSame(ctx, ctx.minusKey(SubDerivedWithKeyAndDifferentBase))
+    }
+
+    @Test
+    fun testSubDerivedWithKeyAndDifferentBase() {
+        val subDerivedWithKeyAndDifferentBase = SubDerivedWithKeyAndDifferentBase()
+        assertSame(Base.Key, subDerivedWithKeyAndDifferentBase.key)
+        testSubDerivedWithKeyAndDifferentBase(EmptyCoroutineContext, subDerivedWithKeyAndDifferentBase)
+        testSubDerivedWithKeyAndDifferentBase(IrrelevantElement, subDerivedWithKeyAndDifferentBase)
+    }
+
+    private fun testSubDerivedWithKeyAndDifferentBase(context: CoroutineContext, element: CoroutineContext.Element) {
+        val ctx = context + element
+        assertEquals(context.size + 1, ctx.size)
+        assertSame(element, ctx[Base]!!)
+        assertSame(element, ctx[DerivedWithKey]!!)
+        assertSame(element, ctx[SubDerivedWithKeyAndDifferentBase]!!)
+        assertNull(ctx[SubDerivedWithKey])
+        assertEquals(context, ctx.minusKey(Base))
+        assertEquals(context, ctx.minusKey(DerivedWithKey))
+        assertEquals(context, ctx.minusKey(SubDerivedWithKeyAndDifferentBase))
+        assertSame(ctx, ctx.minusKey(SubDerivedWithKey))
+    }
+
+    @Test
+    fun testDerivedWithKeyOverridesDerived() {
+        val context = DerivedWithoutKey() + DerivedWithKey()
+        assertEquals(1, context.size)
+        assertTrue { context[Base] is DerivedWithKey }
+        assertTrue { context[DerivedWithKey] is DerivedWithKey }
+        assertEquals(EmptyCoroutineContext, context.minusKey(Base))
+        assertEquals(EmptyCoroutineContext, context.minusKey(DerivedWithKey))
+    }
+
+    @Test
+    fun testSubDerivedOverrides() {
+        val key = SubDerivedWithKeyAndDifferentBase
+        testSubDerivedOverrides<SubDerivedWithKeyAndDifferentBase>(DerivedWithoutKey() + SubDerivedWithKeyAndDifferentBase(), key)
+        testSubDerivedOverrides<SubDerivedWithKeyAndDifferentBase>(DerivedWithKey() + SubDerivedWithKeyAndDifferentBase(), key)
+        testSubDerivedOverrides<SubDerivedWithKeyAndDifferentBase>(SubDerivedWithKeyAndDifferentBase() + SubDerivedWithKeyAndDifferentBase(), key)
+    }
+
+    @Test
+    fun testSubDerivedWithDifferentBaseOverrides() {
+        val key = SubDerivedWithKey
+        testSubDerivedOverrides<SubDerivedWithKey>(DerivedWithoutKey() + SubDerivedWithKey(), key)
+        testSubDerivedOverrides<SubDerivedWithKey>(DerivedWithKey() + SubDerivedWithKey(), key)
+        testSubDerivedOverrides<SubDerivedWithKey>(SubDerivedWithKeyAndDifferentBase() + SubDerivedWithKey(), key)
+    }
+
+    private inline fun <reified T : CoroutineContext.Element> testSubDerivedOverrides(context: CoroutineContext, key: CoroutineContext.Key<T>) {
+        assertEquals(1, context.size)
+        assertTrue { context[Base] is DerivedWithKey }
+        assertTrue { context[DerivedWithKey] is DerivedWithKey }
+        assertTrue { context[DerivedWithKey] is T }
+        assertTrue { context[key] is T }
+        assertEquals(EmptyCoroutineContext, context.minusKey(Base))
+        assertEquals(EmptyCoroutineContext, context.minusKey(DerivedWithKey))
+    }
+}

libraries/stdlib/test/coroutines/ContinuationInterceptorKeyTest.kt

@@ -0,0 +1,132 @@
+/*
+ * 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 test.coroutines
+
+import kotlin.coroutines.*
+import kotlin.test.*
+
+class ContinuationInterceptorKeyTest {
+
+    private val CoroutineContext.size get() = fold(0) { acc, _ -> acc + 1 }
+
+    abstract class BaseElement : AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor {
+        companion object Key :
+            AbstractCoroutineContextKey<ContinuationInterceptor, BaseElement>(ContinuationInterceptor, { it as? BaseElement })
+
+        override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> = continuation
+    }
+
+    // "Legacy" code, BaseElement with ContinuationInterceptor key
+    class DerivedElementWithOldKey : BaseElement() {
+        override val key: CoroutineContext.Key<*>
+            get() = ContinuationInterceptor
+    }
+
+    // "New" code with AbstractCoroutineContextKey
+    class DerivedElementWithPolyKey : BaseElement() {
+        companion object Key :
+            AbstractCoroutineContextKey<BaseElement, DerivedElementWithPolyKey>(BaseElement, { it as? DerivedElementWithPolyKey })
+    }
+
+    // Irrelevant interceptor
+    class CustomInterceptor : ContinuationInterceptor {
+        override val key: CoroutineContext.Key<*>
+            get() = ContinuationInterceptor
+
+        override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> = continuation
+    }
+
+    object IrrelevantElement : AbstractCoroutineContextElement(Key) {
+        object Key : CoroutineContext.Key<IrrelevantElement>
+    }
+
+    @Test
+    fun testKeyIsNotOverridden() {
+        val derivedElementWithOldKey = DerivedElementWithOldKey()
+        testKeyIsNotOverridden(derivedElementWithOldKey, derivedElementWithOldKey)
+        testKeyIsNotOverridden(IrrelevantElement + derivedElementWithOldKey, derivedElementWithOldKey) // test for CombinedContext
+    }
+
+    private fun testKeyIsNotOverridden(context: CoroutineContext, element: CoroutineContext.Element) {
+        run {
+            val interceptor: ContinuationInterceptor = context[ContinuationInterceptor]!!
+            assertSame(element, interceptor)
+        }
+        run {
+            val baseElement: BaseElement = context[BaseElement]!!
+            assertSame(element, baseElement)
+        }
+
+        val subtracted = context.minusKey(ContinuationInterceptor)
+        assertEquals(subtracted, context.minusKey(BaseElement))
+        assertNull(subtracted[ContinuationInterceptor])
+        assertNull(subtracted[BaseElement])
+        assertEquals(context.size - 1, subtracted.size)
+    }
+
+    @Test
+    fun testKeyIsOverridden() {
+        val derivedElementWithPolyKey = DerivedElementWithPolyKey()
+        testKeyIsOverridden(derivedElementWithPolyKey, derivedElementWithPolyKey)
+        testKeyIsOverridden(IrrelevantElement + derivedElementWithPolyKey, derivedElementWithPolyKey) // test for CombinedContext
+    }
+
+    private fun testKeyIsOverridden(context: CoroutineContext, element: CoroutineContext.Element) {
+        testKeyIsNotOverridden(context, element)
+        val derived = context[DerivedElementWithPolyKey]
+        assertNotNull(derived)
+        assertSame(element, derived)
+        val subtracted = context.minusKey(ContinuationInterceptor)
+        assertEquals(subtracted, context.minusKey(BaseElement))
+        assertEquals(subtracted, context.minusKey(DerivedElementWithPolyKey))
+        assertEquals(context.size - 1, subtracted.size)
+    }
+
+    @Test
+    fun testInterceptorKeyIsNotOverridden() {
+        val ci = CustomInterceptor()
+        testInterceptorKeyIsNotOverridden(ci, ci)
+        testInterceptorKeyIsNotOverridden(IrrelevantElement + ci, ci) // test for CombinedContext
+    }
+
+    private fun testInterceptorKeyIsNotOverridden(context: CoroutineContext, element: CoroutineContext.Element) {
+        val interceptor = context[ContinuationInterceptor]
+        assertNotNull(interceptor)
+        assertSame(element, interceptor)
+        assertNull(context[BaseElement])
+        assertNull(context[DerivedElementWithPolyKey])
+        assertEquals(context, context.minusKey(BaseElement))
+        assertEquals(context, context.minusKey(DerivedElementWithPolyKey))
+    }
+
+    @Test
+    fun testContextOperations() {
+        val interceptor = CustomInterceptor()
+        val derivedWithOld = DerivedElementWithOldKey()
+        val derivedWithPoly = DerivedElementWithPolyKey()
+        val e = IrrelevantElement
+        run {
+            assertEquals(interceptor, derivedWithOld + derivedWithPoly + interceptor)
+            assertEquals(interceptor + e, derivedWithOld + derivedWithPoly + interceptor + e)
+            assertEquals(interceptor, derivedWithOld + interceptor)
+            assertEquals(interceptor, derivedWithPoly + interceptor)
+        }
+
+        run {
+            assertEquals(derivedWithOld, derivedWithPoly + interceptor + derivedWithOld)
+            assertEquals(derivedWithOld + e, derivedWithPoly + interceptor + derivedWithOld + e)
+            assertEquals(derivedWithOld, derivedWithPoly + derivedWithOld)
+            assertEquals(derivedWithOld, interceptor + derivedWithOld)
+        }
+
+        run {
+            assertEquals(derivedWithPoly, interceptor + derivedWithOld + derivedWithPoly)
+            assertEquals(derivedWithPoly + e, interceptor + derivedWithOld + derivedWithPoly + e)
+            assertEquals(derivedWithPoly, derivedWithOld + derivedWithPoly)
+            assertEquals(derivedWithPoly, interceptor + derivedWithPoly)
+        }
+    }
+}

libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt

@@ -2683,6 +2683,10 @@ public abstract class kotlin/coroutines/AbstractCoroutineContextElement : kotlin
 	public fun plus (Lkotlin/coroutines/CoroutineContext;)Lkotlin/coroutines/CoroutineContext;
 }
 
+public abstract class kotlin/coroutines/AbstractCoroutineContextKey : kotlin/coroutines/CoroutineContext$Key {
+	public fun <init> (Lkotlin/coroutines/CoroutineContext$Key;Lkotlin/jvm/functions/Function1;)V
+}
+
 public abstract interface class kotlin/coroutines/Continuation {
 	public abstract fun getContext ()Lkotlin/coroutines/CoroutineContext;
 	public abstract fun resumeWith (Ljava/lang/Object;)V
@@ -2742,6 +2746,11 @@ public final class kotlin/coroutines/CoroutineContext$Element$DefaultImpls {
 public abstract interface class kotlin/coroutines/CoroutineContext$Key {
 }
 
+public final class kotlin/coroutines/CoroutineContextImplKt {
+	public static final fun getPolymorphicElement (Lkotlin/coroutines/CoroutineContext$Element;Lkotlin/coroutines/CoroutineContext$Key;)Lkotlin/coroutines/CoroutineContext$Element;
+	public static final fun minusPolymorphicKey (Lkotlin/coroutines/CoroutineContext$Element;Lkotlin/coroutines/CoroutineContext$Key;)Lkotlin/coroutines/CoroutineContext;
+}
+
 public final class kotlin/coroutines/EmptyCoroutineContext : java/io/Serializable, kotlin/coroutines/CoroutineContext {
 	public static final field INSTANCE Lkotlin/coroutines/EmptyCoroutineContext;
 	public fun fold (Ljava/lang/Object;Lkotlin/jvm/functions/Function2;)Ljava/lang/Object;