From 0401b4e0bb3c50e986461824454328f42beb55ef Mon Sep 17 00:00:00 2001 From: Roman Elizarov Date: Sun, 26 Aug 2018 23:01:21 +0300 Subject: [PATCH] Coroutine context shall perform structural equality comparison on keys Fixes KT-26398 --- .../jvm/test/CoroutineContextTest.kt | 100 ++++++++++++++++++ .../coroutines/ContinuationInterceptor.kt | 9 ++ .../src/kotlin/coroutines/CoroutineContext.kt | 6 +- .../kotlin-stdlib-runtime-merged.txt | 2 + 4 files changed, 114 insertions(+), 3 deletions(-) create mode 100644 libraries/stdlib/coroutines/jvm/test/CoroutineContextTest.kt diff --git a/libraries/stdlib/coroutines/jvm/test/CoroutineContextTest.kt b/libraries/stdlib/coroutines/jvm/test/CoroutineContextTest.kt new file mode 100644 index 00000000000..52f276aba2b --- /dev/null +++ b/libraries/stdlib/coroutines/jvm/test/CoroutineContextTest.kt @@ -0,0 +1,100 @@ +/* + * Copyright 2010-2018 JetBrains s.r.o. 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.kotlin.coroutines + +import kotlin.coroutines.* +import kotlin.test.* + +class CoroutineContextTest { + @Test + fun testBasicOperations() { + checkContents(EmptyCoroutineContext) + val de1 = DataElement(1) + val de2 = DataElement(2) + checkContents(de1, de1) + checkContents(de2, de2) + checkContents(de1 + de2, de2) + checkContents(de2 + de1, de1) + val oe3 = OtherElement(3) + val oe4 = OtherElement(4) + val de1oe3 = de1 + oe3 + checkContents(de1oe3, de1, oe3) + checkContents(de1oe3 + de2, de2, oe3) + checkContents(de1oe3 + oe4, de1, oe4) + checkContents(de1oe3.minusKey(DataElement), oe3) + checkContents(de1oe3.minusKey(OtherElement), de1) + } + + @Test + fun testWrapperEquality() { + val we1 = WrapperElement("1") + val we2 = WrapperElement("2") + checkContents(we1, we1) + checkContents(we2, we2) + val we1we2 = we1 + we2 + checkContents(we1we2, we1, we2) + checkContents(we1we2.minusKey(WrapperKey("1")), we2) + checkContents(we1we2.minusKey(WrapperKey("2")), we1) + } + + @Test + fun testInterceptor() { + val de1 = DataElement(1) + val oe2 = OtherElement(2) + val we3 = WrapperElement("3") + val cci = CustomContinuationInterceptor() + // Make sure it works properly with any position of ContinuationInterceptor in the context + checkContentsAndRemoves(de1 + oe2 + we3 + cci, de1, oe2, we3, cci) + checkContentsAndRemoves(de1 + oe2 + cci + we3, de1, oe2, we3, cci) + checkContentsAndRemoves(de1 + cci + oe2 + we3, de1, oe2, we3, cci) + checkContentsAndRemoves(cci + de1 + oe2 + we3, de1, oe2, we3, cci) + } + + private fun checkContentsAndRemoves(context: CoroutineContext, vararg es: CoroutineContext.Element) { + checkContents(context, *es) + for (e in es) { + checkContents(context.minusKey(e.key), *(es.toSet() - e).toTypedArray()) + } + } + + private fun checkContents(context: CoroutineContext, vararg es: CoroutineContext.Element) { + val size = context.fold(0) { c, _ -> c + 1} + assertEquals(es.size, size) + for (e in es) { + val key = e.key + assertSame(e, context[key]) + } + val set = mutableSetOf() + context.fold(set) { s, e -> s.apply { add(e) } } + assertEquals(set, es.toSet()) + when (es.size) { + 0 -> assertSame(context, EmptyCoroutineContext) + 1 -> assertSame(context, es[0]) + } + } + + class DataElement(val data: Int) : AbstractCoroutineContextElement(Key) { + companion object Key : CoroutineContext.Key + } + + class OtherElement(val data: Int) : AbstractCoroutineContextElement(Key) { + companion object Key : CoroutineContext.Key + } + + data class WrapperKey(val key: String) : CoroutineContext.Key + + class WrapperElement(key: String) : CoroutineContext.Element { + override val key = WrapperKey(key) + } + + class CustomContinuationInterceptor : ContinuationInterceptor { + override val key: CoroutineContext.Key<*> + get() = ContinuationInterceptor + + override fun interceptContinuation(continuation: Continuation): Continuation = + continuation + } +} \ No newline at end of file diff --git a/libraries/stdlib/coroutines/src/kotlin/coroutines/ContinuationInterceptor.kt b/libraries/stdlib/coroutines/src/kotlin/coroutines/ContinuationInterceptor.kt index 05d7f88db29..6a824802a04 100644 --- a/libraries/stdlib/coroutines/src/kotlin/coroutines/ContinuationInterceptor.kt +++ b/libraries/stdlib/coroutines/src/kotlin/coroutines/ContinuationInterceptor.kt @@ -42,4 +42,13 @@ public interface ContinuationInterceptor : CoroutineContext.Element { public fun releaseInterceptedContinuation(continuation: Continuation<*>) { /* do nothing by default */ } + + // Performance optimization for a singleton Key + public override operator fun get(key: CoroutineContext.Key): E? = + @Suppress("UNCHECKED_CAST") + if (key === 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 } diff --git a/libraries/stdlib/coroutines/src/kotlin/coroutines/CoroutineContext.kt b/libraries/stdlib/coroutines/src/kotlin/coroutines/CoroutineContext.kt index 615a7f85d02..b76163e0f3f 100644 --- a/libraries/stdlib/coroutines/src/kotlin/coroutines/CoroutineContext.kt +++ b/libraries/stdlib/coroutines/src/kotlin/coroutines/CoroutineContext.kt @@ -66,14 +66,14 @@ public interface CoroutineContext { */ public val key: Key<*> - @Suppress("UNCHECKED_CAST") public override operator fun get(key: Key): E? = - if (this.key === key) this as E else null + @Suppress("UNCHECKED_CAST") + if (this.key == key) this as E else null public override fun fold(initial: R, operation: (R, Element) -> R): R = operation(initial, this) public override fun minusKey(key: Key<*>): CoroutineContext = - if (this.key === key) EmptyCoroutineContext else this + if (this.key == key) EmptyCoroutineContext else this } } diff --git a/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt b/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt index a892ba12946..74b7261b9ad 100644 --- a/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt +++ b/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt @@ -2633,7 +2633,9 @@ public abstract interface class kotlin/coroutines/Continuation { public abstract interface class kotlin/coroutines/ContinuationInterceptor : kotlin/coroutines/CoroutineContext$Element { public static final field Key Lkotlin/coroutines/ContinuationInterceptor$Key; + public abstract fun get (Lkotlin/coroutines/CoroutineContext$Key;)Lkotlin/coroutines/CoroutineContext$Element; public abstract fun interceptContinuation (Lkotlin/coroutines/Continuation;)Lkotlin/coroutines/Continuation; + public abstract fun minusKey (Lkotlin/coroutines/CoroutineContext$Key;)Lkotlin/coroutines/CoroutineContext; public abstract fun releaseInterceptedContinuation (Lkotlin/coroutines/Continuation;)V }