const-size-flatten 0.5.0

Flatten and FlatMap with constant inner iterator size
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 super::ConstSizeIntoIterator;
use core::iter::Fuse;

#[derive(Clone, Debug)]
pub(crate) struct ConstSizeFlattenBase<I, U> {
    pub(crate) base_iter: Fuse<I>,
    pub(crate) front_sub_iter: Option<U>,
    pub(crate) back_sub_iter: Option<U>,
}

impl<I, U> ConstSizeFlattenBase<I, U>
where
    I: Iterator,
    I::Item: ConstSizeIntoIterator,
    U: Iterator,
{
    fn max_size(&self) -> Option<usize> {
        let base_iter_max_size = self.base_iter.size_hint().1?;
        let inner_iter_len = <I::Item as ConstSizeIntoIterator>::SIZE;
        let front_sub_iter_max_size = self
            .front_sub_iter
            .as_ref()
            .map_or(Some(0), |it| it.size_hint().1)?;
        let back_sub_iter_max_size = self
            .back_sub_iter
            .as_ref()
            .map_or(Some(0), |it| it.size_hint().1)?;
        base_iter_max_size
            .checked_mul(inner_iter_len)?
            .checked_add(front_sub_iter_max_size)?
            .checked_add(back_sub_iter_max_size)
    }
    fn min_size(&self) -> usize {
        let base_iter_min_size = self.base_iter.size_hint().0;
        let inner_iter_len = <I::Item as ConstSizeIntoIterator>::SIZE;
        let front_sub_iter_min_size = self
            .front_sub_iter
            .as_ref()
            .map_or(0, |it| it.size_hint().0);
        let back_sub_iter_min_size = self.back_sub_iter.as_ref().map_or(0, |it| it.size_hint().0);
        base_iter_min_size
            .saturating_mul(inner_iter_len)
            .saturating_add(front_sub_iter_min_size)
            .saturating_add(back_sub_iter_min_size)
    }
}

impl<I: Iterator, U> ConstSizeFlattenBase<I, U> {
    pub(crate) fn new<J: IntoIterator<IntoIter = I>>(iter: J) -> Self {
        Self {
            base_iter: iter.into_iter().fuse(),
            front_sub_iter: None,
            back_sub_iter: None,
        }
    }
}

impl<I, U> Iterator for ConstSizeFlattenBase<I, U>
where
    I: Iterator,
    I::Item: ConstSizeIntoIterator<IntoIter = U>,
    U: Iterator,
{
    type Item = <U as Iterator>::Item;
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let next @ Some(_) = and_then_or_clear(&mut self.front_sub_iter, Iterator::next) {
                return next;
            }
            match self.base_iter.next() {
                None => return and_then_or_clear(&mut self.back_sub_iter, Iterator::next),
                Some(next) => self.front_sub_iter = Some(next.into_iter()),
            }
        }
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.min_size(), self.max_size())
    }
}

impl<I, U> DoubleEndedIterator for ConstSizeFlattenBase<I, U>
where
    I: DoubleEndedIterator,
    I::Item: ConstSizeIntoIterator<IntoIter = U>,
    U: DoubleEndedIterator,
{
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            if let next @ Some(_) = and_then_or_clear(&mut self.back_sub_iter, |b| b.next_back()) {
                return next;
            }
            match self.base_iter.next_back() {
                None => return and_then_or_clear(&mut self.front_sub_iter, |f| f.next_back()),
                Some(next) => self.back_sub_iter = Some(next.into_iter()),
            }
        }
    }
}

fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> {
    let x = f(opt.as_mut()?);
    if x.is_none() {
        *opt = None;
    }
    x
}