spike implementation

Author: carllerche

async-stream/src/async_stream.rs

@@ -7,12 +7,13 @@ use std::task::{Context, Poll};
 
 pub struct AsyncStream<T, U> {
     rx: Receiver<T>,
+    done: bool,
     generator: U,
 }
 
 impl<T, U> AsyncStream<T, U> {
     pub fn new(rx: Receiver<T>, generator: U) -> AsyncStream<T, U> {
-        AsyncStream { rx, generator }
+        AsyncStream { rx, done: false, generator }
     }
 }
 
@@ -22,6 +23,30 @@ where U: Future<Output = ()>
     type Item = T;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
-        unimplemented!();
+        unsafe {
+            let me = Pin::get_unchecked_mut(self);
+
+            if me.done {
+                panic!("poll after done");
+            }
+
+            let mut dst = None;
+            let res = {
+                let _enter = me.rx.enter(&mut dst);
+                Pin::new_unchecked(&mut me.generator).poll(cx)
+            };
+
+            me.done = res.is_ready();
+
+            if dst.is_some() {
+                return Poll::Ready(dst.take());
+            }
+
+            if me.done {
+                Poll::Ready(None)
+            } else {
+                Poll::Pending
+            }
+        }
     }
 }

async-stream/src/yielder.rs

@@ -1,4 +1,9 @@
+use std::cell::Cell;
+use std::future::Future;
 use std::marker::PhantomData;
+use std::pin::Pin;
+use std::ptr;
+use std::task::{Context, Poll};
 
 pub struct Sender<T> {
     _p: PhantomData<T>,
@@ -8,12 +13,74 @@ pub struct Receiver<T> {
     _p: PhantomData<T>,
 }
 
+pub(crate) struct Enter<'a, T> {
+    _rx: &'a mut Receiver<T>,
+}
+
 pub fn pair<T>() -> (Sender<T>, Receiver<T>) {
-    unimplemented!();
+    let tx = Sender { _p: PhantomData };
+    let rx = Receiver { _p: PhantomData };
+    (tx, rx)
+}
+
+// Tracks the pointer to `Option<T>`.
+//
+// TODO: Ensure wakers match?
+thread_local!(static STORE: Cell<*mut ()> = Cell::new(ptr::null_mut()));
+
+// ===== impl Sender =====
+
+impl<T: Unpin + 'static> Sender<T> {
+    pub fn send(&mut self, value: T) -> impl Future<Output = ()> {
+        Send { value: Some(value) }
+    }
+}
+
+struct Send<T> {
+    value: Option<T>,
 }
 
-impl<T> Sender<T> {
-    pub async fn send(&mut self, value: T) {
-        unimplemented!();
+impl<T: Unpin + 'static> Future for Send<T> {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        if self.value.is_none() {
+            return Poll::Ready(());
+        }
+
+        STORE.with(|cell| {
+            unsafe {
+                let ptr = cell.get() as *mut Option<T>;
+                let option_ref = ptr.as_mut().expect("invalid usage");
+
+                if option_ref.is_none() {
+                    *option_ref = self.value.take();
+                }
+
+                Poll::Pending
+            }
+        })
+    }
+}
+
+// ===== impl Receiver =====
+
+impl<T> Receiver<T> {
+    pub(crate) fn enter<'a>(&'a mut self, dst: &'a mut Option<T>) -> Enter<'a, T> {
+        STORE.with(|cell| {
+            assert!(cell.get().is_null());
+
+            cell.set(dst as *mut _ as *mut ());
+        });
+
+        Enter { _rx: self }
+    }
+}
+
+// ===== impl Enter =====
+
+impl<'a, T> Drop for Enter<'a, T> {
+    fn drop(&mut self) {
+        STORE.with(|cell| cell.set(ptr::null_mut()))
     }
 }