Merge branch 'sofia/sync-broadcast' of https://github.com/leanprover/lean4 into sofia/sync-broadcast

Author: algebraic-dev

src/Std/Sync/Broadcast.lean

@@ -11,15 +11,9 @@ public import Init.System.Promise
 public import Init.Data.Queue
 public import Std.Sync.Mutex
 public import Std.Internal.Async.Select
-public import Std.Internal.Async.IO
 
 public section
 
-namespace Std
-
-open Std.Internal.Async.IO
-open Std.Internal.IO.Async
-
 /-!
 The `Std.Sync.Broadcast` module implements a broadcasting primitive for sending values
 to multiple consumers. It maintains a queue of values and supports both synchronous
@@ -29,6 +23,8 @@ This module is heavily inspired by `Std.Sync.Channel` as well as
 [tokio’s broadcast implementation](https://github.com/tokio-rs/tokio/blob/master/tokio/src/sync/broadcast.rs).
 -/
 
+namespace Std
+
 /--
 Errors that may be thrown while interacting with the broadcast channel API.
 -/
@@ -333,18 +329,11 @@ private def unsubscribe (bd : Bounded.Receiver α) : IO Unit := do
     let st ← get
 
     let some next := st.receivers.get? bd.id
-      | return Except.error Broadcast.Error.notSubscribed
+      | return Except.error Broadcast.Error.noSubscribers
 
-    let mut currentSt := st
-    let mut currentNext := next
+    discard <| getValueByPosition next
 
-    while currentNext < currentSt.pos ∧ currentSt.size > 0 do
-      let some _val ← getValueByPosition currentNext | break
-
-      currentSt ← get
-      currentNext := currentNext + 1
-
-    set { currentSt with receivers := currentSt.receivers.erase bd.id }
+    set { st with receivers := st.receivers.erase bd.id }
 
     pure <| .ok ()
 
@@ -358,24 +347,21 @@ private def tryRecv'
     let next := st.receivers.get! receiverId
 
     if let some val ← getValueByPosition next then
-      modify ({ · with receivers := st.receivers.modify receiverId (· + 1) })
+      set { st with receivers := st.receivers.modify receiverId (· + 1) }
       return some val
     else
       return none
 
 private def tryRecv (ch : Bounded.Receiver α) : BaseIO (Option α) :=
   ch.state.atomically do
-    if (← get).closed ∨ ¬(← get).receivers.contains ch.id then
-      return none
-
     tryRecv' ch.id
 
 private partial def recv (ch : Bounded.Receiver α) : BaseIO (Task (Option α)) := do
   ch.state.atomically do
-    if (← get).closed ∨ ¬(← get).receivers.contains ch.id then
-      return .pure none
-    else if let some val ← tryRecv' ch.id then
+    if let some val ← tryRecv' ch.id then
       return .pure <| some val
+    else if (← get).closed then
+      return .pure <| none
     else
       let promise ← IO.Promise.new
       modify fun st => { st with waiters := st.waiters.enqueue ⟨promise, none⟩ }
@@ -506,13 +492,6 @@ Subscribes a new `Receiver` from the `Broadcast` channel.
 def subscribe (ch : Broadcast α) : IO (Broadcast.Receiver α) := do
   Broadcast.Receiver.mk <$> ch.inner.subscribe
 
-/--
-Closes a `Broadcast` channel.
--/
-@[inline]
-def close (ch : Broadcast α) : IO Unit := do
-  ch.inner.close
-
 /--
 Send a value through the broadcast channel, returning a task that will resolve once the transmission
 could be completed.
@@ -538,8 +517,11 @@ def tryRecv (ch : Broadcast.Receiver α) : BaseIO (Option α) :=
 Receive a value from the broadcast receiver, returning a task that will resolve with
 the next available message. This will block until a message is available.
 -/
-def recv [Inhabited α] (ch : Broadcast.Receiver α) : BaseIO (Task (Option α)) := do
-  Std.Bounded.Receiver.recv ch.inner
+def recv [Inhabited α] (ch : Broadcast.Receiver α) : BaseIO (Task α) := do
+  BaseIO.bindTask (sync := true) (← Std.Bounded.Receiver.recv ch.inner)
+    fun
+      | some val => return .pure val
+      | none => unreachable!
 
 open Internal.IO.Async in
 
@@ -581,18 +563,6 @@ partial def forAsync (f : α → BaseIO Unit) (ch : Broadcast.Receiver α)
   (prio : Task.Priority := .default) : BaseIO (Task Unit) := do
     ch.inner.forAsync f prio
 
-instance [Inhabited α] : AsyncStream (Broadcast.Receiver α) α where
-  next channel := channel.recvSelector
-  stop channel := channel.unsubscribe
-
-instance [Inhabited α] : AsyncRead (Broadcast.Receiver α) (Option α) where
-  read receiver := Internal.IO.Async.Async.ofIOTask receiver.recv
-
-instance [Inhabited α] : AsyncWrite (Broadcast α) α where
-  write receiver x := do
-    let task ← receiver.send x
-    discard <| Async.ofETask <| task
-
 end Receiver
 
 /--
@@ -638,18 +608,15 @@ def tryRecv (ch : Sync.Receiver α) : BaseIO (Option α) := Broadcast.Receiver.t
 /--
 Receive a value from the channel, blocking until the transmission could be completed.
 -/
-def recv [Inhabited α] (ch : Sync.Receiver α) : BaseIO (Option α) := do
+def recv [Inhabited α] (ch : Sync.Receiver α) : BaseIO α := do
   IO.wait (← Broadcast.Receiver.recv ch)
 
 partial def forIn [Inhabited α] [Monad m] [MonadLiftT BaseIO m]
     (ch : Sync.Receiver α) (f : α → β → m (ForInStep β)) : β → m β := fun b => do
   let a ← ch.recv
-  match a with
-  | none => pure b
-  | some a =>
-    match ← f a b with
-    | .done b => pure b
-    | .yield b => ch.forIn f b
+  match ← f a b with
+  | .done b => pure b
+  | .yield b => ch.forIn f b
 
 /-- `for msg in ch.sync do ...` receives all messages in the channel until it is closed. -/
 instance [Inhabited α] [MonadLiftT BaseIO m] : ForIn m (Sync.Receiver α) α where