From 5895c211a1449e5b803ea8a110103ca566335188 Mon Sep 17 00:00:00 2001
From: Roman Elizarov <elizarov@gmail.com>
Date: Mon, 6 Feb 2017 13:06:39 +0300
Subject: [PATCH] Fixed corotine context performance by using reference
 equality for contexts and keys Added documentation explaining that keys are
 compared by reference

---
 .../coroutines/experimental/CoroutineContext.kt  |  8 ++++++--
 .../experimental/CoroutineContextImpl.kt         | 16 ++++++++--------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/libraries/stdlib/src/kotlin/coroutines/experimental/CoroutineContext.kt b/libraries/stdlib/src/kotlin/coroutines/experimental/CoroutineContext.kt
index dab29bee133..6d21ed81a78 100644
--- a/libraries/stdlib/src/kotlin/coroutines/experimental/CoroutineContext.kt
+++ b/libraries/stdlib/src/kotlin/coroutines/experimental/CoroutineContext.kt
@@ -19,12 +19,14 @@ package kotlin.coroutines.experimental
 /**
  * Persistent context for the coroutine. It is an indexed set of [Element] instances.
  * An indexed set is a mix between a set and a map.
- * Every element in this set has a unique [Key].
+ * Every element in this set has a unique [Key]. Keys are compared _by reference_.
  */
 @SinceKotlin("1.1")
 public interface CoroutineContext {
     /**
      * Returns the element with the given [key] from this context or `null`.
+     * Keys are compared _by reference_, that is to get an element from the context the reference to its actual key
+     * object must be presented to this function.
      */
     public operator fun <E : Element> get(key: Key<E>): E?
 
@@ -42,7 +44,8 @@ public interface CoroutineContext {
 
     /**
      * Returns a context containing elements from this context, but without an element with
-     * the specified [key].
+     * the specified [key]. Keys are compared _by reference_, that is to remove an element from the context
+     * the reference to its actual key object must be presented to this function.
      */
     public fun minusKey(key: Key<*>): CoroutineContext
 
@@ -58,6 +61,7 @@ public interface CoroutineContext {
 
     /**
      * Key for the elements of [CoroutineContext]. [E] is a type of element with this key.
+     * Keys in the context are compared _by reference_.
      */
     public interface Key<E : Element>
 }
diff --git a/libraries/stdlib/src/kotlin/coroutines/experimental/CoroutineContextImpl.kt b/libraries/stdlib/src/kotlin/coroutines/experimental/CoroutineContextImpl.kt
index ed46d89a033..749cfb0edd2 100644
--- a/libraries/stdlib/src/kotlin/coroutines/experimental/CoroutineContextImpl.kt
+++ b/libraries/stdlib/src/kotlin/coroutines/experimental/CoroutineContextImpl.kt
@@ -25,7 +25,7 @@ import kotlin.coroutines.experimental.CoroutineContext.*
 public abstract class AbstractCoroutineContextElement(public override val key: Key<*>) : Element {
     @Suppress("UNCHECKED_CAST")
     public override operator fun <E : Element> get(key: Key<E>): E? =
-            if (this.key == key) this as E else null
+            if (this.key === key) this as E else null
 
     public override fun <R> fold(initial: R, operation: (R, Element) -> R): R =
             operation(initial, this)
@@ -34,7 +34,7 @@ public abstract class AbstractCoroutineContextElement(public override val key: K
             plusImpl(context)
 
     public override fun minusKey(key: Key<*>): CoroutineContext =
-            if (this.key == key) EmptyCoroutineContext else this
+            if (this.key === key) EmptyCoroutineContext else this
 }
 
 /**
@@ -77,9 +77,9 @@ private class CombinedContext(val left: CoroutineContext, val element: Element)
     public override fun minusKey(key: Key<*>): CoroutineContext {
         element[key]?.let { return left }
         val newLeft = left.minusKey(key)
-        return when (newLeft) {
-            left -> this
-            EmptyCoroutineContext -> element
+        return when {
+            newLeft === left -> this
+            newLeft === EmptyCoroutineContext -> element
             else -> CombinedContext(newLeft, element)
         }
     }
@@ -115,15 +115,15 @@ private class CombinedContext(val left: CoroutineContext, val element: Element)
 }
 
 private fun CoroutineContext.plusImpl(context: CoroutineContext): CoroutineContext =
-        if (context == EmptyCoroutineContext) this else // fast path -- avoid lambda creation
+        if (context === EmptyCoroutineContext) this else // fast path -- avoid lambda creation
             context.fold(this) { acc, element ->
                 val removed = acc.minusKey(element.key)
-                if (removed == EmptyCoroutineContext) element else {
+                if (removed === EmptyCoroutineContext) element else {
                     // make sure interceptor is always last in the context (and thus is fast to get when present)
                     val interceptor = removed[ContinuationInterceptor]
                     if (interceptor == null) CombinedContext(removed, element) else {
                         val left = removed.minusKey(ContinuationInterceptor)
-                        if (left == EmptyCoroutineContext) CombinedContext(element, interceptor) else
+                        if (left === EmptyCoroutineContext) CombinedContext(element, interceptor) else
                             CombinedContext(CombinedContext(left, element), interceptor)
                     }
                 }