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
107
108
109
110
111
112
113
114
115

use std::collections::VecDeque;

pub struct RestorableIter<T, I>
where
    T: Iterator<Item = I>,
    I: Clone,
{
    inner: T,
    buf: VecDeque<Option<I>>,
    saving: bool,
}

impl<T, I> RestorableIter<T, I>
where
    T: Iterator<Item = I>,
    I: Clone,
{
    pub fn new(inner: T) -> Self {
        RestorableIter {
            inner,
            buf: VecDeque::with_capacity(4),
            saving: false,
        }
    }

    #[inline]
    pub fn save(&mut self) {
        self.saving = true;
    }

    #[inline]
    pub fn restore(&mut self) {
        self.saving = false;
    }

    #[inline]
    pub fn clear(&mut self) {
        self.buf = VecDeque::with_capacity(4);
        self.saving = false;
    }
}

impl<T, I> Iterator for RestorableIter<T, I>
where
    T: Iterator<Item = I>,
    I: Clone,
{
    type Item = I;
    fn next(&mut self) -> Option<Self::Item> {
        if self.saving {
            let next_item = self.inner.next();
            self.buf.push_back(next_item.clone());
            next_item
        } else {
            match self.buf.pop_front() {
                Some(saved_item) => saved_item,
                None => self.inner.next(),
            }
        }
    }
}

pub trait Restorable<I>: Iterator
where
    I: Clone,
{
    fn iter_restorable(self) -> RestorableIter<Self, I>
    where
        Self: Sized + Iterator<Item = I>,
    {
        RestorableIter::new(self)
    }
}

impl<I, T: ?Sized> Restorable<I> for T
where
    T: Iterator<Item = I>,
    I: Clone,
{
}

#[cfg(test)]
mod tests {
    use crate::*;

    #[test]
    fn test_restorable_new() {
        let mut iterator = RestorableIter::new(0..5);
        assert_eq!(iterator.next(), Some(0));
        assert_eq!(iterator.next(), Some(1));
        assert_eq!(iterator.next(), Some(2));
        iterator.save();
        assert_eq!(iterator.next(), Some(3));
        assert_eq!(iterator.next(), Some(4));
        iterator.restore();
        assert_eq!(iterator.next(), Some(3));
        assert_eq!(iterator.next(), Some(4));
        assert_eq!(iterator.next(), None);
    }

    #[test]
    fn test_iter_restorable() {
        let mut iterator = (0..5).iter_restorable();
        assert_eq!(iterator.next(), Some(0));
        assert_eq!(iterator.next(), Some(1));
        assert_eq!(iterator.next(), Some(2));
        iterator.save();
        assert_eq!(iterator.next(), Some(3));
        assert_eq!(iterator.next(), Some(4));
        iterator.restore();
        assert_eq!(iterator.next(), Some(3));
        assert_eq!(iterator.next(), Some(4));
        assert_eq!(iterator.next(), None);
    }
}