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 (
80 Sender {
81 inner: inner.clone(),
82 },
83 Receiver { inner },
84 )
85}
86
87pub struct Sender<T> {
89 inner: Rc<ChanInner<T>>,
90}
91
92impl<T> Clone for Sender<T> {
93 fn clone(&self) -> Self {
94 self.inner.senders.set(self.inner.senders.get() + 1);
95 Sender {
96 inner: self.inner.clone(),
97 }
98 }
99}
100
101impl<T> Drop for Sender<T> {
102 fn drop(&mut self) {
103 let n = self.inner.senders.get() - 1;
104 self.inner.senders.set(n);
105 if n == 0 {
106 self.inner.wake_receivers();
107 }
108 }
109}
110
111impl<T> Sender<T> {
112 pub fn send(&self, item: T) -> Send_<T> {
114 Send_ {
115 inner: self.inner.clone(),
116 item: Some(item),
117 }
118 }
119
120 pub fn try_send(&self, item: T) -> Result<(), TlmFull<T>> {
122 if self.inner.receivers.get() == 0 {
123 return Err(TlmFull(item)); }
125 if self.inner.has_space() {
126 self.inner.buf.borrow_mut().push_back(item);
127 self.inner.wake_receivers();
128 Ok(())
129 } else {
130 Err(TlmFull(item))
131 }
132 }
133
134 pub fn can_send(&self) -> bool {
135 self.inner.has_space() && self.inner.receivers.get() > 0
136 }
137}
138
139pub struct Send_<T> {
140 inner: Rc<ChanInner<T>>,
141 item: Option<T>,
142}
143
144impl<T> Future for Send_<T> {
145 type Output = Result<(), TlmError>;
146 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
147 let this = unsafe { self.get_unchecked_mut() };
148 if this.inner.receivers.get() == 0 {
149 return Poll::Ready(Err(TlmError::Disconnected));
150 }
151 if this.inner.has_space() {
152 let item = this.item.take().expect("Send polled after completion");
153 this.inner.buf.borrow_mut().push_back(item);
154 this.inner.wake_receivers();
155 Poll::Ready(Ok(()))
156 } else {
157 this.inner
158 .send_waiters
159 .borrow_mut()
160 .push(cx.waker().clone());
161 Poll::Pending
162 }
163 }
164}
165
166pub struct Receiver<T> {
168 inner: Rc<ChanInner<T>>,
169}
170
171impl<T> Clone for Receiver<T> {
172 fn clone(&self) -> Self {
173 self.inner.receivers.set(self.inner.receivers.get() + 1);
174 Receiver {
175 inner: self.inner.clone(),
176 }
177 }
178}
179
180impl<T> Drop for Receiver<T> {
181 fn drop(&mut self) {
182 let n = self.inner.receivers.get() - 1;
183 self.inner.receivers.set(n);
184 if n == 0 {
185 self.inner.wake_senders();
186 }
187 }
188}
189
190impl<T> Receiver<T> {
191 pub fn recv(&self) -> Recv<T> {
193 Recv {
194 inner: self.inner.clone(),
195 }
196 }
197
198 pub fn try_recv(&self) -> Result<T, TlmEmpty> {
199 match self.inner.buf.borrow_mut().pop_front() {
200 Some(v) => {
201 self.inner.wake_senders();
202 Ok(v)
203 }
204 None => Err(TlmEmpty),
205 }
206 }
207
208 pub fn len(&self) -> usize {
209 self.inner.buf.borrow().len()
210 }
211 pub fn is_empty(&self) -> bool {
212 self.inner.buf.borrow().is_empty()
213 }
214}
215
216impl<T: Clone> Receiver<T> {
217 pub fn peek(&self) -> Peek<T> {
219 Peek {
220 inner: self.inner.clone(),
221 }
222 }
223
224 pub fn try_peek(&self) -> Result<T, TlmEmpty> {
225 self.inner.buf.borrow().front().cloned().ok_or(TlmEmpty)
226 }
227}
228
229pub struct Recv<T> {
230 inner: Rc<ChanInner<T>>,
231}
232
233impl<T> Future for Recv<T> {
234 type Output = Result<T, TlmError>;
235 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
236 let item = self.inner.buf.borrow_mut().pop_front();
237 match item {
238 Some(v) => {
239 self.inner.wake_senders();
240 Poll::Ready(Ok(v))
241 }
242 None => {
243 if self.inner.senders.get() == 0 {
244 return Poll::Ready(Err(TlmError::Disconnected));
245 }
246 self.inner
247 .recv_waiters
248 .borrow_mut()
249 .push(cx.waker().clone());
250 Poll::Pending
251 }
252 }
253 }
254}
255
256pub struct Peek<T: Clone> {
257 inner: Rc<ChanInner<T>>,
258}
259
260impl<T: Clone> Future for Peek<T> {
261 type Output = Result<T, TlmError>;
262 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
263 let front = self.inner.buf.borrow().front().cloned();
264 match front {
265 Some(v) => Poll::Ready(Ok(v)),
266 None => {
267 if self.inner.senders.get() == 0 {
268 return Poll::Ready(Err(TlmError::Disconnected));
269 }
270 self.inner
271 .recv_waiters
272 .borrow_mut()
273 .push(cx.waker().clone());
274 Poll::Pending
275 }
276 }
277 }
278}