Use takeSubContOnce in inHandlingContext

kyay10 · Oct 13, 2024 · 86493cd · 86493cd
1 parent 3bae940
commit 86493cd
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/library/src/commonMain/kotlin/cloning.kt b/library/src/commonMain/kotlin/cloning.kt
@@ -171,7 +171,7 @@ internal fun <Start, First, End, P> SplitSeq<Start, First, End>.splitAt(p: Promp
 internal fun <Start, First, End, P> SplitSeq<Start, First, End>.splitAtOnce(p: Prompt<P>): Pair<Segment<Start, *, P>, SplitSeq<P, *, End>> {
   val box = findGuyBefore(p, null)
   return if (box != null) {
-    SingleUseSegment(box, this) to box.sequence!!.also { box.sequence = null } as SplitSeq<P, *, End>
+    SingleUseSegment(box, this) to (box.sequence!! as PromptCont).rest.also { box.sequence = null } as SplitSeq<P, *, End>
   } else {
     EmptySegment to this as SplitSeq<P, *, End>
   }

diff --git a/library/src/commonMain/kotlin/reset.kt b/library/src/commonMain/kotlin/reset.kt
@@ -1,4 +1,5 @@
 import arrow.core.identity
+import arrow.core.nonFatalOrThrow
 import kotlinx.coroutines.CancellationException
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
@@ -136,12 +137,17 @@ public suspend fun <T, R> Prompt<R>.takeSubContWithFinal(
 }
 
 // Acts like shift0/control { it(body()) }
-// TODO make it faster
+// TODO if `body` is multishot, we need to somehow
+//  evolve k to be multishot too. This is a more general issue with
+//  the `once` functionality. Maybe look to Scheme for ideas?
 @ResetDsl
 public suspend fun <T, P> Prompt<P>.inHandlingContext(
-  includeBodyContext: Boolean = false, body: suspend () -> T
-): T = takeSubCont(!includeBodyContext) { k ->
-  k.pushSubContWith(runCatching { body() }, isDelimiting = !includeBodyContext)
+  deleteDelimiter: Boolean = true, body: suspend () -> T
+): T = takeSubContOnce(deleteDelimiter) { k ->
+  val res = runCatching { body() }
+  // TODO test abortWith here
+  res.exceptionOrNull()?.nonFatalOrThrow()
+  k.pushSubContWith(res, isDelimiting = deleteDelimiter)
 }
 
 @Suppress("UNCHECKED_CAST")

diff --git a/library/src/commonTest/kotlin/ContinuationTest.kt b/library/src/commonTest/kotlin/ContinuationTest.kt
@@ -319,7 +319,7 @@ class ContinuationTest {
             shift0 { k -> k(ask()) } shouldBe 1
             ask() shouldBe 2
             inHandlingContext { ask() } shouldBe 1
-            inHandlingContext(includeBodyContext = true) { ask() } shouldBe 1
+            inHandlingContext(deleteDelimiter = true) { ask() } shouldBe 1
             ask() shouldBe 2
           }
         }
@@ -349,8 +349,8 @@ class ContinuationTest {
             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
+            inHandlingContext(deleteDelimiter = false) { inHandlingContext { inHandlingContext { ask() } } } shouldBe 4
+            inHandlingContext(deleteDelimiter = false) { inHandlingContext { inHandlingContext { inHandlingContext { ask() } } } } shouldBe 1
             ask() shouldBe 2
           }
         }