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 =
-        if (key === Key) EmptyCoroutineContext else this
+
+    public override fun minusKey(key: CoroutineContext.Key<*>): CoroutineContext {
+        // 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.io.Serializable
+import kotlin.coroutines.CoroutineContext.Element
+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.
  */