use std::any::type_name;
use std::any::Any;
use std::borrow::Borrow;
use std::cell::Cell;
use std::cell::UnsafeCell;
use std::collections::VecDeque;
use std::fmt;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::ops::Deref;
use std::rc::Rc;
use self::internal as i;
pub type AsyncRef<T> = i::AsyncBorrowImpl<T, i::Shared>;
pub type AsyncMut<T> = i::AsyncBorrowImpl<T, i::Exclusive>;
pub type AsyncRefFuture<T> = i::AsyncBorrowFutureImpl<T, i::Shared>;
pub type AsyncMutFuture<T> = i::AsyncBorrowFutureImpl<T, i::Exclusive>;
pub struct AsyncRefCell<T> {
value: UnsafeCell<T>,
borrow_count: Cell<i::BorrowCount>,
waiters: Cell<VecDeque<Option<i::Waiter>>>,
turn: Cell<usize>,
}
impl<T: 'static> AsyncRefCell<T> {
pub fn new(value: T) -> Self {
Self {
value: UnsafeCell::new(value),
borrow_count: Default::default(),
waiters: Default::default(),
turn: Default::default(),
}
}
pub fn new_rc(value: T) -> Rc<Self> {
Rc::new(Self::new(value))
}
pub fn as_ptr(&self) -> *mut T {
self.value.get()
}
pub fn into_inner(self) -> T {
assert!(self.borrow_count.get().is_empty());
self.value.into_inner()
}
}
impl<T> Debug for AsyncRefCell<T> {
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
write!(f, "AsyncRefCell<{}>", type_name::<T>())
}
}
impl<T: Default + 'static> Default for AsyncRefCell<T> {
fn default() -> Self {
Self::new(Default::default())
}
}
impl<T: Default + 'static> AsyncRefCell<T> {
pub fn default_rc() -> Rc<Self> {
Rc::new(Default::default())
}
}
impl<T: 'static> From<T> for AsyncRefCell<T> {
fn from(value: T) -> Self {
Self::new(value)
}
}
impl<T> AsyncRefCell<T> {
pub fn borrow(self: &Rc<Self>) -> AsyncRefFuture<T> {
AsyncRefFuture::new(self)
}
pub fn borrow_mut(self: &Rc<Self>) -> AsyncMutFuture<T> {
AsyncMutFuture::new(self)
}
pub fn try_borrow(self: &Rc<Self>) -> Option<AsyncRef<T>> {
Self::borrow_sync(self)
}
pub fn try_borrow_mut(self: &Rc<Self>) -> Option<AsyncMut<T>> {
Self::borrow_sync(self)
}
}
impl<T> RcRef<AsyncRefCell<T>> {
pub fn borrow(&self) -> AsyncRefFuture<T> {
AsyncRefFuture::new(self)
}
pub fn borrow_mut(&self) -> AsyncMutFuture<T> {
AsyncMutFuture::new(self)
}
pub fn try_borrow(&self) -> Option<AsyncRef<T>> {
AsyncRefCell::<T>::borrow_sync(self)
}
pub fn try_borrow_mut(&self) -> Option<AsyncMut<T>> {
AsyncRefCell::<T>::borrow_sync(self)
}
}
#[derive(Debug)]
pub struct RcRef<T> {
rc: Rc<dyn Any>,
value: *const T,
}
impl<T: 'static> RcRef<T> {
pub fn new(value: T) -> Self {
Self::from(Rc::new(value))
}
pub fn map<S: 'static, R: RcLike<S>, F: FnOnce(&S) -> &T>(
source: R,
map_fn: F,
) -> RcRef<T> {
let RcRef::<S> { rc, value } = source.into();
#[allow(clippy::undocumented_unsafe_blocks)]
let value = map_fn(unsafe { &*value });
RcRef { rc, value }
}
pub(crate) fn split(rc_ref: &Self) -> (&T, &Rc<dyn Any>) {
let &Self { ref rc, value } = rc_ref;
#[allow(clippy::undocumented_unsafe_blocks)]
(unsafe { &*value }, rc)
}
}
impl<T: Default + 'static> Default for RcRef<T> {
fn default() -> Self {
Self::new(Default::default())
}
}
impl<T> Clone for RcRef<T> {
fn clone(&self) -> Self {
Self {
rc: self.rc.clone(),
value: self.value,
}
}
}
impl<T: 'static> From<&RcRef<T>> for RcRef<T> {
fn from(rc_ref: &RcRef<T>) -> Self {
rc_ref.clone()
}
}
impl<T: 'static> From<Rc<T>> for RcRef<T> {
fn from(rc: Rc<T>) -> Self {
Self {
value: &*rc,
rc: rc as Rc<_>,
}
}
}
impl<T: 'static> From<&Rc<T>> for RcRef<T> {
fn from(rc: &Rc<T>) -> Self {
rc.clone().into()
}
}
impl<T> Deref for RcRef<T> {
type Target = T;
fn deref(&self) -> &Self::Target {
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe {
&*self.value
}
}
}
impl<T> Borrow<T> for RcRef<T> {
fn borrow(&self) -> &T {
self
}
}
impl<T> AsRef<T> for RcRef<T> {
fn as_ref(&self) -> &T {
self
}
}
pub trait RcLike<T>: AsRef<T> + Into<RcRef<T>> {}
impl<T: 'static> RcLike<T> for Rc<T> {}
impl<T: 'static> RcLike<T> for RcRef<T> {}
impl<T: 'static> RcLike<T> for &Rc<T> {}
impl<T: 'static> RcLike<T> for &RcRef<T> {}
mod internal {
use super::AsyncRefCell;
use super::RcLike;
use super::RcRef;
use futures::future::Future;
use futures::ready;
use futures::task::Context;
use futures::task::Poll;
use futures::task::Waker;
use std::borrow::Borrow;
use std::borrow::BorrowMut;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::ops::Deref;
use std::ops::DerefMut;
use std::pin::Pin;
impl<T> AsyncRefCell<T> {
pub fn borrow_sync<M: BorrowModeTrait, R: RcLike<AsyncRefCell<T>>>(
cell: R,
) -> Option<AsyncBorrowImpl<T, M>> {
let cell_ref = cell.as_ref();
#[allow(clippy::undocumented_unsafe_blocks)]
let waiters = unsafe { &mut *cell_ref.waiters.as_ptr() };
if waiters.is_empty() {
let new_borrow_count =
cell_ref.borrow_count.get().try_add(M::borrow_mode())?;
cell_ref.borrow_count.set(new_borrow_count);
Some(AsyncBorrowImpl::<T, M>::new(cell.into()))
} else {
None
}
}
fn drop_borrow<M: BorrowModeTrait>(&self) {
let new_borrow_count = self.borrow_count.get().remove(M::borrow_mode());
self.borrow_count.set(new_borrow_count);
if new_borrow_count.is_empty() {
self.wake_waiters()
}
}
fn create_waiter<M: BorrowModeTrait>(&self) -> usize {
let waiter = Waiter::new(M::borrow_mode());
let turn = self.turn.get();
let index = {
#[allow(clippy::undocumented_unsafe_blocks)]
let waiters = unsafe { &mut *self.waiters.as_ptr() };
waiters.push_back(Some(waiter));
waiters.len() - 1
};
if index == 0 {
self.wake_waiters()
}
turn + index
}
fn poll_waiter<M: BorrowModeTrait>(
&self,
id: usize,
cx: &mut Context,
) -> Poll<()> {
let borrow_count = self.borrow_count.get();
let turn = self.turn.get();
if id < turn {
let _ = borrow_count.remove(M::borrow_mode());
Poll::Ready(())
} else {
#[allow(clippy::undocumented_unsafe_blocks)]
let waiters = unsafe { &mut *self.waiters.as_ptr() };
assert!(id < turn + waiters.len());
assert!(id > turn || borrow_count.try_add(M::borrow_mode()).is_none());
let waiter_mut = waiters[id - turn].as_mut().unwrap();
waiter_mut.set_waker(cx.waker());
Poll::Pending
}
}
fn wake_waiters(&self) {
let mut borrow_count = self.borrow_count.get();
#[allow(clippy::undocumented_unsafe_blocks)]
let waiters = unsafe { &mut *self.waiters.as_ptr() };
let mut turn = self.turn.get();
loop {
let waiter_entry = match waiters.front().map(Option::as_ref) {
None => break, Some(w) => w,
};
let borrow_mode = match waiter_entry {
None => {
waiters.pop_front();
turn += 1;
continue;
}
Some(waiter) => waiter.borrow_mode(),
};
borrow_count = match borrow_count.try_add(borrow_mode) {
None => break, Some(b) => b,
};
let mut waiter = waiters.pop_front().unwrap().unwrap();
turn += 1;
if let Some(waker) = waiter.take_waker() {
waker.wake()
}
}
self.borrow_count.set(borrow_count);
self.turn.set(turn);
}
fn drop_waiter<M: BorrowModeTrait>(&self, id: usize) {
let turn = self.turn.get();
if id < turn {
self.drop_borrow::<M>();
} else {
#[allow(clippy::undocumented_unsafe_blocks)]
let waiters = unsafe { &mut *self.waiters.as_ptr() };
waiters[id - turn].take().unwrap();
}
if id == turn {
self.wake_waiters()
}
}
}
pub struct AsyncBorrowFutureImpl<T: 'static, M: BorrowModeTrait> {
cell: Option<RcRef<AsyncRefCell<T>>>,
id: usize,
_phantom: PhantomData<M>,
}
impl<T, M: BorrowModeTrait> AsyncBorrowFutureImpl<T, M> {
pub fn new<R: RcLike<AsyncRefCell<T>>>(cell: R) -> Self {
Self {
id: cell.as_ref().create_waiter::<M>(),
cell: Some(cell.into()),
_phantom: PhantomData,
}
}
}
impl<T: 'static, M: BorrowModeTrait> Future for AsyncBorrowFutureImpl<T, M> {
type Output = AsyncBorrowImpl<T, M>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
ready!(self.cell.as_ref().unwrap().poll_waiter::<M>(self.id, cx));
#[allow(clippy::undocumented_unsafe_blocks)]
let self_mut = unsafe { Pin::get_unchecked_mut(self) };
let cell = self_mut.cell.take().unwrap();
Poll::Ready(AsyncBorrowImpl::<T, M>::new(cell))
}
}
impl<T, M: BorrowModeTrait> Drop for AsyncBorrowFutureImpl<T, M> {
fn drop(&mut self) {
if let Some(cell) = self.cell.take() {
cell.drop_waiter::<M>(self.id)
}
}
}
pub struct AsyncBorrowImpl<T: 'static, M: BorrowModeTrait> {
cell: RcRef<AsyncRefCell<T>>,
_phantom: PhantomData<M>,
}
impl<T, M: BorrowModeTrait> AsyncBorrowImpl<T, M> {
fn new(cell: RcRef<AsyncRefCell<T>>) -> Self {
Self {
cell,
_phantom: PhantomData,
}
}
}
impl<T, M: BorrowModeTrait> Deref for AsyncBorrowImpl<T, M> {
type Target = T;
fn deref(&self) -> &Self::Target {
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe {
&*self.cell.as_ptr()
}
}
}
impl<T, M: BorrowModeTrait> Borrow<T> for AsyncBorrowImpl<T, M> {
fn borrow(&self) -> &T {
self
}
}
impl<T, M: BorrowModeTrait> AsRef<T> for AsyncBorrowImpl<T, M> {
fn as_ref(&self) -> &T {
self
}
}
impl<T> DerefMut for AsyncBorrowImpl<T, Exclusive> {
fn deref_mut(&mut self) -> &mut Self::Target {
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe {
&mut *self.cell.as_ptr()
}
}
}
impl<T> BorrowMut<T> for AsyncBorrowImpl<T, Exclusive> {
fn borrow_mut(&mut self) -> &mut T {
self
}
}
impl<T> AsMut<T> for AsyncBorrowImpl<T, Exclusive> {
fn as_mut(&mut self) -> &mut T {
self
}
}
impl<T, M: BorrowModeTrait> Drop for AsyncBorrowImpl<T, M> {
fn drop(&mut self) {
self.cell.drop_borrow::<M>()
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum BorrowMode {
Shared,
Exclusive,
}
pub trait BorrowModeTrait: Copy {
fn borrow_mode() -> BorrowMode;
}
#[derive(Copy, Clone, Debug)]
pub struct Shared;
impl BorrowModeTrait for Shared {
fn borrow_mode() -> BorrowMode {
BorrowMode::Shared
}
}
#[derive(Copy, Clone, Debug)]
pub struct Exclusive;
impl BorrowModeTrait for Exclusive {
fn borrow_mode() -> BorrowMode {
BorrowMode::Exclusive
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum BorrowCount {
Shared(usize),
Exclusive,
}
impl Default for BorrowCount {
fn default() -> Self {
Self::Shared(0)
}
}
impl BorrowCount {
pub fn is_empty(self) -> bool {
matches!(self, BorrowCount::Shared(0))
}
pub fn try_add(self, mode: BorrowMode) -> Option<BorrowCount> {
match (self, mode) {
(BorrowCount::Shared(refs), BorrowMode::Shared) => {
Some(BorrowCount::Shared(refs + 1))
}
(BorrowCount::Shared(0), BorrowMode::Exclusive) => {
Some(BorrowCount::Exclusive)
}
_ => None,
}
}
#[allow(dead_code)]
pub fn add(self, mode: BorrowMode) -> BorrowCount {
match self.try_add(mode) {
Some(value) => value,
None => panic!("Can't add {mode:?} to {self:?}"),
}
}
pub fn try_remove(self, mode: BorrowMode) -> Option<BorrowCount> {
match (self, mode) {
(BorrowCount::Shared(refs), BorrowMode::Shared) if refs > 0 => {
Some(BorrowCount::Shared(refs - 1))
}
(BorrowCount::Exclusive, BorrowMode::Exclusive) => {
Some(BorrowCount::Shared(0))
}
_ => None,
}
}
pub fn remove(self, mode: BorrowMode) -> BorrowCount {
match self.try_remove(mode) {
Some(value) => value,
None => panic!("Can't remove {mode:?} from {self:?}"),
}
}
}
pub struct Waiter {
borrow_mode: BorrowMode,
waker: Option<Waker>,
}
impl Waiter {
pub fn new(borrow_mode: BorrowMode) -> Self {
Self {
borrow_mode,
waker: None,
}
}
pub fn borrow_mode(&self) -> BorrowMode {
self.borrow_mode
}
pub fn set_waker(&mut self, new_waker: &Waker) {
if self
.waker
.as_ref()
.filter(|waker| waker.will_wake(new_waker))
.is_none()
{
self.waker.replace(new_waker.clone());
}
}
pub fn take_waker(&mut self) -> Option<Waker> {
self.waker.take()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Default)]
struct Thing {
touch_count: usize,
_private: (),
}
impl Thing {
pub fn look(&self) -> usize {
self.touch_count
}
pub fn touch(&mut self) -> usize {
self.touch_count += 1;
self.touch_count
}
}
#[test]
fn async_ref_cell_borrow() {
let runtime = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
runtime.block_on(async {
let cell = AsyncRefCell::<Thing>::default_rc();
let fut1 = cell.borrow();
let fut2 = cell.borrow_mut();
let fut3 = cell.borrow();
let fut4 = cell.borrow();
let fut5 = cell.borrow();
let fut6 = cell.borrow();
let fut7 = cell.borrow_mut();
let fut8 = cell.borrow();
assert!(cell.try_borrow().is_none());
assert!(cell.try_borrow_mut().is_none());
assert_eq!(fut1.await.look(), 0);
assert_eq!(fut2.await.touch(), 1);
{
let ref5 = fut5.await;
let ref4 = fut4.await;
let ref3 = fut3.await;
let ref6 = fut6.await;
assert_eq!(ref3.look(), 1);
assert_eq!(ref4.look(), 1);
assert_eq!(ref5.look(), 1);
assert_eq!(ref6.look(), 1);
}
{
let mut ref7 = fut7.await;
assert_eq!(ref7.look(), 1);
assert_eq!(ref7.touch(), 2);
}
{
let ref8 = fut8.await;
assert_eq!(ref8.look(), 2);
}
});
}
#[test]
fn async_ref_cell_try_borrow() {
let cell = AsyncRefCell::<Thing>::default_rc();
{
let ref1 = cell.try_borrow().unwrap();
assert_eq!(ref1.look(), 0);
assert!(cell.try_borrow_mut().is_none());
}
{
let mut ref2 = cell.try_borrow_mut().unwrap();
assert_eq!(ref2.touch(), 1);
assert!(cell.try_borrow().is_none());
assert!(cell.try_borrow_mut().is_none());
}
{
let ref3 = cell.try_borrow().unwrap();
let ref4 = cell.try_borrow().unwrap();
let ref5 = cell.try_borrow().unwrap();
let ref6 = cell.try_borrow().unwrap();
assert_eq!(ref3.look(), 1);
assert_eq!(ref4.look(), 1);
assert_eq!(ref5.look(), 1);
assert_eq!(ref6.look(), 1);
assert!(cell.try_borrow_mut().is_none());
}
{
let mut ref7 = cell.try_borrow_mut().unwrap();
assert_eq!(ref7.look(), 1);
assert_eq!(ref7.touch(), 2);
assert!(cell.try_borrow().is_none());
assert!(cell.try_borrow_mut().is_none());
}
{
let ref8 = cell.try_borrow().unwrap();
assert_eq!(ref8.look(), 2);
assert!(cell.try_borrow_mut().is_none());
assert!(cell.try_borrow().is_some());
}
}
#[derive(Default)]
struct ThreeThings {
pub thing1: AsyncRefCell<Thing>,
pub thing2: AsyncRefCell<Thing>,
pub thing3: AsyncRefCell<Thing>,
}
#[test]
fn rc_ref_map() {
let runtime = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
runtime.block_on(async {
let three_cells = Rc::new(ThreeThings::default());
let rc1 = RcRef::map(three_cells.clone(), |things| &things.thing1);
let rc2 = RcRef::map(three_cells.clone(), |things| &things.thing2);
let rc3 = RcRef::map(three_cells, |things| &things.thing3);
let mut ref1 = rc1.borrow_mut().await;
let ref2 = rc2.borrow().await;
let mut ref3 = rc3.borrow_mut().await;
assert_eq!(ref1.look(), 0);
assert_eq!(ref3.touch(), 1);
assert_eq!(ref1.touch(), 1);
assert_eq!(ref2.look(), 0);
assert_eq!(ref3.touch(), 2);
assert_eq!(ref1.look(), 1);
assert_eq!(ref1.touch(), 2);
assert_eq!(ref3.touch(), 3);
assert_eq!(ref1.touch(), 3);
});
}
}