modio/util/upload/
byte_ranges.rs

1use std::iter::{repeat, Chain, Map, Repeat, Zip};
2use std::ops::Range;
3
4use super::MULTIPART_FILE_PART_SIZE;
5
6/// The [`ByteRanges`] type implements the underlying `Iterator` for [`byte_ranges`].
7///
8/// # Example
9///
10/// ```
11/// # use modio::util::upload::ByteRanges;
12/// let mut iter = ByteRanges::new(14).chunks(5).into_iter();
13///
14/// assert_eq!(iter.next(), Some((0, 4)));
15/// assert_eq!(iter.next(), Some((5, 9)));
16/// assert_eq!(iter.next(), Some((10, 13)));
17/// assert_eq!(iter.next(), None);
18/// ```
19pub struct ByteRanges {
20    length: u64,
21    chunk_size: u64,
22}
23
24impl ByteRanges {
25    /// Creates a new `ByteRanges` for the give length.
26    pub const fn new(length: u64) -> Self {
27        Self {
28            length,
29            chunk_size: MULTIPART_FILE_PART_SIZE,
30        }
31    }
32
33    /// Sets the chunk size for the iterator. Defaults to [`MULTIPART_FILE_PART_SIZE`].
34    pub const fn chunks(mut self, chunk_size: u64) -> Self {
35        self.chunk_size = chunk_size;
36        self
37    }
38}
39
40impl IntoIterator for ByteRanges {
41    type Item = (u64, u64);
42    type IntoIter = IntoIter;
43
44    fn into_iter(self) -> Self::IntoIter {
45        fn map_fn((i, chunk_size): (u64, u64)) -> (u64, u64) {
46            (i * chunk_size, (i + 1) * chunk_size - 1)
47        }
48
49        let ByteRanges { length, chunk_size } = self;
50
51        let count = length / chunk_size;
52        let rem = length % chunk_size;
53
54        let iter = (0..count)
55            .zip(repeat(chunk_size))
56            .map(map_fn as MapFn)
57            .chain(Some((count * chunk_size, count * chunk_size + rem - 1)));
58
59        IntoIter { inner: iter }
60    }
61}
62
63type MapFn = fn((u64, u64)) -> (u64, u64);
64type Elements = Map<Zip<Range<u64>, Repeat<u64>>, MapFn>;
65type Last = std::option::IntoIter<(u64, u64)>;
66
67pub struct IntoIter {
68    inner: Chain<Elements, Last>,
69}
70
71impl Iterator for IntoIter {
72    type Item = (u64, u64);
73
74    fn next(&mut self) -> Option<Self::Item> {
75        self.inner.next()
76    }
77}
78
79/// Calculates bytes ranges for the given `length` parameter and the chunk size of 50MB.
80///
81/// # Example
82///
83/// ```
84/// # use modio::util::upload::byte_ranges;
85/// let mut iter = byte_ranges(52 * 1024 * 1024);
86///
87/// assert_eq!(iter.next(), Some((0, 52428799)));
88/// assert_eq!(iter.next(), Some((52428800, 54525951)));
89/// assert_eq!(iter.next(), None);
90/// ```
91pub fn byte_ranges(length: u64) -> impl Iterator<Item = (u64, u64)> {
92    ByteRanges::new(length).into_iter()
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    pub fn multiple_byte_ranges() {
101        const SIZE: u64 = 522 * 1024 * 1024;
102
103        let ranges: Vec<_> = byte_ranges(SIZE).collect();
104
105        assert_eq!(
106            ranges,
107            [
108                (0, 52428799),
109                (52428800, 104857599),
110                (104857600, 157286399),
111                (157286400, 209715199),
112                (209715200, 262143999),
113                (262144000, 314572799),
114                (314572800, 367001599),
115                (367001600, 419430399),
116                (419430400, 471859199),
117                (471859200, 524287999),
118                (524288000, 547356671),
119            ]
120        );
121    }
122
123    #[test]
124    pub fn single_byte_range() {
125        const SIZE: u64 = 25 * 1024 * 1024;
126
127        let ranges: Vec<_> = byte_ranges(SIZE).collect();
128
129        assert_eq!(ranges, [(0, SIZE - 1)]);
130    }
131}