diff --git a/src/libcore/iter.rs b/src/libcore/iter.rs index 855bccb07a741..4a6d95e00d6cd 100644 --- a/src/libcore/iter.rs +++ b/src/libcore/iter.rs @@ -68,7 +68,7 @@ use clone::Clone; use cmp; use cmp::{PartialEq, PartialOrd, Ord}; use mem; -use num::{Zero, One, CheckedAdd, CheckedSub, Saturating, ToPrimitive, Int}; +use num::{Zero, One, CheckedAdd, CheckedSub, CheckedMul, Saturating, ToPrimitive, Int}; use ops::{Add, Mul, Sub}; use option::{Option, Some, None}; use uint; @@ -144,6 +144,25 @@ pub trait Iterator { Zip{a: self, b: other} } + /// Creates an iterator which interleaves the contents of this and the specified + /// iterator. + /// + /// # Example + /// + /// ```rust + /// let a = [0i, 2i]; + /// let b = [1i]; + /// let mut it = a.iter().interleave(b.iter()); + /// assert_eq!(it.next().unwrap(), &0i); + /// assert_eq!(it.next().unwrap(), &1i); + /// assert_eq!(it.next().unwrap(), &1i); + /// assert!(it.next().is_none()); + /// ``` + #[inline] + fn interleave>(self, other: U) -> Interleave { + Interleave{a: self, b: other, using_first: true} + } + /// Creates a new iterator which will apply the specified function to each /// element returned by the first, yielding the mapped element instead. /// @@ -1241,6 +1260,112 @@ RandomAccessIterator<(A, B)> for Zip { } } +/// An iterator which interleaves the contents of two other iterators +#[deriving(Clone)] +pub struct Interleave { + a: T, + b: U, + using_first: bool +} + +fn interleave_len(a_len: uint, b_len: uint, using_first: bool) -> Option { + let off_by_one = using_first as uint; + if a_len < b_len { + a_len.checked_mul(&2).and_then(|x| x.checked_add(&(1 - off_by_one))) + } else if b_len < a_len { + b_len.checked_mul(&2).and_then(|x| x.checked_add(&off_by_one)) + } else { + a_len.checked_mul(&2) + } +} + +// Note to developers: sequence will always look like a,b,a,b,a,... but the last element +// to be yielded will be *either* an a or b. Thus we may yield one more a than b. +// This is a key fact to the design of Interleave. +impl, U: Iterator> Iterator for Interleave { + #[inline] + fn next(&mut self) -> Option { + let use_a = self.using_first; + self.using_first = !use_a; + + if use_a { + self.a.next() + } else { + self.b.next() + } + } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + let (a_lower, a_upper) = self.a.size_hint(); + let (b_lower, b_upper) = self.b.size_hint(); + + let lower = interleave_len(a_lower, b_lower, self.using_first).unwrap_or(uint::MAX); + + let upper = match (a_upper, b_upper) { + (Some(x), Some(y)) => interleave_len(x, y, self.using_first), + // If one is missing, best we can do is assume it's huge + (Some(x), None) => interleave_len(x, uint::MAX, self.using_first), + (None, Some(y)) => interleave_len(uint::MAX, y, self.using_first), + (None, None) => None + }; + + (lower, upper) + } +} + +impl, U: ExactSize> DoubleEndedIterator +for Interleave { + #[inline] + fn next_back(&mut self) -> Option { + let (a_sz, a_upper) = self.a.size_hint(); + let (b_sz, b_upper) = self.b.size_hint(); + assert!(a_upper == Some(a_sz)); + assert!(b_upper == Some(b_sz)); + + let off_by_one = self.using_first as uint; + + if a_sz < b_sz { + for _ in range(1 - off_by_one, b_sz - a_sz) { self.b.next_back(); } + } else if a_sz > b_sz { + for _ in range(off_by_one, a_sz - b_sz) { self.a.next_back(); } + } + let (a_sz, _) = self.a.size_hint(); + let (b_sz, _) = self.b.size_hint(); + assert!(a_sz == b_sz || a_sz + (1 - off_by_one) == b_sz + off_by_one); + + if (self.using_first && a_sz > b_sz) || (!self.using_first && a_sz == b_sz) { + self.a.next_back() + } else { + self.b.next_back() + } + } +} + +impl, U: RandomAccessIterator> +RandomAccessIterator for Interleave { + #[inline] + fn indexable(&self) -> uint { + interleave_len(self.a.indexable(), self.b.indexable(), self.using_first) + .unwrap_or(uint::MAX) + } + + #[inline] + fn idx(&mut self, index: uint) -> Option { + let off_by_one = self.using_first as uint; + + if index < self.indexable() { + if index % 2 == (1 - off_by_one) { + self.a.idx((index - (1 - off_by_one)) / 2) + } else { + self.b.idx((index - off_by_one) / 2) + } + } else { + None + } + } +} + /// An iterator which maps the values of `iter` with `f` #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] pub struct Map<'a, A, B, T> { diff --git a/src/libcoretest/iter.rs b/src/libcoretest/iter.rs index 64c336093996d..e9780398b2b50 100644 --- a/src/libcoretest/iter.rs +++ b/src/libcoretest/iter.rs @@ -336,6 +336,7 @@ fn test_iterator_size_hint() { assert_eq!(c.enumerate().size_hint(), (uint::MAX, None)); assert_eq!(c.chain(vi.map(|&i| i)).size_hint(), (uint::MAX, None)); assert_eq!(c.zip(vi).size_hint(), (10, Some(10))); + assert_eq!(c.iterleave(vi).size_hint(), (20, Some(21))) assert_eq!(c.scan(0i, |_,_| Some(0i)).size_hint(), (0, None)); assert_eq!(c.filter(|_| false).size_hint(), (0, None)); assert_eq!(c.map(|_| 0i).size_hint(), (uint::MAX, None)); @@ -350,6 +351,7 @@ fn test_iterator_size_hint() { assert_eq!(vi.enumerate().size_hint(), (10, Some(10))); assert_eq!(vi.chain(v2.iter()).size_hint(), (13, Some(13))); assert_eq!(vi.zip(v2.iter()).size_hint(), (3, Some(3))); + assert_eq!(vi.interleave(v2.iter()).size_hint(), (7, Some(7))); assert_eq!(vi.scan(0i, |_,_| Some(0i)).size_hint(), (0, Some(10))); assert_eq!(vi.filter(|_| false).size_hint(), (0, Some(10))); assert_eq!(vi.map(|i| i+1).size_hint(), (10, Some(10))); @@ -476,6 +478,23 @@ fn test_double_ended_zip() { assert_eq!(it.next(), None); } +#[test] +fn test_double_ended_interleave() { + let xs = [1i, 3, 5, 7, 8]; + let ys = [2i, 4, 6]; + let a = xs.iter().map(|&x| x); + let b = ys.iter().map(|&x| x); + let mut it = a.interleave(b); + assert_eq!(it.next(), Some(1)); + assert_eq!(it.next(), Some(2)); + assert_eq!(it.next_back(), Some(7)); + assert_eq!(it.next_back(), Some(6)); + assert_eq!(it.next(), Some(3)); + assert_eq!(it.next_back(), Some(5)); + assert_eq!(it.next_back(), Some(4)); + assert_eq!(it.next(), None); +} + #[test] fn test_double_ended_filter() { let xs = [1i, 2, 3, 4, 5, 6]; @@ -621,6 +640,27 @@ fn test_random_access_zip() { check_randacc_iter(xs.iter().zip(ys.iter()), cmp::min(xs.len(), ys.len())); } +#[test] +fn test_random_access_interleave() { + let xs = [1i, 3, 5, 7, 8]; // 8 should not be accessible + let ys = [2i, 4, 6]; + let mut it = xs.iter().interleave(ys.iter()); + assert_eq!(it.idx(0).unwrap(), &1); + assert_eq!(it.idx(3).unwrap(), &4); + assert_eq!(it.idx(6).unwrap(), &7); + assert!(it.idx(7).is_none()); + + it.next(); + it.next(); + it.next_back(); + + assert_eq!(it.idx(0).unwrap(), &3); + assert_eq!(it.idx(3).unwrap(), &7); + assert!(it.idx(4).is_none()); + + check_randacc_iter(it, xs.len() + ys.len() - 3); +} + #[test] fn test_random_access_take() { let xs = [1i, 2, 3, 4, 5];