Skip to main content

object_rainbow/
nested_mut.rs

1use std::{
2    ops::{Deref, DerefMut},
3    pin::Pin,
4    task::{Context, Poll, ready},
5};
6
7use futures_channel::oneshot;
8
9use crate::FailFuture;
10
11struct NestedGuard<'a, T> {
12    original: &'a mut T,
13    returned: oneshot::Receiver<T>,
14}
15
16impl<T> Drop for NestedGuard<'_, T> {
17    fn drop(&mut self) {
18        if let Ok(Some(returned)) = self.returned.try_recv() {
19            *self.original = returned;
20        }
21    }
22}
23
24struct Lent<T> {
25    value: T,
26    return_to: oneshot::Sender<T>,
27}
28
29impl<T> Deref for Lent<T> {
30    type Target = T;
31
32    fn deref(&self) -> &Self::Target {
33        &self.value
34    }
35}
36
37impl<T> DerefMut for Lent<T> {
38    fn deref_mut(&mut self) -> &mut Self::Target {
39        &mut self.value
40    }
41}
42
43impl<T> Lent<T> {
44    fn finish(self) {
45        self.return_to.send(self.value).ok();
46    }
47}
48
49pub struct Borrower<T>(oneshot::Sender<Lent<T>>);
50
51impl<'a, T: Clone> NestedGuard<'a, T> {
52    fn new(original: &'a mut T, borrower: Borrower<T>) -> Self {
53        let (return_to, returned) = oneshot::channel();
54        borrower
55            .0
56            .send(Lent {
57                value: original.clone(),
58                return_to,
59            })
60            .ok();
61        Self { original, returned }
62    }
63}
64
65pub trait LendTo: Clone {
66    fn lend_to<T>(&mut self, borrower: Borrower<Self>) -> impl Future<Output = T> {
67        async move {
68            let _guard = NestedGuard::new(self, borrower);
69            std::future::pending().await
70        }
71    }
72}
73
74impl<T: Clone> LendTo for T {}
75
76pub struct NestedMut<'a, T> {
77    lent: Option<Lent<T>>,
78    _guard: oneshot::Receiver<FailFuture<'a, ()>>,
79}
80
81impl<T> Deref for NestedMut<'_, T> {
82    type Target = T;
83
84    fn deref(&self) -> &Self::Target {
85        self.lent.as_ref().expect("invalid state")
86    }
87}
88
89impl<T> DerefMut for NestedMut<'_, T> {
90    fn deref_mut(&mut self) -> &mut Self::Target {
91        self.lent.as_mut().expect("invalid state")
92    }
93}
94
95impl<T> Drop for NestedMut<'_, T> {
96    fn drop(&mut self) {
97        self.lent.take().expect("invalid state").finish();
98    }
99}
100
101struct WaitingLease<'a, T> {
102    borrowing: oneshot::Receiver<Lent<T>>,
103    future: Option<FailFuture<'a, ()>>,
104}
105
106impl<'a, T> Future for WaitingLease<'a, T> {
107    type Output = object_rainbow::Result<Option<NestedMut<'a, T>>>;
108
109    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
110        let this = self.get_mut();
111        if this
112            .future
113            .as_mut()
114            .expect("invalid state")
115            .as_mut()
116            .poll(cx)?
117            .is_ready()
118        {
119            Poll::Ready(Ok(None))
120        } else {
121            let Ok(lent) = ready!(Pin::new(&mut this.borrowing).poll(cx)) else {
122                return Poll::Ready(Ok(None));
123            };
124            Poll::Ready(Ok(Some(NestedMut::new(
125                lent,
126                this.future.take().expect("invalid state"),
127            ))))
128        }
129    }
130}
131
132impl<'a, T> NestedMut<'a, T> {
133    fn new(lent: Lent<T>, future: FailFuture<'a, ()>) -> Self {
134        let (send, recv) = oneshot::channel();
135        send.send(future).ok();
136        Self {
137            lent: Some(lent),
138            _guard: recv,
139        }
140    }
141
142    pub async fn from_fn<F: 'a + Send + Future<Output = object_rainbow::Result<()>>>(
143        f: impl FnOnce(Borrower<T>) -> F,
144    ) -> object_rainbow::Result<Option<Self>> {
145        let (lending, borrowing) = oneshot::channel();
146        let future = Box::pin(f(Borrower(lending)));
147        WaitingLease {
148            borrowing,
149            future: Some(future),
150        }
151        .await
152    }
153}