Add callsInPlace to ReceiveChannel.consume (#2496)

Fixes #941

2 files changed, 23 insertions, 4 deletions

diff --git a/kotlinx-coroutines-core/common/src/channels/Channels.common.kt b/kotlinx-coroutines-core/common/src/channels/Channels.common.kt
index 60948f62..e3567e31 100644
--- a/kotlinx-coroutines-core/common/src/channels/Channels.common.kt
+++ b/kotlinx-coroutines-core/common/src/channels/Channels.common.kt
@@ -1,14 +1,16 @@
 /*
- * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 @file:JvmMultifileClass
 @file:JvmName("ChannelsKt")
 @file:Suppress("DEPRECATION_ERROR")
+@file:OptIn(ExperimentalContracts::class)
 
 package kotlinx.coroutines.channels
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.selects.*
+import kotlin.contracts.*
 import kotlin.coroutines.*
 import kotlin.jvm.*
 
@@ -164,6 +166,9 @@ public fun consumesAll(vararg channels: ReceiveChannel<*>): CompletionHandler =
  * The operation is _terminal_.
  */
 public inline fun <E, R> ReceiveChannel<E>.consume(block: ReceiveChannel<E>.() -> R): R {
+    contract {
+        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+    }
     var cause: Throwable? = null
     try {
         return block()
diff --git a/kotlinx-coroutines-core/common/test/BuilderContractsTest.kt b/kotlinx-coroutines-core/common/test/BuilderContractsTest.kt
index b20dd6b1..5a96c544 100644
--- a/kotlinx-coroutines-core/common/test/BuilderContractsTest.kt
+++ b/kotlinx-coroutines-core/common/test/BuilderContractsTest.kt
@@ -1,9 +1,10 @@
 /*
- * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.selects.*
 import kotlin.test.*
 
@@ -44,9 +45,22 @@ class BuilderContractsTest : TestBase() {
             Job().apply { complete() }.onJoin {}
         }
         consume(s)
+
+
+        val ch: Int
+        val i = Channel<Int>()
+        i.consume {
+            ch = 321
+        }
+        consume(ch)
     }
 
     private fun consume(a: Int) {
-        a.hashCode() // BE codegen verification
+        /*
+         * Verify the value is actually set correctly
+         * (non-zero, VerificationError is not triggered, can be read)
+         */
+        assertNotEquals(0, a)
+        assertEquals(a.hashCode(), a)
     }
-}
\ No newline at end of file
+}