Remove unit from handlers

Add handleStateful
kyay10 · Jun 29, 2024 · 90e5a5d · 90e5a5d
1 parent 02c57f5
commit 90e5a5d
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 67 deletions.
diff --git a/library/src/commonMain/kotlin/effekt/handler.kt b/library/src/commonMain/kotlin/effekt/handler.kt
@@ -2,31 +2,46 @@ package effekt
 
 import Prompt
 import abortS0
+import pushPrompt
 import reset
 import shift0
+import takeSubCont
+import kotlin.coroutines.CoroutineContext
 import kotlin.jvm.JvmInline
 
-public interface Handler<R, E> {
-  public fun prompt(): ObscurePrompt<E>
-  public suspend fun unit(value: R): E
-  public suspend fun <A> use(body: suspend (Cont<A, E>) -> E): A = prompt().prompt.shift0(body)
-  public fun <A> useAbort(body: suspend () -> E): A = prompt().prompt.abortS0(body)
+public interface Handler<E> {
+  public fun prompt(): HandlerPrompt<E>
 }
 
-public suspend fun <R, E, H : Handler<R, E>> handle(
-  handler: ((() -> ObscurePrompt<E>) -> H), body: suspend H.() -> R
-): E {
-  val prompt = ObscurePrompt(Prompt<E>())
-  val handler = handler { prompt }
-  return handler.handle(body)
-}
-
-public suspend fun <R, E, H : Handler<R, E>> H.handle(body: suspend H.() -> R): E {
-  return prompt().prompt.reset {
-    val res = body()
-    prompt().prompt.abortS0 { unit(res) }
+public suspend fun <A, E> Handler<E>.use(body: suspend (Cont<A, E>) -> E): A = prompt().prompt.shift0(body)
+public fun <A, E> Handler<E>.discard(body: suspend () -> E): A = prompt().prompt.abortS0(body)
+public suspend fun <A, E> Handler<E>.useStateful(body: suspend (suspend (A, CoroutineContext) -> E) -> E): A =
+  prompt().prompt.takeSubCont { sk ->
+    body { a, ctx -> sk.pushSubContWith(Result.success(a), isDelimiting = true, extraContext = ctx) }
   }
-}
+
+public suspend fun <E, H> handle(
+  handler: ((() -> HandlerPrompt<E>) -> H), body: suspend H.() -> E
+): E = handle { handler { this }.body() }
+
+public suspend fun <E> handle(body: suspend HandlerPrompt<E>.() -> E): E = HandlerPrompt<E>().handle(body)
+
+public suspend fun <E> HandlerPrompt<E>.handle(body: suspend HandlerPrompt<E>.() -> E): E = prompt.reset { body() }
+
+public suspend fun <E, H> handleStateful(
+  handler: ((() -> HandlerPrompt<E>) -> H), extraContext: CoroutineContext, body: suspend H.() -> E
+): E = handleStateful(extraContext) { handler { this }.body() }
+
+public suspend fun <E> handleStateful(extraContext: CoroutineContext, body: suspend HandlerPrompt<E>.() -> E): E =
+  HandlerPrompt<E>().handleStateful(extraContext, body)
+
+public suspend fun <E> HandlerPrompt<E>.handleStateful(
+  extraContext: CoroutineContext, body: suspend HandlerPrompt<E>.() -> E
+): E = prompt.pushPrompt(extraContext) { body() }
 
 @JvmInline
-public value class ObscurePrompt<E> internal constructor(internal val prompt: Prompt<E>)
+public value class HandlerPrompt<E> internal constructor(internal val prompt: Prompt<E>) : Handler<E> {
+  public constructor() : this(Prompt())
+
+  override fun prompt(): HandlerPrompt<E> = this
+}
diff --git a/library/src/commonTest/kotlin/effekt/HandlerTest.kt b/library/src/commonTest/kotlin/effekt/HandlerTest.kt
@@ -1,12 +1,16 @@
 package effekt
 
+import Reader
 import arrow.core.None
 import arrow.core.Option
 import arrow.core.Some
+import ask
+import context
 import io.kotest.matchers.shouldBe
 import kotlinx.coroutines.test.runTest
 import runCC
 import kotlin.collections.plus
+import kotlin.coroutines.CoroutineContext
 import kotlin.test.Test
 
 class HandlerTest {
@@ -128,40 +132,32 @@ interface Amb {
   suspend fun flip(): Boolean
 }
 
-interface Choose {
-  suspend fun <A> choose(first: A, second: A): A
-}
-
 private suspend fun drunkFlip(amb: Amb, exc: Exc): String {
   val caught = amb.flip()
   val heads = if (caught) amb.flip() else exc.raise("We dropped the coin.")
   return if (heads) "Heads" else "Tails"
 }
 
-fun interface Maybe<R> : Exc, Handler<R, Option<R>> {
-  override suspend fun unit(value: R): Option<R> = Some(value)
-  override suspend fun raise(msg: String): Nothing = useAbort { None }
+fun interface Maybe<R> : Exc, Handler<Option<R>> {
+  override suspend fun raise(msg: String): Nothing = discard { None }
 }
 
-suspend fun <R> maybe(block: suspend Exc.() -> R): Option<R> = handle(::Maybe, block)
+suspend fun <R> maybe(block: suspend Exc.() -> R): Option<R> = handle(::Maybe) {
+  Some(block())
+}
 
-fun interface Collect<R> : Amb, Handler<R, List<R>> {
-  override suspend fun unit(value: R): List<R> = listOf(value)
+fun interface Collect<R> : Amb, Handler<List<R>> {
   override suspend fun flip(): Boolean = use { resume ->
     val ts = resume(true)
     val fs = resume(false)
     ts + fs
   }
 }
 
-suspend fun <R> collect(block: suspend Amb.() -> R): List<R> = handle(::Collect, block)
-
-fun interface CollectChoose<R> : Collect<R>, Choose {
-  override suspend fun <A> choose(first: A, second: A): A = if (flip()) first else second
+suspend fun <R> collect(block: suspend Amb.() -> R): List<R> = handle(::Collect) {
+  listOf(block())
 }
 
-suspend fun <R> collectChoose(block: suspend CollectChoose<R>.() -> R): List<R> = handle(::CollectChoose, block)
-
 interface Fiber {
   suspend fun suspend()
   suspend fun fork(): Boolean
@@ -212,52 +208,55 @@ suspend inline fun poll(state: StateScope, fiber: Fiber, block: suspend Async.()
 
 interface NonDetermined : Amb, Exc
 
-interface Backtrack<R> : Amb, Handler<R, Option<R>> {
+interface Backtrack<R> : Amb, Handler<Option<R>> {
   override suspend fun flip(): Boolean = use { resume ->
     resume(true).onNone { return@use resume(false) }
   }
 }
 
 fun interface FirstResult<R> : NonDetermined, Maybe<R>, Backtrack<R>
 
-suspend fun <R> firstResult(block: suspend NonDetermined.() -> R) = handle(::FirstResult, block)
+suspend fun <R> firstResult(block: suspend NonDetermined.() -> R): Option<R> = handle(::FirstResult) {
+  Some(block())
+}
 
-interface Scheduler<R> : Fiber, Handler<R, Unit> {
-  val queue: StateScope.Field<Queue>
-  override suspend fun unit(value: R) {}
-  override suspend fun exit(): Nothing = useAbort {}
-  override suspend fun fork(): Boolean = use { resume ->
-    queue.update { listOf(suspend { resume(true) }, suspend { resume(false) }) + it }
-    run()
+interface Scheduler : Fiber, Handler<Unit> {
+  val queue: Reader<Queue>
+  override suspend fun exit(): Nothing = discard {}
+  override suspend fun fork(): Boolean {
+    val q = queue.ask()
+    return useStateful { resume ->
+      run(listOf<suspend (CoroutineContext) -> Unit>({ resume(true, it) }, { resume(false, it) }) + q)
+    }
   }
 
-  override suspend fun suspend() = use { resume ->
-    queue.update { it + suspend { resume(Unit) } }
-    run()
+  override suspend fun suspend() {
+    val q = queue.ask()
+    useStateful { resume ->
+      run(q + { resume(Unit, it) })
+    }
   }
 }
 
-suspend fun StateScope.scheduler(block: suspend Fiber.() -> Unit) {
-  val queue = field(emptyList<suspend () -> Unit>())
-  handle<Unit, Unit, Scheduler<Unit>>({
-    object : Scheduler<Unit> {
+suspend fun scheduler(block: suspend Fiber.() -> Unit) {
+  val queue = Reader<Queue>()
+  handleStateful<Unit, Fiber>({
+    object : Scheduler {
       override val queue = queue
-      override fun prompt(): ObscurePrompt<Unit> = it()
+      override fun prompt(): HandlerPrompt<Unit> = it()
     }
-  }, block)
+  }, queue.context(emptyList()), block)
 }
 
-private tailrec suspend fun Scheduler<*>.run() {
-  val q = queue.get()
+private suspend fun Scheduler.run(q: Queue) {
   if (q.isNotEmpty()) {
     val p = q.first()
-    queue.set(q.drop(1))
-    p()
-    run()
+    val newQ = q.drop(1)
+    p(queue.context(newQ))
   }
 }
 
-typealias Queue = List<suspend () -> Unit>
+typealias Queue = List<suspend (CoroutineContext) -> Unit>
 
 interface Input {
   suspend fun read(): Char

diff --git a/library/src/jvmTest/kotlin/effekt/HandlerJvmTest.kt b/library/src/jvmTest/kotlin/effekt/HandlerJvmTest.kt
@@ -3,6 +3,7 @@ package effekt
 import arrow.core.None
 import arrow.core.Option
 import arrow.core.Some
+import arrow.core.recover
 import io.kotest.matchers.shouldBe
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -120,18 +121,22 @@ suspend fun <R> stringInput(input: String, block: suspend StringInput<R>.() -> R
 }
 
 fun interface Nondet<R> : Collect<R>, NonDetermined {
-  override suspend fun raise(msg: String): Nothing = useAbort { emptyList() }
+  override suspend fun raise(msg: String): Nothing = discard { emptyList() }
 }
 
-suspend fun <R> nondet(block: suspend NonDetermined.() -> R): List<R> = handle(::Nondet, block)
+suspend fun <R> nondet(block: suspend NonDetermined.() -> R): List<R> = handle(::Nondet) {
+  listOf(block())
+}
 
 fun interface Backtrack2<R> : Maybe<R>, NonDetermined {
   override suspend fun flip(): Boolean = use { resume ->
-    resume(true).onNone { return@use resume(false) }
+    resume(true).recover { resume(false).bind() }
   }
 }
 
-suspend fun <R> backtrack(block: suspend NonDetermined.() -> R): Option<R> = handle(::Backtrack2, block)
+suspend fun <R> backtrack(block: suspend NonDetermined.() -> R): Option<R> = handle(::Backtrack2) {
+  Some(block())
+}
 
 interface Generator<A> {
   suspend fun yield(value: A)
@@ -141,9 +146,8 @@ private suspend fun Generator<Int>.numbers(upto: Int) {
   for (i in 0..upto) yield(i)
 }
 
-interface Iterate<A> : Generator<A>, Handler<Unit, EffectfulIterator<A>> {
+interface Iterate<A> : Generator<A>, Handler<EffectfulIterator<A>> {
   val nextValue: StateScope.Field<(suspend (Unit) -> A)?>
-  override suspend fun unit(value: Unit): EffectfulIterator<A> = EffectfulIteratorImpl(nextValue)
   override suspend fun yield(value: A) = use { resume ->
     nextValue.set {
       resume(Unit)
@@ -164,9 +168,17 @@ suspend fun <A> StateScope.iterate(block: suspend Generator<A>.() -> Unit): Effe
   return handle({
     object : Iterate<A> {
       override val nextValue = nextValue
-      override fun prompt(): ObscurePrompt<EffectfulIterator<A>> = it()
+      override fun prompt(): HandlerPrompt<EffectfulIterator<A>> = it()
     }
-  }, block)
+  }) {
+    block()
+    EmptyEffectfulIterator
+  }
+}
+
+private object EmptyEffectfulIterator : EffectfulIterator<Nothing> {
+  override suspend fun hasNext(): Boolean = false
+  override suspend fun next(): Nothing = throw NoSuchElementException()
 }
 
 interface EffectfulIterator<out A> {