rustdv_methodology/
channel.rs1use std::cell::{Cell, RefCell};
17use std::collections::VecDeque;
18use std::fmt;
19use std::future::Future;
20use std::pin::Pin;
21use std::rc::Rc;
22use std::task::{Context, Poll, Waker};
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum TlmError {
26 Disconnected,
28}
29
30impl fmt::Display for TlmError {
31 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32 write!(f, "channel disconnected")
33 }
34}
35impl std::error::Error for TlmError {}
36
37#[derive(Debug)]
39pub struct TlmFull<T>(pub T);
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub struct TlmEmpty;
42
43struct ChanInner<T> {
44 buf: RefCell<VecDeque<T>>,
45 cap: usize,
46 senders: Cell<usize>,
47 receivers: Cell<usize>,
48 send_waiters: RefCell<Vec<Waker>>,
49 recv_waiters: RefCell<Vec<Waker>>,
50}
51
52impl<T> ChanInner<T> {
53 fn has_space(&self) -> bool {
54 self.buf.borrow().len() < self.cap
55 }
56 fn wake_senders(&self) {
57 for w in self.send_waiters.borrow_mut().drain(..) {
58 w.wake();
59 }
60 }
61 fn wake_receivers(&self) {
62 for w in self.recv_waiters.borrow_mut().drain(..) {
63 w.wake();
64 }
65 }
66}
67
68pub fn channel<T>(capacity: usize) -> (Sender<T>, Receiver<T>) {
71 let inner = Rc::new(ChanInner {
72 buf: RefCell::new(VecDeque::new()),
73 cap: capacity.max(1),
74 senders: Cell::new(1),
75 receivers: Cell::new(1),
76 send_waiters: RefCell::new(Vec::new()),
77 recv_waiters: RefCell::new(Vec::new()),
78 });
79 (Sender { inner: inner.clone() }, Receiver { inner })
80}
81
82pub struct Sender<T> {
84 inner: Rc<ChanInner<T>>,
85}
86
87impl<T> Clone for Sender<T> {
88 fn clone(&self) -> Self {
89 self.inner.senders.set(self.inner.senders.get() + 1);
90 Sender { inner: self.inner.clone() }
91 }
92}
93
94impl<T> Drop for Sender<T> {
95 fn drop(&mut self) {
96 let n = self.inner.senders.get() - 1;
97 self.inner.senders.set(n);
98 if n == 0 {
99 self.inner.wake_receivers();
100 }
101 }
102}
103
104impl<T> Sender<T> {
105 pub fn send(&self, item: T) -> Send_<T> {
107 Send_ { inner: self.inner.clone(), item: Some(item) }
108 }
109
110 pub fn try_send(&self, item: T) -> Result<(), TlmFull<T>> {
112 if self.inner.receivers.get() == 0 {
113 return Err(TlmFull(item)); }
115 if self.inner.has_space() {
116 self.inner.buf.borrow_mut().push_back(item);
117 self.inner.wake_receivers();
118 Ok(())
119 } else {
120 Err(TlmFull(item))
121 }
122 }
123
124 pub fn can_send(&self) -> bool {
125 self.inner.has_space() && self.inner.receivers.get() > 0
126 }
127}
128
129pub struct Send_<T> {
130 inner: Rc<ChanInner<T>>,
131 item: Option<T>,
132}
133
134impl<T> Future for Send_<T> {
135 type Output = Result<(), TlmError>;
136 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
137 let this = unsafe { self.get_unchecked_mut() };
138 if this.inner.receivers.get() == 0 {
139 return Poll::Ready(Err(TlmError::Disconnected));
140 }
141 if this.inner.has_space() {
142 let item = this.item.take().expect("Send polled after completion");
143 this.inner.buf.borrow_mut().push_back(item);
144 this.inner.wake_receivers();
145 Poll::Ready(Ok(()))
146 } else {
147 this.inner.send_waiters.borrow_mut().push(cx.waker().clone());
148 Poll::Pending
149 }
150 }
151}
152
153pub struct Receiver<T> {
155 inner: Rc<ChanInner<T>>,
156}
157
158impl<T> Clone for Receiver<T> {
159 fn clone(&self) -> Self {
160 self.inner.receivers.set(self.inner.receivers.get() + 1);
161 Receiver { inner: self.inner.clone() }
162 }
163}
164
165impl<T> Drop for Receiver<T> {
166 fn drop(&mut self) {
167 let n = self.inner.receivers.get() - 1;
168 self.inner.receivers.set(n);
169 if n == 0 {
170 self.inner.wake_senders();
171 }
172 }
173}
174
175impl<T> Receiver<T> {
176 pub fn recv(&self) -> Recv<T> {
178 Recv { inner: self.inner.clone() }
179 }
180
181 pub fn try_recv(&self) -> Result<T, TlmEmpty> {
182 match self.inner.buf.borrow_mut().pop_front() {
183 Some(v) => {
184 self.inner.wake_senders();
185 Ok(v)
186 }
187 None => Err(TlmEmpty),
188 }
189 }
190
191 pub fn len(&self) -> usize {
192 self.inner.buf.borrow().len()
193 }
194 pub fn is_empty(&self) -> bool {
195 self.inner.buf.borrow().is_empty()
196 }
197}
198
199impl<T: Clone> Receiver<T> {
200 pub fn peek(&self) -> Peek<T> {
202 Peek { inner: self.inner.clone() }
203 }
204
205 pub fn try_peek(&self) -> Result<T, TlmEmpty> {
206 self.inner.buf.borrow().front().cloned().ok_or(TlmEmpty)
207 }
208}
209
210pub struct Recv<T> {
211 inner: Rc<ChanInner<T>>,
212}
213
214impl<T> Future for Recv<T> {
215 type Output = Result<T, TlmError>;
216 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
217 let item = self.inner.buf.borrow_mut().pop_front();
218 match item {
219 Some(v) => {
220 self.inner.wake_senders();
221 Poll::Ready(Ok(v))
222 }
223 None => {
224 if self.inner.senders.get() == 0 {
225 return Poll::Ready(Err(TlmError::Disconnected));
226 }
227 self.inner.recv_waiters.borrow_mut().push(cx.waker().clone());
228 Poll::Pending
229 }
230 }
231 }
232}
233
234pub struct Peek<T: Clone> {
235 inner: Rc<ChanInner<T>>,
236}
237
238impl<T: Clone> Future for Peek<T> {
239 type Output = Result<T, TlmError>;
240 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
241 let front = self.inner.buf.borrow().front().cloned();
242 match front {
243 Some(v) => Poll::Ready(Ok(v)),
244 None => {
245 if self.inner.senders.get() == 0 {
246 return Poll::Ready(Err(TlmError::Disconnected));
247 }
248 self.inner.recv_waiters.borrow_mut().push(cx.waker().clone());
249 Poll::Pending
250 }
251 }
252 }
253}