Add notify

Author: Virv12

pixie-uefi/src/os/mod.rs

@@ -18,6 +18,7 @@ pub mod input;
 mod logger;
 pub mod memory;
 pub mod net;
+mod notify;
 mod send_wrapper;
 mod timer;
 pub mod ui;

pixie-uefi/src/os/notify.rs

@@ -0,0 +1,87 @@
+use alloc::sync::Arc;
+use alloc::vec::Vec;
+use core::future::Future;
+use core::pin::Pin;
+use core::task::{Context, Poll, Waker};
+
+use spin::Mutex;
+
+struct NotifyInner {
+    vec: Vec<Arc<Mutex<NotifiedInner>>>,
+}
+
+pub struct Notify {
+    inner: Mutex<NotifyInner>,
+}
+
+impl Notify {
+    pub fn new() -> Self {
+        let inner = Mutex::new(NotifyInner { vec: Vec::new() });
+        Notify { inner }
+    }
+
+    pub fn wait(&self) -> Notified<'_> {
+        let mut n = self.inner.lock();
+        let inner = Arc::new(Mutex::new(NotifiedInner {
+            notified: false,
+            position: n.vec.len(),
+            waker: None,
+        }));
+        n.vec.push(inner.clone());
+        Notified {
+            notify: self,
+            inner,
+        }
+    }
+
+    pub fn notify_all(&self) {
+        let mut n = self.inner.lock();
+        for inner in n.vec.drain(..) {
+            let mut guard = inner.lock();
+            guard.notified = true;
+            if let Some(waker) = guard.waker.take() {
+                waker.wake();
+            }
+        }
+    }
+}
+
+struct NotifiedInner {
+    notified: bool,
+    position: usize,
+    waker: Option<Waker>,
+}
+
+pub struct Notified<'a> {
+    notify: &'a Notify,
+    inner: Arc<Mutex<NotifiedInner>>,
+}
+
+impl Future for Notified<'_> {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        let mut guard = self.inner.lock();
+        if guard.notified {
+            Poll::Ready(())
+        } else {
+            guard.waker = Some(cx.waker().clone());
+            Poll::Pending
+        }
+    }
+}
+
+impl Drop for Notified<'_> {
+    fn drop(&mut self) {
+        let mut n = self.notify.inner.lock();
+        let i = self.inner.lock();
+        if i.notified {
+            return;
+        }
+        n.vec.swap_remove(i.position);
+        if let Some(swapped) = n.vec.get(i.position) {
+            let mut swapped_guard = swapped.lock();
+            swapped_guard.position = i.position;
+        }
+    }
+}