add core::iter::repeat_with

2 files changed, 148 insertions, 0 deletions

diff --git a/src/libcore/iter/sources.rs b/src/libcore/iter/sources.rs
index b05a893e661..980f3fc7443 100644
--- a/src/libcore/iter/sources.rs
+++ b/src/libcore/iter/sources.rs
@@ -57,6 +57,12 @@ unsafe impl<A: Clone> TrustedLen for Repeat<A> {}
 ///
 /// [`take`]: trait.Iterator.html#method.take
 ///
+/// If the element type of the iterator you need does not implement `Clone`,
+/// or if you do not want to keep the repeated element in memory, you can
+/// instead use the [`repeat_with`] function.
+///
+/// [`repeat_with`]: fn.repeat_with.html
+///
 /// # Examples
 ///
 /// Basic usage:
@@ -99,6 +105,104 @@ pub fn repeat<T: Clone>(elt: T) -> Repeat<T> {
     Repeat{element: elt}
 }
 
+/// An iterator that repeats elements of type `A` endlessly by
+/// applying the provided closure `F: FnMut() -> A`.
+///
+/// This `struct` is created by the [`repeat_with`] function.
+/// See its documentation for more.
+///
+/// [`repeat_with`]: fn.repeat_with.html
+#[unstable(feature = "iterator_repeat_with", issue = "0")]
+pub struct RepeatWith<F> {
+    repeater: F
+}
+
+#[unstable(feature = "iterator_repeat_with", issue = "0")]
+impl<A, F: FnMut() -> A> Iterator for RepeatWith<F> {
+    type Item = A;
+
+    #[inline]
+    fn next(&mut self) -> Option<A> { Some((self.repeater)()) }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) { (usize::MAX, None) }
+}
+
+#[unstable(feature = "iterator_repeat_with", issue = "0")]
+impl<A, F: FnMut() -> A> DoubleEndedIterator for RepeatWith<F> {
+    #[inline]
+    fn next_back(&mut self) -> Option<A> { self.next() }
+}
+
+#[unstable(feature = "fused", issue = "35602")]
+impl<A, F: FnMut() -> A> FusedIterator for RepeatWith<F> {}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A, F: FnMut() -> A> TrustedLen for RepeatWith<F> {}
+
+/// Creates a new that repeats elements of type `A` endlessly by
+/// applying the provided closure, the repeater, `F: FnMut() -> A`.
+///
+/// The `repeat_with()` function calls the repeater over and over and over and
+/// over and over and 🔁.
+///
+/// Infinite iterators like `repeat_with()` are often used with adapters like
+/// [`take`], in order to make them finite.
+///
+/// [`take`]: trait.Iterator.html#method.take
+///
+/// If the element type of the iterator you need implements `Clone`, and
+/// it is OK to keep the source element in memory, you should instead use
+/// the [`repeat`] function.
+///
+/// [`repeat`]: fn.repeat.html
+///
+/// # Examples
+///
+/// Basic usage:
+///
+/// ```
+/// use std::iter;
+///
+/// // let's assume we have some value of a type that is not `Clone`
+/// // or which don't want to have in memory just yet because it is expensive:
+/// #[derive(PartialEq, Debug)]
+/// struct Expensive;
+///
+/// // a particular value forever:
+/// let mut things = iter::repeat_with(|| Expensive);
+///
+/// assert_eq!(Some(Expensive), things.next());
+/// assert_eq!(Some(Expensive), things.next());
+/// assert_eq!(Some(Expensive), things.next());
+/// assert_eq!(Some(Expensive), things.next());
+/// assert_eq!(Some(Expensive), things.next());
+/// ```
+///
+/// Using mutation and going finite:
+///
+/// ```rust
+/// use std::iter;
+///
+/// // From the zeroth to the third power of two:
+/// let mut curr = 1;
+/// let mut pow2 = iter::repeat_with(|| { let tmp = curr; curr *= 2; tmp })
+///                     .take(4);
+///
+/// assert_eq!(Some(1), pow2.next());
+/// assert_eq!(Some(2), pow2.next());
+/// assert_eq!(Some(4), pow2.next());
+/// assert_eq!(Some(8), pow2.next());
+///
+/// // ... and now we're done
+/// assert_eq!(None, pow2.next());
+/// ```
+#[inline]
+#[unstable(feature = "iterator_repeat_with", issue = "0")]
+pub fn repeat_with<A, F: FnMut() -> A>(repeater: F) -> RepeatWith<F> {
+    RepeatWith { repeater }
+}
+
 /// An iterator that yields nothing.
 ///
 /// This `struct` is created by the [`empty`] function. See its documentation for more.
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index b2a5243d5e6..ca5318d198e 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -1550,6 +1550,50 @@ fn test_repeat_take_collect() {
 }
 
 #[test]
+fn test_repeat_with() {
+    struct NotClone(usize);
+    let mut it = repeat_with(|| NotClone(42));
+    assert_eq!(it.next(), Some(NotClone(42)));
+    assert_eq!(it.next(), Some(NotClone(42)));
+    assert_eq!(it.next(), Some(NotClone(42)));
+    assert_eq!(repeat_with(|| NotClone(42)).size_hint(), (usize::MAX, None));
+}
+
+#[test]
+fn test_repeat_with_rev() {
+    let mut curr = 1;
+    let mut pow2 = repeat_with(|| { let tmp = curr; curr *= 2; tmp })
+                    .rev().take(4);
+    assert_eq!(it.next(), Some(1));
+    assert_eq!(it.next(), Some(2));
+    assert_eq!(it.next(), Some(4));
+    assert_eq!(it.next(), Some(8));
+    assert_eq!(it.next(), None);
+}
+
+#[test]
+fn test_repeat_with_take() {
+    let mut it = repeat_with(|| 42).take(3);
+    assert_eq!(it.next(), Some(42));
+    assert_eq!(it.next(), Some(42));
+    assert_eq!(it.next(), Some(42));
+    assert_eq!(it.next(), None);
+    is_trusted_len(repeat_with(|| 42).take(3));
+    assert_eq!(repeat_with(|| 42).take(3).size_hint(), (3, Some(3)));
+    assert_eq!(repeat_with(|| 42).take(0).size_hint(), (0, Some(0)));
+    assert_eq!(repeat_with(|| 42).take(usize::MAX).size_hint(),
+               (usize::MAX, Some(usize::MAX)));
+}
+
+#[test]
+fn test_repeat_take_collect() {
+    let mut curr = 1;
+    let v: Vec<_> = repeat_with(|| { let tmp = curr; curr *= 2; tmp })
+                      .take(5).collect();
+    assert_eq!(v, vec![1, 2, 4, 8, 16]);
+}
+
+#[test]
 fn test_fuse() {
     let mut it = 0..3;
     assert_eq!(it.len(), 3);