1 use std::cell::Cell;
2 use std::future::Future;
3 use std::marker::PhantomData;
4 use std::pin::Pin;
5 use std::ptr;
6 use std::task::{Context, Poll};
7
8 #[derive(Debug)]
9 pub struct Sender<T> {
10 _p: PhantomData<T>,
11 }
12
13 #[derive(Debug)]
14 pub struct Receiver<T> {
15 _p: PhantomData<T>,
16 }
17
18 pub(crate) struct Enter<'a, T> {
19 _rx: &'a mut Receiver<T>,
20 prev: *mut (),
21 }
22
pair<T>() -> (Sender<T>, Receiver<T>)23 pub fn pair<T>() -> (Sender<T>, Receiver<T>) {
24 let tx = Sender { _p: PhantomData };
25 let rx = Receiver { _p: PhantomData };
26 (tx, rx)
27 }
28
29 // Tracks the pointer to `Option<T>`.
30 //
31 // TODO: Ensure wakers match?
32 thread_local!(static STORE: Cell<*mut ()> = Cell::new(ptr::null_mut()));
33
34 // ===== impl Sender =====
35
36 impl<T: Unpin> Sender<T> {
send(&mut self, value: T) -> impl Future<Output = ()>37 pub fn send(&mut self, value: T) -> impl Future<Output = ()> {
38 Send { value: Some(value) }
39 }
40 }
41
42 struct Send<T> {
43 value: Option<T>,
44 }
45
46 impl<T: Unpin> Future for Send<T> {
47 type Output = ();
48
poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()>49 fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
50 if self.value.is_none() {
51 return Poll::Ready(());
52 }
53
54 STORE.with(|cell| unsafe {
55 let ptr = cell.get() as *mut Option<T>;
56 let option_ref = ptr.as_mut().expect("invalid usage");
57
58 if option_ref.is_none() {
59 *option_ref = self.value.take();
60 }
61
62 Poll::Pending
63 })
64 }
65 }
66
67 // ===== impl Receiver =====
68
69 impl<T> Receiver<T> {
enter<'a>(&'a mut self, dst: &'a mut Option<T>) -> Enter<'a, T>70 pub(crate) fn enter<'a>(&'a mut self, dst: &'a mut Option<T>) -> Enter<'a, T> {
71 let prev = STORE.with(|cell| {
72 let prev = cell.get();
73 cell.set(dst as *mut _ as *mut ());
74 prev
75 });
76
77 Enter { _rx: self, prev }
78 }
79 }
80
81 // ===== impl Enter =====
82
83 impl<'a, T> Drop for Enter<'a, T> {
drop(&mut self)84 fn drop(&mut self) {
85 STORE.with(|cell| cell.set(self.prev));
86 }
87 }
88