fix: remove async from notify and test

Author: algebraic-dev

src/Std/Sync/Notify.lean

@@ -106,13 +106,13 @@ def notifyOne (x : Notify) : BaseIO Bool := do
 Wait to be notified. Returns a task that completes when notify is called.
 Note: if notify was called before wait, this will wait for the next notify call.
 -/
-def wait (x : Notify) : Async Unit :=
+def wait (x : Notify) : IO (AsyncTask Unit) :=
   x.state.atomically do
     let promise ← IO.Promise.new
     modify fun st => { st with consumers := st.consumers.enqueue (.normal promise) }
-    match ← await promise.result? with
-    | some res => pure res
-    | none => throw (IO.userError "notify dropped")
+    IO.bindTask promise.result? fun
+      | some res => pure <| Task.pure (.ok res)
+      | none => throw (IO.userError "notify dropped")
 
 /--
 Creates a selector that waits for notifications

tests/lean/run/sync_notify.lean

@@ -0,0 +1,108 @@
+import Std.Internal.Async
+import Std.Sync
+
+open Std.Internal.IO Async
+
+-- Test basic wait and notifyOne functionality
+def testBasicWaitNotifyOne : Async Unit := do
+  let notify ← Std.Notify.new
+  let waitTask ← notify.wait
+
+  assert! (← waitTask.getState) = .waiting
+  discard <| notify.notifyOne
+  await waitTask
+  assert! (← waitTask.getState) = .finished
+
+#eval testBasicWaitNotifyOne.block
+
+-- Test multiple waiters with notifyOne (only one should be notified)
+def testMultipleWaitersNotifyOne : Async Unit := do
+  let notify ← Std.Notify.new
+  let task1 ← notify.wait
+  let task2 ← notify.wait
+  let task3 ← notify.wait
+
+  assert! (← task1.getState) = .waiting
+  assert! (← task2.getState) = .waiting
+  assert! (← task3.getState) = .waiting
+
+  discard <| notify.notifyOne
+
+  IO.sleep 100
+
+  let states ← [task1, task2, task3].mapM (fun t => t.getState)
+  let finishedCount := states.filter (· == .finished) |>.length
+  let waitingCount := states.filter (· == .waiting) |>.length
+
+  assert! finishedCount == 1
+  assert! waitingCount == 2
+
+  discard <| notify.notifyOne
+
+#eval testMultipleWaitersNotifyOne.block
+
+-- Test multiple waiters with notify (all should be notified)
+def testMultipleWaitersNotifyAll : Async Unit := do
+  let notify ← Std.Notify.new
+  let task1 ← notify.wait
+  let task2 ← notify.wait
+  let task3 ← notify.wait
+
+  assert! (← task1.getState) = .waiting
+  assert! (← task2.getState) = .waiting
+  assert! (← task3.getState) = .waiting
+
+  discard <| notify.notify
+
+  await task1
+  await task2
+  await task3
+
+  assert! (← task1.getState) = .finished
+  assert! (← task2.getState) = .finished
+  assert! (← task3.getState) = .finished
+
+#eval testMultipleWaitersNotifyAll.block
+
+-- Test sequential notification
+def testSequentialNotification : Async Unit := do
+  let notify ← Std.Notify.new
+  let task1 ← notify.wait
+  let task2 ← notify.wait
+  let task3 ← notify.wait
+
+  discard <| notify.notifyOne
+  await task1
+  assert! (← task1.getState) = .finished
+  assert! (← task2.getState) = .waiting
+  assert! (← task3.getState) = .waiting
+
+  discard <| notify.notifyOne
+  await task2
+  assert! (← task2.getState) = .finished
+  assert! (← task3.getState) = .waiting
+
+  discard <| notify.notifyOne
+  await task3
+  assert! (← task3.getState) = .finished
+
+#eval testSequentialNotification.block
+
+def testReuseAfterCompletion : Async Unit := do
+  let notify ← Std.Notify.new
+
+  let task1 ← notify.wait
+  discard <| notify.notifyOne
+  await task1
+  assert! (← task1.getState) = .finished
+
+  let task2 ← notify.wait
+  let task3 ← notify.wait
+  discard <| notify.notify
+  await task2
+  await task3
+
+  assert! (← task2.getState) = .finished
+  assert! (← task3.getState) = .finished
+
+#eval testReuseAfterCompletion.block