foyer_common/
future.rs

1// Copyright 2025 foyer Project Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::{
16    future::Future,
17    marker::PhantomData,
18    ops::Deref,
19    pin::Pin,
20    task::{ready, Context, Poll},
21};
22
23use pin_project::pin_project;
24
25/// Result that the inner future of a [`DiversionFuture`] should return.
26///
27/// - The `target` will be further returned by [`DiversionFuture`].
28/// - The `store` will be stored in the [`DiversionFuture`].
29pub struct Diversion<T, S> {
30    /// The `target` will be further returned by [`DiversionFuture`].
31    pub target: T,
32    /// The `store` will be stored in the [`DiversionFuture`].
33    pub store: Option<S>,
34}
35
36impl<T, S> From<T> for Diversion<T, S> {
37    fn from(value: T) -> Self {
38        Self {
39            target: value,
40            store: None,
41        }
42    }
43}
44
45/// [`DiversionFuture`] is a future wrapper that partially store and partially return the future result.
46#[must_use]
47#[pin_project]
48pub struct DiversionFuture<FU, T, S> {
49    #[pin]
50    inner: FU,
51    store: Option<S>,
52    _marker: PhantomData<T>,
53}
54
55impl<FU, T, S> DiversionFuture<FU, T, S> {
56    /// Create a new [`DiversionFuture`] wrapper.
57    pub fn new(future: FU) -> Self {
58        Self {
59            inner: future,
60            store: None,
61            _marker: PhantomData,
62        }
63    }
64
65    /// Get the stored state.
66    pub fn store(&self) -> &Option<S> {
67        &self.store
68    }
69}
70
71impl<FU, T, S> Deref for DiversionFuture<FU, T, S> {
72    type Target = FU;
73
74    fn deref(&self) -> &Self::Target {
75        &self.inner
76    }
77}
78
79impl<FU, T, S, I> Future for DiversionFuture<FU, T, S>
80where
81    FU: Future<Output = I>,
82    I: Into<Diversion<T, S>>,
83{
84    type Output = T;
85
86    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
87        let this = self.project();
88        let Diversion { target, store } = ready!(this.inner.poll(cx)).into();
89        *this.store = store;
90        Poll::Ready(target)
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use std::{future::poll_fn, pin::pin};
97
98    use super::*;
99
100    #[tokio::test]
101    async fn test_diversion_future() {
102        let mut f = pin!(DiversionFuture::new(async move {
103            Diversion {
104                target: "The answer to life, the universe, and everything.".to_string(),
105                store: Some(42),
106            }
107        },));
108
109        let question: String = poll_fn(|cx| f.as_mut().poll(cx)).await;
110        let answer = f.store().unwrap();
111
112        assert_eq!(
113            (question.as_str(), answer),
114            ("The answer to life, the universe, and everything.", 42)
115        );
116    }
117}
foyer_common/future.rs

foyer_common/
future.rs
