Ensure that non-local returns work properly in inline functions (Kotlin#3986)

Also, ensure that the workaround for KT-58685 still works.
It was enough to move the `return` a tiny bit to fix the issue.

Fixes Kotlin#3985

Author: dkhalanskyjb

kotlinx-coroutines-core/common/src/sync/Mutex.kt

@@ -121,19 +121,12 @@ public suspend inline fun <T> Mutex.withLock(owner: Any? = null, action: () -> T
     contract {
         callsInPlace(action, InvocationKind.EXACTLY_ONCE)
     }
-
-    // Cannot use 'finally' in this function because of KT-58685
-    // See kotlinx.coroutines.sync.MutexTest.testWithLockJsMiscompilation
-
     lock(owner)
-    val result = try {
+    return try {
         action()
-    } catch (e: Throwable) {
+    } finally {
         unlock(owner)
-        throw e
     }
-    unlock(owner)
-    return result
 }
 
 

kotlinx-coroutines-core/common/src/sync/Semaphore.kt

@@ -82,19 +82,12 @@ public suspend inline fun <T> Semaphore.withPermit(action: () -> T): T {
     contract {
         callsInPlace(action, InvocationKind.EXACTLY_ONCE)
     }
-
-    // Cannot use 'finally' in this function because of KT-58685
-    // See kotlinx.coroutines.sync.SemaphoreTest.testWithPermitJsMiscompilation
-
     acquire()
-    val result = try {
+    return try {
         action()
-    } catch (e: Throwable) {
+    } finally {
         release()
-        throw e
     }
-    release()
-    return result
 }
 
 @Suppress("UNCHECKED_CAST")

kotlinx-coroutines-core/common/test/channels/ConsumeTest.kt

@@ -0,0 +1,111 @@
+/*
+ * Copyright 2016-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+@file:OptIn(DelicateCoroutinesApi::class)
+package kotlinx.coroutines.channels
+
+import kotlinx.coroutines.*
+import kotlin.test.*
+
+class ConsumeTest: TestBase() {
+
+    /** Check that [ReceiveChannel.consume] does not suffer from KT-58685 */
+    @Test
+    fun testConsumeJsMiscompilation() = runTest {
+        val channel = Channel<Int>()
+        assertFailsWith<IndexOutOfBoundsException> {
+            try {
+                channel.consume { null } ?: throw IndexOutOfBoundsException() // should throw…
+            } catch (e: Exception) {
+                throw e // …but instead fails here
+            }
+        }
+    }
+
+    /** Checks that [ReceiveChannel.consume] closes the channel when the block executes successfully. */
+    @Test
+    fun testConsumeClosesOnSuccess() = runTest {
+        val channel = Channel<Int>()
+        channel.consume { }
+        assertTrue(channel.isClosedForReceive)
+    }
+
+    /** Checks that [ReceiveChannel.consume] closes the channel when the block executes successfully. */
+    @Test
+    fun testConsumeClosesOnFailure() = runTest {
+        val channel = Channel<Int>()
+        try {
+            channel.consume { throw TestException() }
+        } catch (e: TestException) {
+            // Expected
+        }
+        assertTrue(channel.isClosedForReceive)
+    }
+
+    /** Checks that [ReceiveChannel.consume] closes the channel when the block does an early return. */
+    @Test
+    fun testConsumeClosesOnEarlyReturn() = runTest {
+        val channel = Channel<Int>()
+        fun f() {
+            try {
+                channel.consume { return }
+            } catch (e: TestException) {
+                // Expected
+            }
+        }
+        f()
+        assertTrue(channel.isClosedForReceive)
+    }
+
+    /** Checks that [ReceiveChannel.consume] closes the channel when the block executes successfully. */
+    @Test
+    fun testConsumeEachClosesOnSuccess() = runTest {
+        val channel = Channel<Int>(Channel.UNLIMITED)
+        launch { channel.close() }
+        channel.consumeEach { fail("unreached") }
+        assertTrue(channel.isClosedForReceive)
+    }
+
+    /** Checks that [ReceiveChannel.consume] closes the channel when the block executes successfully. */
+    @Test
+    fun testConsumeEachClosesOnFailure() = runTest {
+        val channel = Channel<Unit>(Channel.UNLIMITED)
+        channel.send(Unit)
+        try {
+            channel.consumeEach { throw TestException() }
+        } catch (e: TestException) {
+            // Expected
+        }
+        assertTrue(channel.isClosedForReceive)
+    }
+
+    /** Checks that [ReceiveChannel.consume] closes the channel when the block does an early return. */
+    @Test
+    fun testConsumeEachClosesOnEarlyReturn() = runTest {
+        val channel = Channel<Unit>(Channel.UNLIMITED)
+        channel.send(Unit)
+        suspend fun f() {
+            channel.consumeEach {
+                return@f
+            }
+        }
+        f()
+        assertTrue(channel.isClosedForReceive)
+    }
+
+    /** Check that [BroadcastChannel.consume] does not suffer from KT-58685 */
+    @OptIn(ObsoleteCoroutinesApi::class)
+    @Suppress("DEPRECATION")
+    @Test
+    fun testBroadcastChannelConsumeJsMiscompilation() = runTest {
+        val channel = BroadcastChannel<Int>(1)
+        assertFailsWith<IndexOutOfBoundsException> {
+            try {
+                channel.consume { null } ?: throw IndexOutOfBoundsException() // should throw…
+            } catch (e: Exception) {
+                throw e // …but instead fails here
+            }
+        }
+    }
+}

kotlinx-coroutines-core/common/test/sync/MutexTest.kt

@@ -66,11 +66,27 @@ class MutexTest : TestBase() {
         try {
             mutex.withLock {
                 expect(1)
+                assertTrue(mutex.isLocked)
                 throw TestException()
             }
         } catch (e: TestException) {
-            finish(2)
+            expect(2)
+        }
+        assertFalse(mutex.isLocked)
+        finish(3)
+    }
+
+    @Test
+    fun withLockOnEarlyReturnTest() = runTest {
+        val mutex = Mutex()
+        assertFalse(mutex.isLocked)
+        suspend fun f() {
+            mutex.withLock {
+                assertTrue(mutex.isLocked)
+                return@f
+            }
         }
+        f()
         assertFalse(mutex.isLocked)
     }
 

kotlinx-coroutines-core/common/test/sync/SemaphoreTest.kt

@@ -69,6 +69,35 @@ class SemaphoreTest : TestBase() {
         assertEquals(1, semaphore.availablePermits)
     }
 
+    @Test
+    fun withSemaphoreOnFailureTest() = runTest {
+        val semaphore = Semaphore(1)
+        assertEquals(1, semaphore.availablePermits)
+        try {
+            semaphore.withPermit {
+                assertEquals(0, semaphore.availablePermits)
+                throw TestException()
+            }
+        } catch (e: TestException) {
+            // Expected
+        }
+        assertEquals(1, semaphore.availablePermits)
+    }
+
+    @Test
+    fun withSemaphoreOnEarlyReturnTest() = runTest {
+        val semaphore = Semaphore(1)
+        assertEquals(1, semaphore.availablePermits)
+        suspend fun f() {
+            semaphore.withPermit {
+                assertEquals(0, semaphore.availablePermits)
+                return@f
+            }
+        }
+        f()
+        assertEquals(1, semaphore.availablePermits)
+    }
+
     @Test
     fun fairnessTest() = runTest {
         val semaphore = Semaphore(1)