mtorrent_utils/
loop_select.rs

1use std::ops::ControlFlow;
2use std::pin::Pin;
3use std::task::{Context, Poll, ready};
4use tokio::task;
5
6/// Type alias for the poll function signature used in [`loop_select`].
7pub type LoopSelectPollFn<D, O> = fn(&mut D, &mut Context<'_>) -> Poll<ControlFlow<O>>;
8
9/// Future that runs multiple poll functions in a loop until one of them returns
10/// [`ControlFlow::Break`].
11pub struct LoopSelectUnpin<'c, D: Unpin, O, const N: usize> {
12    data: &'c mut D,
13    poll_fns: [LoopSelectPollFn<D, O>; N],
14}
15
16impl<'c, D: Unpin, O, const N: usize> Future for LoopSelectUnpin<'c, D, O, N> {
17    type Output = O;
18
19    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
20        loop {
21            let coop = ready!(task::coop::poll_proceed(cx));
22            let mut made_progress = false;
23            for poll_fn in self.poll_fns {
24                match poll_fn(self.data, cx) {
25                    Poll::Ready(ControlFlow::Break(out)) => {
26                        return Poll::Ready(out);
27                    }
28                    Poll::Ready(ControlFlow::Continue(())) => {
29                        made_progress = true;
30                    }
31                    Poll::Pending => {}
32                }
33            }
34            if !made_progress {
35                return Poll::Pending;
36            }
37            coop.made_progress();
38        }
39    }
40}
41
42/// Run multiple poll functions in a loop until one of them returns [`ControlFlow::Break`].
43/// All functions have access to a mutable reference to the same data.
44///
45/// A poll function should return `ControlFlow::Continue(())` if it made progress,
46/// or `ControlFlow::Break(output)` to finish the loop with the given output.
47pub fn loop_select<Data: Unpin, Ret, const N: usize>(
48    data: &mut Data,
49    poll_fns: [LoopSelectPollFn<Data, Ret>; N],
50) -> LoopSelectUnpin<'_, Data, Ret, N> {
51    LoopSelectUnpin { data, poll_fns }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57    use std::ops::ControlFlow;
58    use std::task::{Context, Poll};
59
60    struct TestCtx {
61        counter1: usize,
62        counter2: usize,
63        target1: usize,
64        target2: usize,
65    }
66
67    fn poll_fn1(ctx: &mut TestCtx, _cx: &mut Context<'_>) -> Poll<ControlFlow<()>> {
68        if ctx.counter1 < ctx.target1 {
69            ctx.counter1 += 1;
70            Poll::Ready(ControlFlow::Continue(()))
71        } else {
72            Poll::Ready(ControlFlow::Break(()))
73        }
74    }
75    fn poll_fn2(ctx: &mut TestCtx, _cx: &mut Context<'_>) -> Poll<ControlFlow<()>> {
76        if ctx.counter2 < ctx.target2 {
77            ctx.counter2 += 1;
78            Poll::Ready(ControlFlow::Continue(()))
79        } else {
80            Poll::Ready(ControlFlow::Break(()))
81        }
82    }
83
84    #[tokio::test]
85    async fn test_loop_select() {
86        let mut ctx = TestCtx {
87            counter1: 0,
88            counter2: 0,
89            target1: 5,
90            target2: 3,
91        };
92        loop_select(&mut ctx, [poll_fn1, poll_fn2]).await;
93        assert_eq!(ctx.counter1, 4);
94        assert_eq!(ctx.counter2, 3);
95
96        let mut ctx = TestCtx {
97            counter1: 0,
98            counter2: 0,
99            target1: 5,
100            target2: 3,
101        };
102        loop_select(&mut ctx, [poll_fn2, poll_fn1]).await;
103        assert_eq!(ctx.counter2, 3);
104        assert_eq!(ctx.counter1, 3);
105    }
106}
mtorrent_utils/loop_select.rs

mtorrent_utils/
loop_select.rs
