Add inHandlingContext and tests for it

Fix debug message in Continuation.forEach Fix logic in compilerCloning for JVM when running in debugger
kyay10 · Jun 8, 2024 · fdcb28e · fdcb28e
1 parent f013aba
commit fdcb28e
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 3 deletions.
diff --git a/library/src/commonMain/kotlin/cloning.kt b/library/src/commonMain/kotlin/cloning.kt
@@ -19,7 +19,7 @@ internal inline fun Continuation<*>.forEach(block: (Continuation<*>) -> Unit) {
     current = when (current) {
       in CompilerGenerated -> current.completion
       is CopyableContinuation -> current.completion
-      else -> error("Continuation $this is not see-through, so its stack can't be traversed")
+      else -> error("Continuation $current is not see-through, so its stack can't be traversed")
     }
   }
 }

diff --git a/library/src/commonMain/kotlin/reset.kt b/library/src/commonMain/kotlin/reset.kt
@@ -114,6 +114,20 @@ public suspend fun <T, R> Prompt<R>.takeSubCont(
   body.startCoroutine(SubCont(subchain), if(deleteDelimiter) hole.completion else hole)
 }
 
+// Acts like shift0/shift { it(body()) }
+// This is NOT multishot
+@ResetDsl
+public suspend fun <T> Prompt<*>.inHandlingContext(
+  includeBodyContext: Boolean = false, body: suspend () -> T
+): T = suspendCoroutineUnintercepted { k ->
+  val hole = k.context.holeFor(this, !includeBodyContext)
+  body.startCoroutine(Continuation(hole.context) {
+    val exception = it.exceptionOrNull()
+    if (exception is SeekingCoroutineContextException) exception.use(hole.context)
+    else k.resumeWith(it)
+  })
+}
+
 @Suppress("UNCHECKED_CAST")
 internal fun <R> Prompt<R>.abortWith(deleteDelimiter: Boolean, value: Result<R>): Nothing =
   throw AbortWithValueException(this as Prompt<Any?>, value, deleteDelimiter)

diff --git a/library/src/commonTest/kotlin/ContinuationTest.kt b/library/src/commonTest/kotlin/ContinuationTest.kt
@@ -306,6 +306,56 @@ class ContinuationTest {
       x + 1
     } shouldBe 4
   }
+
+  @Test
+  fun handlingContext() = runTest {
+    runCC {
+      runReader(1) {
+        newReset<Int> {
+          pushReader(2) {
+            ask() shouldBe 2
+            shift { k -> k(ask()) } shouldBe 1
+            shift0 { k -> k(ask()) } shouldBe 1
+            ask() shouldBe 2
+            inHandlingContext { ask() } shouldBe 1
+            inHandlingContext(includeBodyContext = true) { ask() } shouldBe 1
+            ask() shouldBe 2
+          }
+        }
+      }
+    }
+    runCC {
+      runReader(1) {
+        newReset<Int> {
+          pushReader(2) {
+            ask() shouldBe 2
+            shift { k ->
+              val v = ask()
+              pushReader(3) { k(v) }
+            } shouldBe 1
+            ask() shouldBe 2
+            shift0 { k -> k(ask()) } shouldBe 3
+            shift0 { k -> shift0 { k(ask()) } } shouldBe 1
+            ask() shouldBe 2
+            shift { k ->
+              val v = ask()
+              pushReader(4) { k(v) }
+            } shouldBe 1
+            ask() shouldBe 2
+            // Introduces an extra delimiter
+            shift { k -> k(ask()) } shouldBe 4
+            ask() shouldBe 2
+            inHandlingContext { ask() } shouldBe 4
+            inHandlingContext { inHandlingContext { ask() } } shouldBe 4
+            inHandlingContext { inHandlingContext { inHandlingContext { ask() } } } shouldBe 1
+            inHandlingContext(includeBodyContext = true) { inHandlingContext { inHandlingContext { ask() } } } shouldBe 4
+            inHandlingContext(includeBodyContext = true) { inHandlingContext { inHandlingContext { inHandlingContext { ask() } } } } shouldBe 1
+            ask() shouldBe 2
+          }
+        }
+      }
+    }
+  }
 }
 
 val stackSafeIterations: Int = when (platform) {

diff --git a/library/src/jvmMain/kotlin/compilerCloning.kt b/library/src/jvmMain/kotlin/compilerCloning.kt
@@ -47,7 +47,6 @@ internal actual fun <T> Continuation<T>.copy(completion: Continuation<*>): Conti
   }
 
   else -> {
-    val delegate = delegateField.get(this) as Continuation<*>
-    coroutineOwnerConstructor.newInstance(delegate.copy(completion), infoField.get(this)) as Continuation<T>
+    coroutineOwnerConstructor.newInstance(completion, infoField.get(this)) as Continuation<T>
   }
 }