refactor(moonbit): simplify async bindings

Author: peter-jerry-ye

crates/moonbit/README.md

@@ -30,7 +30,7 @@ Stream + Sink (batched reads/writes):
 let (s, sink) = @async.Stream::new[Byte]()
 @async.spawn(async fn() {
   let chunk : Array[Byte] = [1, 2, 3, 4]
-  let _ = sink.write(chunk[:])
+  assert_true(sink.write_all(chunk[:]))
   sink.close()
 })
 let chunk = s.read(4096)
@@ -43,8 +43,9 @@ match chunk {
 ```
 
 `Stream::read(count)` returns up to `count` elements; `Sink::write` accepts
-`ArrayView[T]` so byte streams can batch data efficiently. `Stream::new`
-accepts an optional `capacity` (<= 0 means unbounded).
+`ArrayView[T]` so byte streams can batch data efficiently, and
+`Sink::write_all` loops until the whole slice is accepted or the sink closes.
+`Stream::new` accepts an optional `capacity` (<= 0 means unbounded).
 
 ## Testing
 

crates/moonbit/src/async/coroutine.mbt

@@ -50,7 +50,7 @@ fn Coroutine::wake(self : Coroutine) -> Unit {
 ///|
 pub fn is_being_cancelled() -> Bool {
   let coro = current_coroutine()
-  coro.cancelled && not(coro.shielded)
+  coro.cancelled && !coro.shielded
 }
 
 ///|
@@ -59,15 +59,15 @@ pub(all) suberror Cancelled derive(Show)
 ///|
 fn Coroutine::cancel(self : Coroutine) -> Unit {
   self.cancelled = true
-  if not(self.shielded || self.ready) {
+  if !(self.shielded || self.ready) {
     self.wake()
   }
 }
 
 ///|
 pub async fn pause() -> Unit raise Cancelled {
   guard scheduler.curr_coro is Some(coro)
-  if coro.cancelled && not(coro.shielded) {
+  if coro.cancelled && !coro.shielded {
     raise Cancelled::Cancelled
   }
   async_suspend(fn(ok_cont, err_cont) {
@@ -81,7 +81,7 @@ pub async fn pause() -> Unit raise Cancelled {
 ///|
 pub async fn suspend() -> Unit raise Cancelled {
   guard scheduler.curr_coro is Some(coro)
-  if coro.cancelled && not(coro.shielded) {
+  if coro.cancelled && !coro.shielded {
     raise Cancelled::Cancelled
   }
   scheduler.blocking += 1
@@ -147,7 +147,7 @@ fn Coroutine::unwrap(self : Coroutine) -> Unit raise {
 ///|
 async fn Coroutine::wait(target : Coroutine) -> Unit {
   guard scheduler.curr_coro is Some(coro)
-  guard not(physical_equal(coro, target))
+  guard !physical_equal(coro, target)
   match target.state {
     Done => return
     Fail(err) => raise err

crates/moonbit/src/async/ev.mbt

@@ -16,6 +16,7 @@
 priv struct EventLoop {
   subscribes : Map[Int, Subscriber]
   tasks : Map[WaitableSet, Coroutine]
+  owned_coroutines : Map[WaitableSet, @set.Set[Coroutine]]
   finished : Map[WaitableSet, Bool]
 }
 
@@ -39,6 +40,7 @@ fn should_complete(waitable_set : WaitableSet) -> Bool {
 fn next_callback(waitable_set : WaitableSet) -> Int {
   if should_complete(waitable_set) {
     ev.tasks.remove(waitable_set)
+    ev.owned_coroutines.remove(waitable_set)
     ev.finished.remove(waitable_set)
     waitable_set.drop()
     return CallbackCode::Completed.encode()
@@ -52,12 +54,17 @@ fn next_callback(waitable_set : WaitableSet) -> Int {
 ///|
 pub fn with_waitableset(f : async () -> Unit) -> Int {
   let waitable_set = WaitableSet::new()
+  let owned = @set.Set::new()
   tls_set(waitable_set.0)
   let coro = spawn(async fn() -> Unit {
+    let coro = current_coroutine()
+    defer owned.remove(coro)
     defer ev.finished.set(waitable_set, true)
     f()
   })
   ev.tasks.set(waitable_set, coro)
+  owned.add(coro)
+  ev.owned_coroutines.set(waitable_set, owned)
   ev.finished.set(waitable_set, false)
   reschedule()
   next_callback(waitable_set)
@@ -75,13 +82,17 @@ pub fn cb(event : Int, waitable_id : Int, code : Int) -> Int {
     TaskCancelled => {
       guard ev.tasks.get(waitable_set) is Some(coro)
       coro.cancel()
-      for {
+      if ev.owned_coroutines.get(waitable_set) is Some(owned) {
+        owned.each(Coroutine::cancel)
+      }
+      for ;; {
         reschedule()
         if !has_immediately_ready_task() {
           break
         }
       }
       ev.tasks.remove(waitable_set)
+      ev.owned_coroutines.remove(waitable_set)
       ev.finished.remove(waitable_set)
       if ev.subscribes.is_empty() {
         waitable_set.drop()
@@ -104,7 +115,32 @@ pub fn cb(event : Int, waitable_id : Int, code : Int) -> Int {
 }
 
 ///|
-let ev : EventLoop = { subscribes: {}, tasks: {}, finished: {} }
+let ev : EventLoop = {
+  subscribes: {},
+  tasks: {},
+  owned_coroutines: {},
+  finished: {},
+}
+
+///|
+/// Spawn a coroutine owned by the current component-model async task.
+///
+/// This is intended for runtime bridge work such as lowering local futures and
+/// streams to component-model handles. The coroutine does not participate in
+/// `TaskGroup` structured concurrency, but it is cancelled when the current
+/// component-model task is cancelled and keeps the waitable-set alive until it
+/// terminates.
+pub fn spawn_component_task_current(f : async () -> Unit) -> Unit {
+  let waitable_set = current_waitableset()
+  guard ev.tasks.get(waitable_set) is Some(_)
+  guard ev.owned_coroutines.get(waitable_set) is Some(owned)
+  let coro = spawn(async fn() -> Unit {
+    let coro = current_coroutine()
+    defer owned.remove(coro)
+    f()
+  })
+  owned.add(coro)
+}
 
 ///|
 pub fn detach_waitable(waitable_id : Int) -> Unit {
@@ -125,7 +161,7 @@ pub async fn suspend_for_subtask(
 
   // Helper: ensure cleanup is called once we've moved past Starting state
   fn ensure_cleanup(state : SubTaskState) -> Unit {
-    if not(cleaned) && !(state is Starting) {
+    if !cleaned && !(state is Starting) {
       cleanup_after_started()
       cleaned = true
     }
@@ -159,7 +195,7 @@ pub async fn suspend_for_subtask(
   let set = @set.Set::new()
   let subscriber = { event: None, coro: set }
   ev.subscribes.set(task.handle, subscriber)
-  for {
+  for ;; {
     set.add(current_coroutine())
     waitable_join(task.handle, current_waitableset().0)
     defer subscriber.coro.remove(current_coroutine())

crates/moonbit/src/async/moon.pkg.json

@@ -5,5 +5,5 @@
     { "path": "moonbitlang/core/ref", "alias": "ref" },
     { "path": "moonbitlang/core/set", "alias": "set" }
   ],
-  "supported-targets": ["wasm"]
+  "supported-targets": "+wasm"
 }

crates/moonbit/src/async/scheduler.mbt

@@ -54,7 +54,7 @@ pub fn reschedule() -> Unit {
     coro.state = Running
     let last_coro = scheduler.curr_coro
     scheduler.curr_coro = Some(coro)
-    if coro.cancelled && not(coro.shielded) {
+    if coro.cancelled && !coro.shielded {
       err_cont(Cancelled::Cancelled)
     } else {
       ok_cont(())

crates/moonbit/src/async/task_group.mbt

@@ -64,14 +64,14 @@ fn[X] TaskGroup::spawn_coroutine(
   guard self.state is Running else {
     abort("trying to spawn from a terminated task group")
   }
-  if not(no_wait) {
+  if !no_wait {
     self.waiting += 1
   }
   async fn worker() {
     let coro = current_coroutine()
     defer {
       self.children.remove(coro)
-      if not(no_wait) {
+      if !no_wait {
         self.waiting -= 1
         if self.waiting == 0 && self.state is Running {
           for child in self.children {
@@ -93,7 +93,7 @@ fn[X] TaskGroup::spawn_coroutine(
             child.cancel()
           }
           self.state = Fail(err)
-        } else if not(err is Cancelled::Cancelled) {
+        } else if !(err is Cancelled::Cancelled) {
           self.state = Fail(err)
         }
         raise err
@@ -212,7 +212,7 @@ pub async fn[X] with_task_group(f : async (TaskGroup[X]) -> X) -> X {
       tg.result = Some(value)
     }
   })
-  if not(tg.children.is_empty()) {
+  if !tg.children.is_empty() {
     suspend() catch {
       err =>
         if tg.state is Running {
@@ -223,7 +223,7 @@ pub async fn[X] with_task_group(f : async (TaskGroup[X]) -> X) -> X {
         }
     }
   }
-  if not(tg.children.is_empty()) {
+  if !tg.children.is_empty() {
     protect_from_cancel(() => suspend()) catch {
       _ => ()
     }