diff --git a/src/adaptors.rs b/src/adaptors.rs index 7a7792758..538583d8a 100644 --- a/src/adaptors.rs +++ b/src/adaptors.rs @@ -70,6 +70,84 @@ impl Iterator for Interleave where } } +/// An iterator adaptor that alternates elements from the two iterators until +/// one of them runs out. +/// +/// This iterator is *fused*. +/// +/// See [*.interleave_shortest()*](trait.Itertools.html#method.interleave_shortest) +/// for more information. +#[derive(Clone)] +pub struct InterleaveShortest where + I: Iterator, + J: Iterator, +{ + it0: Fuse, + it1: Fuse, + phase: bool, // false ==> it0, true ==> it1 +} + +impl InterleaveShortest where + I: Iterator, + J: Iterator, +{ + /// Create a new **InterleaveShortest** iterator. + pub fn new(a: I, b: J) -> InterleaveShortest { + InterleaveShortest { + it0: a.fuse(), + it1: b.fuse(), + phase: false, + } + } +} + +impl Iterator for InterleaveShortest where + I: Iterator, + J: Iterator, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option { + match self.phase { + false => match self.it0.next() { + None => None, + e => { + self.phase = true; + e + } + }, + true => match self.it1.next() { + None => None, + e => { + self.phase = false; + e + } + }, + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + fn bound(a: usize, b: usize) -> Option { + use std::cmp::min; + 2usize.checked_mul(min(a, b)) + .and_then(|lhs| lhs.checked_add(if a > b { 1 } else { 0 })) + } + + let (l0, u0) = self.it0.size_hint(); + let (l1, u1) = self.it1.size_hint(); + let lb = bound(l0, l1).unwrap_or(usize::max_value()); + let ub = match (u0, u1) { + (None, None) => None, + (Some(u0), None) => 2usize.checked_mul(u0), + (None, Some(u1)) => 2usize.checked_mul(u1).and_then(|l| l.checked_add(1)), + (Some(u0), Some(u1)) => bound(u0, u1) + }; + (lb, ub) + } +} + /// **Deprecated:** Use *.map_fn()* instead. pub struct FnMap where I: Iterator, diff --git a/src/lib.rs b/src/lib.rs index eb7a6ccb1..565e8e61d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -39,6 +39,7 @@ use std::fmt; pub use adaptors::{ Interleave, + InterleaveShortest, Product, PutBack, PutBackN, @@ -57,6 +58,7 @@ pub use adaptors::{ pub use adaptors::EnumerateFrom; pub use intersperse::Intersperse; pub use islice::{ISlice}; +pub use pad_tail::PadTailUsing; pub use repeatn::RepeatN; pub use rciter::RcIter; pub use stride::Stride; @@ -77,6 +79,7 @@ mod intersperse; mod islice; mod linspace; pub mod misc; +mod pad_tail; mod rciter; mod repeatn; mod sources; @@ -238,6 +241,25 @@ pub trait Itertools : Iterator { Interleave::new(self, other.into_iter()) } + /// Alternate elements from two iterators until one of them runs out. + /// + /// Iterator element type is **Self::Item**. + /// + /// This iterator is *fused*. + /// + /// ``` + /// use itertools::Itertools; + /// + /// let it = (0..5).interleave_shortest(vec![7, 8]); + /// itertools::assert_equal(it, vec![0, 7, 1, 8, 2]); + /// ``` + fn interleave_shortest(self, other: J) -> InterleaveShortest where + J: IntoIterator, + Self: Sized + { + InterleaveShortest::new(self, other.into_iter()) + } + /// An iterator adaptor to insert a particular value /// between each element of the adapted iterator. /// @@ -668,6 +690,32 @@ pub trait Itertools : Iterator { Combinations::new(self) } + /** + Return an iterator adaptor that pads the sequence to a minimum length of + **min** by filling missing elements using a closure **f**. + + Iterator element type is **Self::Item**. + + ``` + use itertools::Itertools; + + let it = (0..5).pad_tail_using(10, |i| 2*i); + itertools::assert_equal(it, vec![0, 1, 2, 3, 4, 10, 12, 14, 16, 18]); + + let it = (0..10).pad_tail_using(5, |i| 2*i); + itertools::assert_equal(it, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + + let it = (0..5).pad_tail_using(10, |i| 2*i).rev(); + itertools::assert_equal(it, vec![18, 16, 14, 12, 10, 4, 3, 2, 1, 0]); + ``` + */ + fn pad_tail_using(self, min: usize, f: F) -> PadTailUsing where + Self: Sized, + F: FnMut(usize) -> Self::Item, + { + PadTailUsing::new(self, min, f) + } + /// Like regular *.map()*, specialized to using a simple function pointer instead, /// so that the resulting **Map** iterator value can be cloned. /// diff --git a/src/pad_tail.rs b/src/pad_tail.rs new file mode 100644 index 000000000..62875d86d --- /dev/null +++ b/src/pad_tail.rs @@ -0,0 +1,74 @@ +/// An iterator adaptor that pads a sequence to a minimum length by filling +/// missing elements using a closure. +/// +/// Iterator element type is **I::Item**. +/// +/// See [*.pad_tail_using()*](trait.Itertools.html#method.pad_tail_using) for more information. +#[derive(Clone, Debug)] +pub struct PadTailUsing where + I: Iterator, + F: FnMut(usize) -> I::Item, +{ + iter: I, + min: usize, + pos: usize, + filler: F, +} + +impl PadTailUsing where + I: Iterator, + F: FnMut(usize) -> I::Item, +{ + /// Create a new `PadTailUsing` iterator. + pub fn new(iter: I, min: usize, filler: F) -> PadTailUsing { + PadTailUsing { + iter: iter, + min: min, + pos: 0, + filler: filler, + } + } +} + +impl Iterator for PadTailUsing where + I: Iterator, + F: FnMut(usize) -> I::Item, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option { + match self.iter.next() { + None => { + if self.pos < self.min { + let e = Some((self.filler)(self.pos)); + self.pos += 1; + e + } else { + None + } + }, + e => { + self.pos += 1; + e + } + } + } +} + +impl DoubleEndedIterator for PadTailUsing where + I: DoubleEndedIterator + ExactSizeIterator, + F: FnMut(usize) -> I::Item, +{ + fn next_back(&mut self) -> Option { + if self.min == 0 { + self.iter.next_back() + } else if self.iter.len() >= self.min { + self.min -= 1; + self.iter.next_back() + } else { + self.min -= 1; + Some((self.filler)(self.min)) + } + } +} diff --git a/tests/tests.rs b/tests/tests.rs index 7045894ed..467dd3c25 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -96,6 +96,31 @@ fn interleave() { it::assert_equal(it, rs.iter()); } +#[test] +fn interleave_shortest() { + let v0: Vec = vec![0, 2, 4]; + let v1: Vec = vec![1, 3, 5, 7]; + let it = v0.into_iter().interleave_shortest(v1.into_iter()); + assert_eq!(it.size_hint(), (6, Some(6))); + assert_eq!(it.collect_vec(), vec![0, 1, 2, 3, 4, 5]); + + let v0: Vec = vec![0, 2, 4, 6, 8]; + let v1: Vec = vec![1, 3, 5]; + let it = v0.into_iter().interleave_shortest(v1.into_iter()); + assert_eq!(it.size_hint(), (7, Some(7))); + assert_eq!(it.collect_vec(), vec![0, 1, 2, 3, 4, 5, 6]); + + let i0 = ::std::iter::repeat(0); + let v1: Vec<_> = vec![1, 3, 5]; + let it = i0.interleave_shortest(v1.into_iter()); + assert_eq!(it.size_hint(), (7, Some(7))); + + let v0: Vec<_> = vec![0, 2, 4]; + let i1 = ::std::iter::repeat(1); + let it = v0.into_iter().interleave_shortest(i1); + assert_eq!(it.size_hint(), (6, Some(6))); +} + #[test] fn times() { assert!(it::times(0).count() == 0); @@ -556,3 +581,14 @@ fn part() { assert_eq!(i, 3); assert_eq!(data, [9, 6, 3, 4, 5, 2, 7, 8, 1]); } + +#[test] +fn pad_tail_using() { + let v: Vec = vec![0, 1, 2]; + let r: Vec<_> = v.into_iter().pad_tail_using(5, |n| n).collect(); + assert_eq!(r, vec![0, 1, 2, 3, 4]); + + let v: Vec = vec![0, 1, 2]; + let r: Vec<_> = v.into_iter().pad_tail_using(1, |_| panic!()).collect(); + assert_eq!(r, vec![0, 1, 2]); +}