Leverage polymorphic keys un CoroutineDispatcher and ExecutorCoroutin…

…eDispatcher

kotlinx-coroutines-core/api/kotlinx-coroutines-core.api

@@ -151,6 +151,7 @@ public final class kotlinx/coroutines/CoroutineContextKt {
 }
 
 public abstract class kotlinx/coroutines/CoroutineDispatcher : kotlin/coroutines/AbstractCoroutineContextElement, kotlin/coroutines/ContinuationInterceptor {
+	public static final field Key Lkotlinx/coroutines/CoroutineDispatcher$Key;
 	public fun <init> ()V
 	public abstract fun dispatch (Lkotlin/coroutines/CoroutineContext;Ljava/lang/Runnable;)V
 	public fun dispatchYield (Lkotlin/coroutines/CoroutineContext;Ljava/lang/Runnable;)V
@@ -163,6 +164,9 @@ public abstract class kotlinx/coroutines/CoroutineDispatcher : kotlin/coroutines
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class kotlinx/coroutines/CoroutineDispatcher$Key : kotlin/coroutines/AbstractCoroutineContextKey {
+}
+
 public abstract interface class kotlinx/coroutines/CoroutineExceptionHandler : kotlin/coroutines/CoroutineContext$Element {
 	public static final field Key Lkotlinx/coroutines/CoroutineExceptionHandler$Key;
 	public abstract fun handleException (Lkotlin/coroutines/CoroutineContext;Ljava/lang/Throwable;)V
@@ -294,11 +298,15 @@ public final class kotlinx/coroutines/ExceptionsKt {
 }
 
 public abstract class kotlinx/coroutines/ExecutorCoroutineDispatcher : kotlinx/coroutines/CoroutineDispatcher, java/io/Closeable {
+	public static final field Key Lkotlinx/coroutines/ExecutorCoroutineDispatcher$Key;
 	public fun <init> ()V
 	public abstract fun close ()V
 	public abstract fun getExecutor ()Ljava/util/concurrent/Executor;
 }
 
+public final class kotlinx/coroutines/ExecutorCoroutineDispatcher$Key : kotlin/coroutines/AbstractCoroutineContextKey {
+}
+
 public final class kotlinx/coroutines/ExecutorsKt {
 	public static final fun asExecutor (Lkotlinx/coroutines/CoroutineDispatcher;)Ljava/util/concurrent/Executor;
 	public static final fun from (Ljava/util/concurrent/Executor;)Lkotlinx/coroutines/CoroutineDispatcher;

kotlinx-coroutines-core/common/src/CoroutineDispatcher.kt

@@ -30,6 +30,12 @@ import kotlin.coroutines.*
 public abstract class CoroutineDispatcher :
     AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor {
 
+    /** @suppress */
+    @ExperimentalStdlibApi
+    public companion object Key : AbstractCoroutineContextKey<ContinuationInterceptor, CoroutineDispatcher>(
+        ContinuationInterceptor,
+        { it as? CoroutineDispatcher })
+
     /**
      * Returns `true` if the execution of the coroutine should be performed with [dispatch] method.
      * The default behavior for most dispatchers is to return `true`.

kotlinx-coroutines-core/jvm/src/Executors.kt

@@ -5,7 +5,7 @@
 package kotlinx.coroutines
 
 import kotlinx.coroutines.internal.*
-import java.io.Closeable
+import java.io.*
 import java.util.concurrent.*
 import kotlin.coroutines.*
 
@@ -17,17 +17,23 @@ import kotlin.coroutines.*
  * asynchronous API which requires instance of the [Executor].
  */
 public abstract class ExecutorCoroutineDispatcher: CoroutineDispatcher(), Closeable {
+    /** @suppress */
+    @ExperimentalStdlibApi
+    public companion object Key : AbstractCoroutineContextKey<CoroutineDispatcher, ExecutorCoroutineDispatcher>(
+        CoroutineDispatcher,
+        { it as? ExecutorCoroutineDispatcher })
+
+    /**
+     * Underlying executor of current [CoroutineDispatcher].
+     */
+    public abstract val executor: Executor
+
     /**
      * Closes this coroutine dispatcher and shuts down its executor.
      *
      * It may throw an exception if this dispatcher is global and cannot be closed.
      */
     public abstract override fun close()
-
-    /**
-     * Underlying executor of current [CoroutineDispatcher].
-     */
-    public abstract val executor: Executor
 }
 
 /**

kotlinx-coroutines-core/jvm/test/DispatcherKeyTest.kt

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import org.junit.Test
+import kotlin.coroutines.*
+import kotlin.test.*
+
+@UseExperimental(ExperimentalStdlibApi::class)
+class DispatcherKeyTest : TestBase() {
+
+    companion object CustomInterceptor : AbstractCoroutineContextElement(ContinuationInterceptor),
+        ContinuationInterceptor {
+        override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> {
+            return continuation
+        }
+    }
+
+    private val name = CoroutineName("test")
+
+    @Test
+    fun testDispatcher() {
+        val context = name + CustomInterceptor
+        assertNull(context[CoroutineDispatcher])
+        assertSame(CustomInterceptor, context[ContinuationInterceptor])
+
+        val updated = context + Dispatchers.Main
+        val result: CoroutineDispatcher? = updated[CoroutineDispatcher]
+        assertSame(Dispatchers.Main, result)
+        assertSame(Dispatchers.Main, updated[ContinuationInterceptor])
+        assertEquals(name, updated.minusKey(CoroutineDispatcher))
+        assertEquals(name, updated.minusKey(ContinuationInterceptor))
+    }
+
+    @Test
+    fun testExecutorCoroutineDispatcher() {
+        val context = name + CustomInterceptor
+        assertNull(context[ExecutorCoroutineDispatcher])
+        val updated = context + Dispatchers.Main
+        assertNull(updated[ExecutorCoroutineDispatcher])
+        val executor = Dispatchers.Default
+        val updated2 = updated + executor
+        assertSame(Dispatchers.Default, updated2[ContinuationInterceptor])
+        assertSame(Dispatchers.Default, updated2[CoroutineDispatcher])
+        assertSame(Dispatchers.Default as ExecutorCoroutineDispatcher, updated2[ExecutorCoroutineDispatcher])
+        assertEquals(name, updated2.minusKey(ContinuationInterceptor))
+        assertEquals(name, updated2.minusKey(CoroutineDispatcher))
+        assertEquals(name, updated2.minusKey(ExecutorCoroutineDispatcher))
+    }
+}