mtorrent-utils 0.5.2

Miscellaneous utilities for BitTorrent protocol
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

use std::ops::ControlFlow;
use std::pin::Pin;
use std::task::{Context, Poll, ready};
use tokio::task;

/// Type alias for the poll function signature used in [`loop_select`].
pub type LoopSelectPollFn<D, O> = fn(&mut D, &mut Context<'_>) -> Poll<ControlFlow<O>>;

/// Future that runs multiple poll functions in a loop until one of them returns
/// [`ControlFlow::Break`].
pub struct LoopSelectUnpin<'c, D: Unpin, O, const N: usize> {
    data: &'c mut D,
    poll_fns: [LoopSelectPollFn<D, O>; N],
}

impl<'c, D: Unpin, O, const N: usize> Future for LoopSelectUnpin<'c, D, O, N> {
    type Output = O;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            let coop = ready!(task::coop::poll_proceed(cx));
            let mut made_progress = false;
            for poll_fn in self.poll_fns {
                match poll_fn(self.data, cx) {
                    Poll::Ready(ControlFlow::Break(out)) => {
                        return Poll::Ready(out);
                    }
                    Poll::Ready(ControlFlow::Continue(())) => {
                        made_progress = true;
                    }
                    Poll::Pending => {}
                }
            }
            if !made_progress {
                return Poll::Pending;
            }
            coop.made_progress();
        }
    }
}

/// Run multiple poll functions in a loop until one of them returns [`ControlFlow::Break`].
/// All functions have access to a mutable reference to the same data.
///
/// A poll function should return `ControlFlow::Continue(())` if it made progress,
/// or `ControlFlow::Break(output)` to finish the loop with the given output.
pub fn loop_select<Data: Unpin, Ret, const N: usize>(
    data: &mut Data,
    poll_fns: [LoopSelectPollFn<Data, Ret>; N],
) -> LoopSelectUnpin<'_, Data, Ret, N> {
    LoopSelectUnpin { data, poll_fns }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::ops::ControlFlow;
    use std::task::{Context, Poll};

    struct TestCtx {
        counter1: usize,
        counter2: usize,
        target1: usize,
        target2: usize,
    }

    fn poll_fn1(ctx: &mut TestCtx, _cx: &mut Context<'_>) -> Poll<ControlFlow<()>> {
        if ctx.counter1 < ctx.target1 {
            ctx.counter1 += 1;
            Poll::Ready(ControlFlow::Continue(()))
        } else {
            Poll::Ready(ControlFlow::Break(()))
        }
    }
    fn poll_fn2(ctx: &mut TestCtx, _cx: &mut Context<'_>) -> Poll<ControlFlow<()>> {
        if ctx.counter2 < ctx.target2 {
            ctx.counter2 += 1;
            Poll::Ready(ControlFlow::Continue(()))
        } else {
            Poll::Ready(ControlFlow::Break(()))
        }
    }

    #[tokio::test]
    async fn test_loop_select() {
        let mut ctx = TestCtx {
            counter1: 0,
            counter2: 0,
            target1: 5,
            target2: 3,
        };
        loop_select(&mut ctx, [poll_fn1, poll_fn2]).await;
        assert_eq!(ctx.counter1, 4);
        assert_eq!(ctx.counter2, 3);

        let mut ctx = TestCtx {
            counter1: 0,
            counter2: 0,
            target1: 5,
            target2: 3,
        };
        loop_select(&mut ctx, [poll_fn2, poll_fn1]).await;
        assert_eq!(ctx.counter2, 3);
        assert_eq!(ctx.counter1, 3);
    }
}