v6/src/generator/util/one_more.rs

33 lines
697 B
Rust

pub struct OneMore<I, P: Iterator<Item = I>> {
parent: P,
last: Option<I>,
done: bool,
}
impl<I, P: Iterator<Item = I>> OneMore<I, P> {
pub fn new(parent: P, last: I) -> Self {
Self {
parent,
last: Some(last),
done: false,
}
}
}
impl<I, P: Iterator<Item = I>> Iterator for OneMore<I, P> {
type Item = I;
fn next(&mut self) -> Option<Self::Item> {
if self.done {
return None;
}
match self.parent.next() {
Some(item) => Some(item),
None => {
self.done = true;
Some(self.last.take().unwrap())
}
}
}
}