diff --git a/src/libextra/treemap.rs b/src/libextra/treemap.rs index 424492a3cfe..f5a61692509 100644 --- a/src/libextra/treemap.rs +++ b/src/libextra/treemap.rs @@ -14,7 +14,8 @@ use std::util::{swap, replace}; -use std::iterator::{FromIterator, Extendable}; +use std::iterator::{FromIterator, Extendable, Peekable}; +use std::cmp::Ordering; // This is implemented as an AA tree, which is a simplified variation of // a red-black tree where red (horizontal) nodes can only be added @@ -529,24 +530,24 @@ impl TreeSet { /// Visit the values (in-order) representing the difference pub fn difference<'a>(&'a self, other: &'a TreeSet) -> Difference<'a, T> { - Difference{a: Focus::new(self.iter()), b: Focus::new(other.iter())} + Difference{a: self.iter().peekable(), b: other.iter().peekable()} } /// Visit the values (in-order) representing the symmetric difference pub fn symmetric_difference<'a>(&'a self, other: &'a TreeSet) -> SymDifference<'a, T> { - SymDifference{a: Focus::new(self.iter()), b: Focus::new(other.iter())} + SymDifference{a: self.iter().peekable(), b: other.iter().peekable()} } /// Visit the values (in-order) representing the intersection pub fn intersection<'a>(&'a self, other: &'a TreeSet) -> Intersection<'a, T> { - Intersection{a: Focus::new(self.iter()), b: Focus::new(other.iter())} + Intersection{a: self.iter().peekable(), b: other.iter().peekable()} } /// Visit the values (in-order) representing the union pub fn union<'a>(&'a self, other: &'a TreeSet) -> Union<'a, T> { - Union{a: Focus::new(self.iter()), b: Focus::new(other.iter())} + Union{a: self.iter().peekable(), b: other.iter().peekable()} } } @@ -560,61 +561,47 @@ pub struct TreeSetRevIterator<'self, T> { priv iter: TreeMapRevIterator<'self, T, ()> } -// Encapsulate an iterator and hold its latest value until stepped forward -struct Focus { - priv iter: T, - priv focus: Option, -} - -impl> Focus { - fn new(mut it: T) -> Focus { - Focus{focus: it.next(), iter: it} - } - fn step(&mut self) { - self.focus = self.iter.next() - } -} - /// Lazy iterator producing elements in the set difference (in-order) pub struct Difference<'self, T> { - priv a: Focus<&'self T, TreeSetIterator<'self, T>>, - priv b: Focus<&'self T, TreeSetIterator<'self, T>>, + priv a: Peekable<&'self T, TreeSetIterator<'self, T>>, + priv b: Peekable<&'self T, TreeSetIterator<'self, T>>, } /// Lazy iterator producing elements in the set symmetric difference (in-order) pub struct SymDifference<'self, T> { - priv a: Focus<&'self T, TreeSetIterator<'self, T>>, - priv b: Focus<&'self T, TreeSetIterator<'self, T>>, + priv a: Peekable<&'self T, TreeSetIterator<'self, T>>, + priv b: Peekable<&'self T, TreeSetIterator<'self, T>>, } /// Lazy iterator producing elements in the set intersection (in-order) pub struct Intersection<'self, T> { - priv a: Focus<&'self T, TreeSetIterator<'self, T>>, - priv b: Focus<&'self T, TreeSetIterator<'self, T>>, + priv a: Peekable<&'self T, TreeSetIterator<'self, T>>, + priv b: Peekable<&'self T, TreeSetIterator<'self, T>>, } /// Lazy iterator producing elements in the set intersection (in-order) pub struct Union<'self, T> { - priv a: Focus<&'self T, TreeSetIterator<'self, T>>, - priv b: Focus<&'self T, TreeSetIterator<'self, T>>, + priv a: Peekable<&'self T, TreeSetIterator<'self, T>>, + priv b: Peekable<&'self T, TreeSetIterator<'self, T>>, +} + +/// Compare `x` and `y`, but return `short` if x is None and `long` if y is None +fn cmp_opt(x: Option<&T>, y: Option<&T>, + short: Ordering, long: Ordering) -> Ordering { + match (x, y) { + (None , _ ) => short, + (_ , None ) => long, + (Some(x1), Some(y1)) => x1.cmp(y1), + } } impl<'self, T: TotalOrd> Iterator<&'self T> for Difference<'self, T> { fn next(&mut self) -> Option<&'self T> { loop { - match (self.a.focus, self.b.focus) { - (None , _ ) => return None, - (ret , None ) => { self.a.step(); return ret }, - (Some(a1), Some(b1)) => { - let cmp = a1.cmp(b1); - if cmp == Less { - self.a.step(); - return Some(a1); - } else { - if cmp == Equal { self.a.step() } - self.b.step(); - } - } + match cmp_opt(self.a.peek(), self.b.peek(), Less, Less) { + Less => return self.a.next(), + Equal => { self.a.next(); self.b.next(); } + Greater => { self.b.next(); } } } } @@ -623,23 +610,10 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for Difference<'self, T> { impl<'self, T: TotalOrd> Iterator<&'self T> for SymDifference<'self, T> { fn next(&mut self) -> Option<&'self T> { loop { - match (self.a.focus, self.b.focus) { - (ret , None ) => { self.a.step(); return ret }, - (None , ret ) => { self.b.step(); return ret }, - (Some(a1), Some(b1)) => { - let cmp = a1.cmp(b1); - if cmp == Less { - self.a.step(); - return Some(a1); - } else { - self.b.step(); - if cmp == Greater { - return Some(b1); - } else { - self.a.step(); - } - } - } + match cmp_opt(self.a.peek(), self.b.peek(), Greater, Less) { + Less => return self.a.next(), + Equal => { self.a.next(); self.b.next(); } + Greater => return self.b.next(), } } } @@ -648,20 +622,16 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for SymDifference<'self, T> { impl<'self, T: TotalOrd> Iterator<&'self T> for Intersection<'self, T> { fn next(&mut self) -> Option<&'self T> { loop { - match (self.a.focus, self.b.focus) { - (None , _ ) => return None, - (_ , None ) => return None, - (Some(a1), Some(b1)) => { - let cmp = a1.cmp(b1); - if cmp == Less { - self.a.step(); - } else { - self.b.step(); - if cmp == Equal { - return Some(a1); - } - } - }, + let o_cmp = match (self.a.peek(), self.b.peek()) { + (None , _ ) => None, + (_ , None ) => None, + (Some(a1), Some(b1)) => Some(a1.cmp(b1)), + }; + match o_cmp { + None => return None, + Some(Less) => { self.a.next(); } + Some(Equal) => { self.b.next(); return self.a.next() } + Some(Greater) => { self.b.next(); } } } } @@ -670,22 +640,10 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for Intersection<'self, T> { impl<'self, T: TotalOrd> Iterator<&'self T> for Union<'self, T> { fn next(&mut self) -> Option<&'self T> { loop { - match (self.a.focus, self.b.focus) { - (ret , None) => { self.a.step(); return ret }, - (None , ret ) => { self.b.step(); return ret }, - (Some(a1), Some(b1)) => { - let cmp = a1.cmp(b1); - if cmp == Greater { - self.b.step(); - return Some(b1); - } else { - self.a.step(); - if cmp == Equal { - self.b.step(); - } - return Some(a1); - } - } + match cmp_opt(self.a.peek(), self.b.peek(), Greater, Less) { + Less => return self.a.next(), + Equal => { self.b.next(); return self.a.next() } + Greater => return self.b.next(), } } } diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs index a7a1c0bede8..e2d157c2a59 100644 --- a/src/libstd/iterator.rs +++ b/src/libstd/iterator.rs @@ -156,6 +156,28 @@ pub trait Iterator { Enumerate{iter: self, count: 0} } + + /// Creates an iterator that has a `.peek()` method + /// that returns a optional reference to the next element. + /// + /// # Example + /// + /// ~~~ {.rust} + /// let a = [100, 200, 300]; + /// let mut it = xs.iter().map(|&x|x).peekable(); + /// assert_eq!(it.peek().unwrap(), &100); + /// assert_eq!(it.next().unwrap(), 100); + /// assert_eq!(it.next().unwrap(), 200); + /// assert_eq!(it.peek().unwrap(), &300); + /// assert_eq!(it.peek().unwrap(), &300); + /// assert_eq!(it.next().unwrap(), 300); + /// assert!(it.peek().is_none()); + /// assert!(it.next().is_none()); + /// ~~~ + fn peekable(self) -> Peekable { + Peekable{iter: self, peeked: None} + } + /// Creates an iterator which invokes the predicate on elements until it /// returns false. Once the predicate returns false, all further elements are /// yielded. @@ -286,15 +308,15 @@ pub trait Iterator { ///let xs = [1u, 4, 2, 3, 8, 9, 6]; ///let sum = xs.iter() /// .map(|&x| x) - /// .peek(|&x| debug!("filtering %u", x)) + /// .inspect(|&x| debug!("filtering %u", x)) /// .filter(|&x| x % 2 == 0) - /// .peek(|&x| debug!("%u made it through", x)) + /// .inspect(|&x| debug!("%u made it through", x)) /// .sum(); ///println(sum.to_str()); /// ~~~ #[inline] - fn peek<'r>(self, f: &'r fn(&A)) -> Peek<'r, A, Self> { - Peek{iter: self, f: f} + fn inspect<'r>(self, f: &'r fn(&A)) -> Inspect<'r, A, Self> { + Inspect{iter: self, f: f} } /// An adaptation of an external iterator to the for-loop protocol of rust. @@ -1059,6 +1081,38 @@ impl> RandomAccessIterator<(uint, A)> for Enumerat } } +/// An iterator with a `peek()` that returns an optional reference to the next element. +pub struct Peekable { + priv iter: T, + priv peeked: Option, +} + +impl> Iterator for Peekable { + #[inline] + fn next(&mut self) -> Option { + if self.peeked.is_some() { self.peeked.take() } + else { self.iter.next() } + } +} + +impl<'self, A, T: Iterator> Peekable { + /// Return a reference to the next element of the iterator with out advancing it, + /// or None if the iterator is exhausted. + #[inline] + pub fn peek(&'self mut self) -> Option<&'self A> { + match self.peeked { + Some(ref value) => Some(value), + None => { + self.peeked = self.iter.next(); + match self.peeked { + Some(ref value) => Some(value), + None => None, + } + }, + } + } +} + /// An iterator which rejects elements while `predicate` is true pub struct SkipWhile<'self, A, T> { priv iter: T, @@ -1329,14 +1383,14 @@ impl<'self, /// An iterator that calls a function with a reference to each /// element before yielding it. -pub struct Peek<'self, A, T> { +pub struct Inspect<'self, A, T> { priv iter: T, priv f: &'self fn(&A) } -impl<'self, A, T> Peek<'self, A, T> { +impl<'self, A, T> Inspect<'self, A, T> { #[inline] - fn do_peek(&self, elt: Option) -> Option { + fn do_inspect(&self, elt: Option) -> Option { match elt { Some(ref a) => (self.f)(a), None => () @@ -1346,11 +1400,11 @@ impl<'self, A, T> Peek<'self, A, T> { } } -impl<'self, A, T: Iterator> Iterator for Peek<'self, A, T> { +impl<'self, A, T: Iterator> Iterator for Inspect<'self, A, T> { #[inline] fn next(&mut self) -> Option { let next = self.iter.next(); - self.do_peek(next) + self.do_inspect(next) } #[inline] @@ -1359,15 +1413,17 @@ impl<'self, A, T: Iterator> Iterator for Peek<'self, A, T> { } } -impl<'self, A, T: DoubleEndedIterator> DoubleEndedIterator for Peek<'self, A, T> { +impl<'self, A, T: DoubleEndedIterator> DoubleEndedIterator +for Inspect<'self, A, T> { #[inline] fn next_back(&mut self) -> Option { let next = self.iter.next_back(); - self.do_peek(next) + self.do_inspect(next) } } -impl<'self, A, T: RandomAccessIterator> RandomAccessIterator for Peek<'self, A, T> { +impl<'self, A, T: RandomAccessIterator> RandomAccessIterator +for Inspect<'self, A, T> { #[inline] fn indexable(&self) -> uint { self.iter.indexable() @@ -1375,7 +1431,7 @@ impl<'self, A, T: RandomAccessIterator> RandomAccessIterator for Peek<'sel #[inline] fn idx(&self, index: uint) -> Option { - self.do_peek(self.iter.idx(index)) + self.do_inspect(self.iter.idx(index)) } } @@ -1566,6 +1622,24 @@ mod tests { } } + #[test] + fn test_iterator_peekable() { + let xs = ~[0u, 1, 2, 3, 4, 5]; + let mut it = xs.iter().map(|&x|x).peekable(); + assert_eq!(it.peek().unwrap(), &0); + assert_eq!(it.next().unwrap(), 0); + assert_eq!(it.next().unwrap(), 1); + assert_eq!(it.next().unwrap(), 2); + assert_eq!(it.peek().unwrap(), &3); + assert_eq!(it.peek().unwrap(), &3); + assert_eq!(it.next().unwrap(), 3); + assert_eq!(it.next().unwrap(), 4); + assert_eq!(it.peek().unwrap(), &5); + assert_eq!(it.next().unwrap(), 5); + assert!(it.peek().is_none()); + assert!(it.next().is_none()); + } + #[test] fn test_iterator_take_while() { let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19]; @@ -1651,13 +1725,13 @@ mod tests { } #[test] - fn test_peek() { + fn test_inspect() { let xs = [1u, 2, 3, 4]; let mut n = 0; let ys = xs.iter() .map(|&x| x) - .peek(|_| n += 1) + .inspect(|_| n += 1) .collect::<~[uint]>(); assert_eq!(n, xs.len()); @@ -2011,11 +2085,11 @@ mod tests { } #[test] - fn test_random_access_peek() { + fn test_random_access_inspect() { let xs = [1, 2, 3, 4, 5]; - // test .map and .peek that don't implement Clone - let it = xs.iter().peek(|_| {}); + // test .map and .inspect that don't implement Clone + let it = xs.iter().inspect(|_| {}); assert_eq!(xs.len(), it.indexable()); for (i, elt) in xs.iter().enumerate() { assert_eq!(Some(elt), it.idx(i)); @@ -2027,7 +2101,6 @@ mod tests { fn test_random_access_map() { let xs = [1, 2, 3, 4, 5]; - // test .map and .peek that don't implement Clone let it = xs.iter().map(|x| *x); assert_eq!(xs.len(), it.indexable()); for (i, elt) in xs.iter().enumerate() {