diff --git a/src/lib.rs b/src/lib.rs index 5d5e2e3..ffe2c87 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -592,6 +592,19 @@ pub trait StreamingIteratorMut: StreamingIterator { { MapDerefMut { it: self, f } } + + /// Creates an iterator which flattens nested streaming iterators. + #[inline] + fn flatten(self) -> Flatten + where + Self: Sized, + Self::Item: StreamingIterator, + { + Flatten { + iter: self, + first: true, + } + } } impl<'a, I: ?Sized> StreamingIteratorMut for &'a mut I @@ -1268,6 +1281,80 @@ where } } +/// A streaming iterator that flattens nested streaming iterators. +#[derive(Debug)] +pub struct Flatten { + iter: I, + first: bool, +} + +impl StreamingIterator for Flatten +where + I: StreamingIteratorMut, + I::Item: StreamingIterator, +{ + type Item = ::Item; + + #[inline] + fn advance(&mut self) { + if self.first { + self.first = false; + self.iter.advance(); + } + while let Some(iter) = self.iter.get_mut() { + iter.advance(); + if !iter.is_done() { + break; + } + self.iter.advance(); // since we got Some, self.iter is not done and can be advanced + } + } + + #[inline] + fn is_done(&self) -> bool { + match self.iter.get() { + Some(iter) => iter.is_done(), + None => true, + } + } + + #[inline] + fn get(&self) -> Option<&Self::Item> { + self.iter.get().and_then(I::Item::get) + } + + #[inline] + fn fold(self, init: Acc, mut fold: Fold) -> Acc + where + Self: Sized, + Fold: FnMut(Acc, &Self::Item) -> Acc, + { + self.iter + .fold_mut(init, |acc, item| item.fold(acc, &mut fold)) + } +} + +impl StreamingIteratorMut for Flatten +where + I: StreamingIteratorMut, + I::Item: StreamingIteratorMut, +{ + #[inline] + fn get_mut(&mut self) -> Option<&mut Self::Item> { + self.iter.get_mut().and_then(I::Item::get_mut) + } + + #[inline] + fn fold_mut(self, init: Acc, mut fold: Fold) -> Acc + where + Self: Sized, + Fold: FnMut(Acc, &mut Self::Item) -> Acc, + { + self.iter + .fold_mut(init, |acc, item| item.fold_mut(acc, &mut fold)) + } +} + /// A regular, non-streaming iterator which both filters and maps elements of a streaming iterator with a closure. #[derive(Debug)] pub struct FilterMapDeref { @@ -2549,6 +2636,34 @@ mod test { test(it, &[0, 1, 2, 3, 4, 5]); } + #[test] + fn flatten() { + let mut items = [ + convert_ref([].as_ref()), + convert_ref([1].as_ref()), + convert_ref([].as_ref()), + convert_ref([2, 3].as_ref()), + convert_ref([].as_ref()), + ]; + let it = convert_mut(&mut items).flatten(); + + test(it, &[1, 2, 3]); + } + + #[test] + fn flatten_unsized() { + type DynI32 = dyn StreamingIterator; + let mut items = [ + &mut once(1) as &mut DynI32, + &mut empty(), + &mut convert(2..=3), + ]; + let iters = items.iter_mut().map(|iter| &mut **iter); + let it = convert_mut(iters).flatten(); + + test(it, &[1, 2, 3]); + } + #[test] fn nth() { let items = [0, 1];