Add facility to check number of suspensions in coroutine tests

compiler/testData/codegen/boxInline/suspend/stateMachine/numberOfSuspentions.kt

@@ -21,53 +21,20 @@ import COROUTINES_PACKAGE.*
 import COROUTINES_PACKAGE.intrinsics.*
 import helpers.*
 
-var result = "FAIL"
-var i = 0
-var finished = false
-
-var proceed: () -> Unit = {}
-
-suspend fun suspendHere() = suspendCoroutine<Unit> { c ->
-    i++
-    proceed = { c.resume(Unit) }
-}
-
 fun builder(c: suspend () -> Unit) {
-    val continuation = object: ContinuationAdapter<Unit>() {
+    c.startCoroutine(CheckStateMachineContinuation)
-        override val context: CoroutineContext
-            get() = EmptyCoroutineContext
-
-        override fun resume(value: Unit) {
-            proceed = {
-                result = "OK"
-                finished = true
-            }
-        }
-
-        override fun resumeWithException(exception: Throwable) {
-            throw exception
-        }
-    }
-    c.startCoroutine(continuation)
 }
 
 fun box(): String {
     builder {
         crossinlineMe2 {
-            suspendHere()
+            StateMachineChecker.suspendHere()
-            suspendHere()
+            StateMachineChecker.suspendHere()
-            suspendHere()
+            StateMachineChecker.suspendHere()
-            suspendHere()
+            StateMachineChecker.suspendHere()
-            suspendHere()
+            StateMachineChecker.suspendHere()
         }
     }
-    for (counter in 0 until 10) {
+    StateMachineChecker.check(numberOfSuspensions = 10)
-        if (i != counter + 1) return "Expected ${counter + 1}, got $i"
+    return "OK"
-        proceed()
-    }
-    if (i != 10) return "FAIL $i"
-    if (finished) return "resume on root continuation is called"
-    proceed()
-    if (!finished) return "resume on root continuation is not called"
-    return result
 }

compiler/tests-common/tests/org/jetbrains/kotlin/coroutineTestUtil.kt

@@ -76,6 +76,46 @@ fun createTextForHelpers(isReleaseCoroutines: Boolean): String {
         else
             ""
 
+    val checkStateMachine = """
+    object StateMachineChecker {
+        private var counter = 0
+        var finished = false
+
+        var proceed: () -> Unit = {}
+
+        suspend fun suspendHere() = suspendCoroutine<Unit> { c ->
+            counter++
+            proceed = { c.resume(Unit) }
+        }
+
+        fun check(numberOfSuspensions: Int) {
+            for (i in 1..numberOfSuspensions) {
+                if (counter != i) error("Wrong state-machine generated: suspendHere called should be called exactly once in one state. Expected " + i + ", got " + counter)
+                proceed()
+            }
+            if (counter != numberOfSuspensions)
+                error("Wrong state-machine generated: suspendHere called should be called exactly once in one state. Expected " + numberOfSuspensions + ", got " + counter)
+            if (finished) error("Wrong state-machine generated: it is finished early")
+            proceed()
+            if (!finished) error("Wrong state-machine generated: it is not finished yet")
+        }
+    }
+    object CheckStateMachineContinuation: ContinuationAdapter<Unit>() {
+        override val context: CoroutineContext
+            get() = EmptyCoroutineContext
+
+        override fun resume(value: Unit) {
+            StateMachineChecker.proceed = {
+                StateMachineChecker.finished = true
+            }
+        }
+
+        override fun resumeWithException(exception: Throwable) {
+            throw exception
+        }
+    }
+    """.trimIndent()
+
     return """
             |package helpers
             |import $coroutinesPackage.*
@@ -100,5 +140,7 @@ fun createTextForHelpers(isReleaseCoroutines: Boolean): String {
             |    override val context: CoroutineContext = EmptyCoroutineContext
             |    $continuationAdapterBody
             |}
+            |
+            |$checkStateMachine
         """.trimMargin()
 }