diff options
| -rw-r--r-- | src/libcore/iterator.rs | 43 |
1 files changed, 43 insertions, 0 deletions
diff --git a/src/libcore/iterator.rs b/src/libcore/iterator.rs index 50939ba7faf..7ca40ae3a30 100644 --- a/src/libcore/iterator.rs +++ b/src/libcore/iterator.rs @@ -261,6 +261,28 @@ impl<A, T: Iterator<A>> Iterator<A> for TakeIterator<T> { } } +pub struct UnfoldrIterator<'self, A, St> { + priv f: &'self fn(&mut St) -> Option<A>, + priv state: St +} + +pub impl<'self, A, St> UnfoldrIterator<'self, A, St> { + #[inline] + fn new(f: &'self fn(&mut St) -> Option<A>, initial_state: St) -> UnfoldrIterator<'self, A, St> { + UnfoldrIterator { + f: f, + state: initial_state + } + } +} + +impl<'self, A, St> Iterator<A> for UnfoldrIterator<'self, A, St> { + #[inline] + fn next(&mut self) -> Option<A> { + (self.f)(&mut self.state) + } +} + #[cfg(test)] mod tests { use super::*; @@ -326,4 +348,25 @@ mod tests { } assert_eq!(i, ys.len()); } + + #[test] + fn test_unfoldr() { + fn count(st: &mut uint) -> Option<uint> { + if *st < 10 { + let ret = Some(*st); + *st += 1; + ret + } else { + None + } + } + + let mut it = UnfoldrIterator::new(count, 0); + let mut i = 0; + for it.advance |counted| { + assert_eq!(counted, i); + i += 1; + } + assert_eq!(i, 10); + } } |
