1use std::any::type_name;
4use std::any::Any;
5use std::borrow::Borrow;
6use std::cell::Cell;
7use std::cell::UnsafeCell;
8use std::collections::VecDeque;
9use std::fmt;
10use std::fmt::Debug;
11use std::fmt::Formatter;
12use std::ops::Deref;
13use std::rc::Rc;
14
15use self::internal as i;
16
17pub type AsyncRef<T> = i::AsyncBorrowImpl<T, i::Shared>;
18pub type AsyncMut<T> = i::AsyncBorrowImpl<T, i::Exclusive>;
19
20pub type AsyncRefFuture<T> = i::AsyncBorrowFutureImpl<T, i::Shared>;
21pub type AsyncMutFuture<T> = i::AsyncBorrowFutureImpl<T, i::Exclusive>;
22
23pub struct AsyncRefCell<T> {
24 value: UnsafeCell<T>,
25 borrow_count: Cell<i::BorrowCount>,
26 waiters: Cell<VecDeque<Option<i::Waiter>>>,
27 turn: Cell<usize>,
28}
29
30impl<T: 'static> AsyncRefCell<T> {
31 pub fn new(value: T) -> Self {
37 Self {
38 value: UnsafeCell::new(value),
39 borrow_count: Default::default(),
40 waiters: Default::default(),
41 turn: Default::default(),
42 }
43 }
44
45 pub fn new_rc(value: T) -> Rc<Self> {
46 Rc::new(Self::new(value))
47 }
48
49 pub fn as_ptr(&self) -> *mut T {
50 self.value.get()
51 }
52
53 pub fn into_inner(self) -> T {
54 assert!(self.borrow_count.get().is_empty());
55 self.value.into_inner()
56 }
57}
58
59impl<T> Debug for AsyncRefCell<T> {
60 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
61 write!(f, "AsyncRefCell<{}>", type_name::<T>())
62 }
63}
64
65impl<T: Default + 'static> Default for AsyncRefCell<T> {
66 fn default() -> Self {
67 Self::new(Default::default())
68 }
69}
70
71impl<T: Default + 'static> AsyncRefCell<T> {
72 pub fn default_rc() -> Rc<Self> {
73 Rc::new(Default::default())
74 }
75}
76
77impl<T: 'static> From<T> for AsyncRefCell<T> {
78 fn from(value: T) -> Self {
79 Self::new(value)
80 }
81}
82
83impl<T> AsyncRefCell<T> {
84 pub fn borrow(self: &Rc<Self>) -> AsyncRefFuture<T> {
85 AsyncRefFuture::new(self)
86 }
87
88 pub fn borrow_mut(self: &Rc<Self>) -> AsyncMutFuture<T> {
89 AsyncMutFuture::new(self)
90 }
91
92 pub fn try_borrow(self: &Rc<Self>) -> Option<AsyncRef<T>> {
93 Self::borrow_sync(self)
94 }
95
96 pub fn try_borrow_mut(self: &Rc<Self>) -> Option<AsyncMut<T>> {
97 Self::borrow_sync(self)
98 }
99}
100
101impl<T> RcRef<AsyncRefCell<T>> {
102 pub fn borrow(&self) -> AsyncRefFuture<T> {
103 AsyncRefFuture::new(self)
104 }
105
106 pub fn borrow_mut(&self) -> AsyncMutFuture<T> {
107 AsyncMutFuture::new(self)
108 }
109
110 pub fn try_borrow(&self) -> Option<AsyncRef<T>> {
111 AsyncRefCell::<T>::borrow_sync(self)
112 }
113
114 pub fn try_borrow_mut(&self) -> Option<AsyncMut<T>> {
115 AsyncRefCell::<T>::borrow_sync(self)
116 }
117}
118
119#[derive(Debug)]
145pub struct RcRef<T> {
146 rc: Rc<dyn Any>,
147 value: *const T,
148}
149
150impl<T: 'static> RcRef<T> {
151 pub fn new(value: T) -> Self {
152 Self::from(Rc::new(value))
153 }
154
155 pub fn map<S: 'static, R: RcLike<S>, F: FnOnce(&S) -> &T>(
156 source: R,
157 map_fn: F,
158 ) -> RcRef<T> {
159 let RcRef::<S> { rc, value } = source.into();
160 #[allow(clippy::undocumented_unsafe_blocks)]
162 let value = map_fn(unsafe { &*value });
163 RcRef { rc, value }
164 }
165
166 pub(crate) fn split(rc_ref: &Self) -> (&T, &Rc<dyn Any>) {
167 let &Self { ref rc, value } = rc_ref;
168 #[allow(clippy::undocumented_unsafe_blocks)]
170 (unsafe { &*value }, rc)
171 }
172}
173
174impl<T: Default + 'static> Default for RcRef<T> {
175 fn default() -> Self {
176 Self::new(Default::default())
177 }
178}
179
180impl<T> Clone for RcRef<T> {
181 fn clone(&self) -> Self {
182 Self {
183 rc: self.rc.clone(),
184 value: self.value,
185 }
186 }
187}
188
189impl<T: 'static> From<&RcRef<T>> for RcRef<T> {
190 fn from(rc_ref: &RcRef<T>) -> Self {
191 rc_ref.clone()
192 }
193}
194
195impl<T: 'static> From<Rc<T>> for RcRef<T> {
196 fn from(rc: Rc<T>) -> Self {
197 Self {
198 value: &*rc,
199 rc: rc as Rc<_>,
200 }
201 }
202}
203
204impl<T: 'static> From<&Rc<T>> for RcRef<T> {
205 fn from(rc: &Rc<T>) -> Self {
206 rc.clone().into()
207 }
208}
209
210impl<T> Deref for RcRef<T> {
211 type Target = T;
212 fn deref(&self) -> &Self::Target {
213 #[allow(clippy::undocumented_unsafe_blocks)]
215 unsafe {
216 &*self.value
217 }
218 }
219}
220
221impl<T> Borrow<T> for RcRef<T> {
222 fn borrow(&self) -> &T {
223 self
224 }
225}
226
227impl<T> AsRef<T> for RcRef<T> {
228 fn as_ref(&self) -> &T {
229 self
230 }
231}
232
233pub trait RcLike<T>: AsRef<T> + Into<RcRef<T>> {}
236
237impl<T: 'static> RcLike<T> for Rc<T> {}
238impl<T: 'static> RcLike<T> for RcRef<T> {}
239impl<T: 'static> RcLike<T> for &Rc<T> {}
240impl<T: 'static> RcLike<T> for &RcRef<T> {}
241
242mod internal {
243 use super::AsyncRefCell;
244 use super::RcLike;
245 use super::RcRef;
246 use futures::future::Future;
247 use futures::ready;
248 use futures::task::Context;
249 use futures::task::Poll;
250 use futures::task::Waker;
251 use std::borrow::Borrow;
252 use std::borrow::BorrowMut;
253 use std::fmt::Debug;
254 use std::marker::PhantomData;
255 use std::ops::Deref;
256 use std::ops::DerefMut;
257 use std::pin::Pin;
258
259 impl<T> AsyncRefCell<T> {
260 pub fn borrow_sync<M: BorrowModeTrait, R: RcLike<AsyncRefCell<T>>>(
265 cell: R,
266 ) -> Option<AsyncBorrowImpl<T, M>> {
267 let cell_ref = cell.as_ref();
268 #[allow(clippy::undocumented_unsafe_blocks)]
273 let waiters = unsafe { &mut *cell_ref.waiters.as_ptr() };
274 if waiters.is_empty() {
275 let new_borrow_count =
280 cell_ref.borrow_count.get().try_add(M::borrow_mode())?;
281 cell_ref.borrow_count.set(new_borrow_count);
282 Some(AsyncBorrowImpl::<T, M>::new(cell.into()))
283 } else {
284 None
285 }
286 }
287
288 fn drop_borrow<M: BorrowModeTrait>(&self) {
289 let new_borrow_count = self.borrow_count.get().remove(M::borrow_mode());
290 self.borrow_count.set(new_borrow_count);
291
292 if new_borrow_count.is_empty() {
293 self.wake_waiters()
294 }
295 }
296
297 fn create_waiter<M: BorrowModeTrait>(&self) -> usize {
298 let waiter = Waiter::new(M::borrow_mode());
299 let turn = self.turn.get();
300 let index = {
301 #[allow(clippy::undocumented_unsafe_blocks)]
303 let waiters = unsafe { &mut *self.waiters.as_ptr() };
304 waiters.push_back(Some(waiter));
305 waiters.len() - 1
306 };
307 if index == 0 {
308 self.wake_waiters()
310 }
311 turn + index
313 }
314
315 fn poll_waiter<M: BorrowModeTrait>(
316 &self,
317 id: usize,
318 cx: &mut Context,
319 ) -> Poll<()> {
320 let borrow_count = self.borrow_count.get();
321 let turn = self.turn.get();
322 if id < turn {
323 let _ = borrow_count.remove(M::borrow_mode());
327 Poll::Ready(())
328 } else {
329 #[allow(clippy::undocumented_unsafe_blocks)]
332 let waiters = unsafe { &mut *self.waiters.as_ptr() };
333 assert!(id < turn + waiters.len());
335 assert!(id > turn || borrow_count.try_add(M::borrow_mode()).is_none());
338 let waiter_mut = waiters[id - turn].as_mut().unwrap();
340 waiter_mut.set_waker(cx.waker());
341 Poll::Pending
342 }
343 }
344
345 fn wake_waiters(&self) {
346 let mut borrow_count = self.borrow_count.get();
347 #[allow(clippy::undocumented_unsafe_blocks)]
349 let waiters = unsafe { &mut *self.waiters.as_ptr() };
350 let mut turn = self.turn.get();
351
352 loop {
353 let waiter_entry = match waiters.front().map(Option::as_ref) {
354 None => break, Some(w) => w,
356 };
357 let borrow_mode = match waiter_entry {
358 None => {
359 waiters.pop_front();
362 turn += 1;
363 continue;
364 }
365 Some(waiter) => waiter.borrow_mode(),
366 };
367 borrow_count = match borrow_count.try_add(borrow_mode) {
373 None => break, Some(b) => b,
375 };
376 let mut waiter = waiters.pop_front().unwrap().unwrap();
378 turn += 1;
379 if let Some(waker) = waiter.take_waker() {
381 waker.wake()
382 }
383 }
384 self.borrow_count.set(borrow_count);
386 self.turn.set(turn);
387 }
388
389 fn drop_waiter<M: BorrowModeTrait>(&self, id: usize) {
390 let turn = self.turn.get();
391 if id < turn {
392 self.drop_borrow::<M>();
396 } else {
397 #[allow(clippy::undocumented_unsafe_blocks)]
400 let waiters = unsafe { &mut *self.waiters.as_ptr() };
401 waiters[id - turn].take().unwrap();
402 }
403
404 if id == turn {
405 self.wake_waiters()
408 }
409 }
410 }
411
412 pub struct AsyncBorrowFutureImpl<T: 'static, M: BorrowModeTrait> {
413 cell: Option<RcRef<AsyncRefCell<T>>>,
414 id: usize,
415 _phantom: PhantomData<M>,
416 }
417
418 impl<T, M: BorrowModeTrait> AsyncBorrowFutureImpl<T, M> {
419 pub fn new<R: RcLike<AsyncRefCell<T>>>(cell: R) -> Self {
420 Self {
421 id: cell.as_ref().create_waiter::<M>(),
422 cell: Some(cell.into()),
423 _phantom: PhantomData,
424 }
425 }
426 }
427
428 impl<T: 'static, M: BorrowModeTrait> Future for AsyncBorrowFutureImpl<T, M> {
429 type Output = AsyncBorrowImpl<T, M>;
430 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
431 ready!(self.cell.as_ref().unwrap().poll_waiter::<M>(self.id, cx));
432 #[allow(clippy::undocumented_unsafe_blocks)]
434 let self_mut = unsafe { Pin::get_unchecked_mut(self) };
435 let cell = self_mut.cell.take().unwrap();
436 Poll::Ready(AsyncBorrowImpl::<T, M>::new(cell))
437 }
438 }
439
440 impl<T, M: BorrowModeTrait> Drop for AsyncBorrowFutureImpl<T, M> {
441 fn drop(&mut self) {
442 if let Some(cell) = self.cell.take() {
449 cell.drop_waiter::<M>(self.id)
450 }
451 }
452 }
453
454 pub struct AsyncBorrowImpl<T: 'static, M: BorrowModeTrait> {
455 cell: RcRef<AsyncRefCell<T>>,
456 _phantom: PhantomData<M>,
457 }
458
459 impl<T, M: BorrowModeTrait> AsyncBorrowImpl<T, M> {
460 fn new(cell: RcRef<AsyncRefCell<T>>) -> Self {
461 Self {
462 cell,
463 _phantom: PhantomData,
464 }
465 }
466 }
467
468 impl<T, M: BorrowModeTrait> Deref for AsyncBorrowImpl<T, M> {
469 type Target = T;
470 fn deref(&self) -> &Self::Target {
471 #[allow(clippy::undocumented_unsafe_blocks)]
473 unsafe {
474 &*self.cell.as_ptr()
475 }
476 }
477 }
478
479 impl<T, M: BorrowModeTrait> Borrow<T> for AsyncBorrowImpl<T, M> {
480 fn borrow(&self) -> &T {
481 self
482 }
483 }
484
485 impl<T, M: BorrowModeTrait> AsRef<T> for AsyncBorrowImpl<T, M> {
486 fn as_ref(&self) -> &T {
487 self
488 }
489 }
490
491 impl<T> DerefMut for AsyncBorrowImpl<T, Exclusive> {
492 fn deref_mut(&mut self) -> &mut Self::Target {
493 #[allow(clippy::undocumented_unsafe_blocks)]
495 unsafe {
496 &mut *self.cell.as_ptr()
497 }
498 }
499 }
500
501 impl<T> BorrowMut<T> for AsyncBorrowImpl<T, Exclusive> {
502 fn borrow_mut(&mut self) -> &mut T {
503 self
504 }
505 }
506
507 impl<T> AsMut<T> for AsyncBorrowImpl<T, Exclusive> {
508 fn as_mut(&mut self) -> &mut T {
509 self
510 }
511 }
512
513 impl<T, M: BorrowModeTrait> Drop for AsyncBorrowImpl<T, M> {
514 fn drop(&mut self) {
515 self.cell.drop_borrow::<M>()
516 }
517 }
518
519 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
520 pub enum BorrowMode {
521 Shared,
522 Exclusive,
523 }
524
525 pub trait BorrowModeTrait: Copy {
526 fn borrow_mode() -> BorrowMode;
527 }
528
529 #[derive(Copy, Clone, Debug)]
530 pub struct Shared;
531
532 impl BorrowModeTrait for Shared {
533 fn borrow_mode() -> BorrowMode {
534 BorrowMode::Shared
535 }
536 }
537
538 #[derive(Copy, Clone, Debug)]
539 pub struct Exclusive;
540
541 impl BorrowModeTrait for Exclusive {
542 fn borrow_mode() -> BorrowMode {
543 BorrowMode::Exclusive
544 }
545 }
546
547 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
548 pub enum BorrowCount {
549 Shared(usize),
550 Exclusive,
551 }
552
553 impl Default for BorrowCount {
554 fn default() -> Self {
555 Self::Shared(0)
556 }
557 }
558
559 impl BorrowCount {
560 pub fn is_empty(self) -> bool {
561 matches!(self, BorrowCount::Shared(0))
562 }
563
564 pub fn try_add(self, mode: BorrowMode) -> Option<BorrowCount> {
565 match (self, mode) {
566 (BorrowCount::Shared(refs), BorrowMode::Shared) => {
567 Some(BorrowCount::Shared(refs + 1))
568 }
569 (BorrowCount::Shared(0), BorrowMode::Exclusive) => {
570 Some(BorrowCount::Exclusive)
571 }
572 _ => None,
573 }
574 }
575
576 #[allow(dead_code)]
577 pub fn add(self, mode: BorrowMode) -> BorrowCount {
578 match self.try_add(mode) {
579 Some(value) => value,
580 None => panic!("Can't add {mode:?} to {self:?}"),
581 }
582 }
583
584 pub fn try_remove(self, mode: BorrowMode) -> Option<BorrowCount> {
585 match (self, mode) {
586 (BorrowCount::Shared(refs), BorrowMode::Shared) if refs > 0 => {
587 Some(BorrowCount::Shared(refs - 1))
588 }
589 (BorrowCount::Exclusive, BorrowMode::Exclusive) => {
590 Some(BorrowCount::Shared(0))
591 }
592 _ => None,
593 }
594 }
595
596 pub fn remove(self, mode: BorrowMode) -> BorrowCount {
597 match self.try_remove(mode) {
598 Some(value) => value,
599 None => panic!("Can't remove {mode:?} from {self:?}"),
600 }
601 }
602 }
603
604 pub struct Waiter {
607 borrow_mode: BorrowMode,
608 waker: Option<Waker>,
609 }
610
611 impl Waiter {
612 pub fn new(borrow_mode: BorrowMode) -> Self {
613 Self {
614 borrow_mode,
615 waker: None,
616 }
617 }
618
619 pub fn borrow_mode(&self) -> BorrowMode {
620 self.borrow_mode
621 }
622
623 pub fn set_waker(&mut self, new_waker: &Waker) {
624 if self
625 .waker
626 .as_ref()
627 .filter(|waker| waker.will_wake(new_waker))
628 .is_none()
629 {
630 self.waker.replace(new_waker.clone());
631 }
632 }
633
634 pub fn take_waker(&mut self) -> Option<Waker> {
635 self.waker.take()
636 }
637 }
638}
639
640#[cfg(test)]
641mod tests {
642 use super::*;
643
644 #[derive(Default)]
645 struct Thing {
646 touch_count: usize,
647 _private: (),
648 }
649
650 impl Thing {
651 pub fn look(&self) -> usize {
652 self.touch_count
653 }
654
655 pub fn touch(&mut self) -> usize {
656 self.touch_count += 1;
657 self.touch_count
658 }
659 }
660
661 #[test]
662 fn async_ref_cell_borrow() {
663 let runtime = tokio::runtime::Builder::new_current_thread()
664 .build()
665 .unwrap();
666 runtime.block_on(async {
667 let cell = AsyncRefCell::<Thing>::default_rc();
668
669 let fut1 = cell.borrow();
670 let fut2 = cell.borrow_mut();
671 let fut3 = cell.borrow();
672 let fut4 = cell.borrow();
673 let fut5 = cell.borrow();
674 let fut6 = cell.borrow();
675 let fut7 = cell.borrow_mut();
676 let fut8 = cell.borrow();
677
678 assert!(cell.try_borrow().is_none());
681 assert!(cell.try_borrow_mut().is_none());
682
683 assert_eq!(fut1.await.look(), 0);
684
685 assert_eq!(fut2.await.touch(), 1);
686
687 {
688 let ref5 = fut5.await;
689 let ref4 = fut4.await;
690 let ref3 = fut3.await;
691 let ref6 = fut6.await;
692 assert_eq!(ref3.look(), 1);
693 assert_eq!(ref4.look(), 1);
694 assert_eq!(ref5.look(), 1);
695 assert_eq!(ref6.look(), 1);
696 }
697
698 {
699 let mut ref7 = fut7.await;
700 assert_eq!(ref7.look(), 1);
701 assert_eq!(ref7.touch(), 2);
702 }
703
704 {
705 let ref8 = fut8.await;
706 assert_eq!(ref8.look(), 2);
707 }
708 });
709 }
710
711 #[test]
712 fn async_ref_cell_try_borrow() {
713 let cell = AsyncRefCell::<Thing>::default_rc();
714
715 {
716 let ref1 = cell.try_borrow().unwrap();
717 assert_eq!(ref1.look(), 0);
718 assert!(cell.try_borrow_mut().is_none());
719 }
720
721 {
722 let mut ref2 = cell.try_borrow_mut().unwrap();
723 assert_eq!(ref2.touch(), 1);
724 assert!(cell.try_borrow().is_none());
725 assert!(cell.try_borrow_mut().is_none());
726 }
727
728 {
729 let ref3 = cell.try_borrow().unwrap();
730 let ref4 = cell.try_borrow().unwrap();
731 let ref5 = cell.try_borrow().unwrap();
732 let ref6 = cell.try_borrow().unwrap();
733 assert_eq!(ref3.look(), 1);
734 assert_eq!(ref4.look(), 1);
735 assert_eq!(ref5.look(), 1);
736 assert_eq!(ref6.look(), 1);
737 assert!(cell.try_borrow_mut().is_none());
738 }
739
740 {
741 let mut ref7 = cell.try_borrow_mut().unwrap();
742 assert_eq!(ref7.look(), 1);
743 assert_eq!(ref7.touch(), 2);
744 assert!(cell.try_borrow().is_none());
745 assert!(cell.try_borrow_mut().is_none());
746 }
747
748 {
749 let ref8 = cell.try_borrow().unwrap();
750 assert_eq!(ref8.look(), 2);
751 assert!(cell.try_borrow_mut().is_none());
752 assert!(cell.try_borrow().is_some());
753 }
754 }
755
756 #[derive(Default)]
757 struct ThreeThings {
758 pub thing1: AsyncRefCell<Thing>,
759 pub thing2: AsyncRefCell<Thing>,
760 pub thing3: AsyncRefCell<Thing>,
761 }
762
763 #[test]
764 fn rc_ref_map() {
765 let runtime = tokio::runtime::Builder::new_current_thread()
766 .build()
767 .unwrap();
768 runtime.block_on(async {
769 let three_cells = Rc::new(ThreeThings::default());
770
771 let rc1 = RcRef::map(three_cells.clone(), |things| &things.thing1);
772 let rc2 = RcRef::map(three_cells.clone(), |things| &things.thing2);
773 let rc3 = RcRef::map(three_cells, |things| &things.thing3);
774
775 let mut ref1 = rc1.borrow_mut().await;
776 let ref2 = rc2.borrow().await;
777 let mut ref3 = rc3.borrow_mut().await;
778
779 assert_eq!(ref1.look(), 0);
780 assert_eq!(ref3.touch(), 1);
781 assert_eq!(ref1.touch(), 1);
782 assert_eq!(ref2.look(), 0);
783 assert_eq!(ref3.touch(), 2);
784 assert_eq!(ref1.look(), 1);
785 assert_eq!(ref1.touch(), 2);
786 assert_eq!(ref3.touch(), 3);
787 assert_eq!(ref1.touch(), 3);
788 });
789 }
790}