repeat_n.rs - source
Description: Source of the Rust file `library/core/src/iter/sources/repeat_n.rs`.

#
core/iter/sources/
repeat_n.rs

    1use crate::fmt;
    2use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
    3use crate::num::NonZero;
    4use crate::ops::Try;
    5
    6/// Creates a new iterator that repeats a single element a given number of times.
    7///
    8/// The `repeat_n()` function repeats a single value exactly `n` times.
    9///
    10/// This is very similar to using [`repeat()`] with [`Iterator::take()`],
    11/// but `repeat_n()` can return the original value, rather than always cloning.
    12///
    13/// [`repeat()`]: crate::iter::repeat
    14///
    15/// # Examples
    16///
    17/// Basic usage:
    18///
    19/// ```
    20/// use std::iter;
    21///
    22/// // four of the number four:
    23/// let mut four_fours = iter::repeat_n(4, 4);
    24///
    25/// assert_eq!(Some(4), four_fours.next());
    26/// assert_eq!(Some(4), four_fours.next());
    27/// assert_eq!(Some(4), four_fours.next());
    28/// assert_eq!(Some(4), four_fours.next());
    29///
    30/// // no more fours
    31/// assert_eq!(None, four_fours.next());
    32/// ```
    33///
    34/// For non-`Copy` types,
    35///
    36/// ```
    37/// use std::iter;
    38///
    39/// let v: Vec<i32> = Vec::with_capacity(123);
    40/// let mut it = iter::repeat_n(v, 5);
    41///
    42/// for i in 0..4 {
    43///     // It starts by cloning things
    44///     let cloned = it.next().unwrap();
    45///     assert_eq!(cloned.len(), 0);
    46///     assert_eq!(cloned.capacity(), 0);
    47/// }
    48///
    49/// // ... but the last item is the original one
    50/// let last = it.next().unwrap();
    51/// assert_eq!(last.len(), 0);
    52/// assert_eq!(last.capacity(), 123);
    53///
    54/// // ... and now we're done
    55/// assert_eq!(None, it.next());
    56/// ```
    57#[inline]
    58#[stable(feature = "iter_repeat_n", since = "1.82.0")]
    59pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
    60    RepeatN { inner: RepeatNInner::new(element, count) }
    61}
    62
    63#[derive(Clone, Copy)]
    64struct RepeatNInner<T> {
    65    count: NonZero<usize>,
    66    element: T,
    67}
    68
    69impl<T> RepeatNInner<T> {
    70    fn new(element: T, count: usize) -> Option<Self> {
    71        let count = NonZero::<usize>::new(count)?;
    72        Some(Self { element, count })
    73    }
    74}
    75
    76/// An iterator that repeats an element an exact number of times.
    77///
    78/// This `struct` is created by the [`repeat_n()`] function.
    79/// See its documentation for more.
    80#[stable(feature = "iter_repeat_n", since = "1.82.0")]
    81#[derive(Clone)]
    82pub struct RepeatN<A> {
    83    inner: Option<RepeatNInner<A>>,
    84}
    85
    86impl<A> RepeatN<A> {
    87    /// If we haven't already dropped the element, return it in an option.
    88    #[inline]
    89    fn take_element(&mut self) -> Option<A> {
    90        self.inner.take().map(|inner| inner.element)
    91    }
    92}
    93
    94#[stable(feature = "iter_repeat_n", since = "1.82.0")]
    95impl<A: fmt::Debug> fmt::Debug for RepeatN<A> {
    96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    97        let (count, element) = match self.inner.as_ref() {
    98            Some(inner) => (inner.count.get(), Some(&inner.element)),
    99            None => (0, None),
    100        };
    101        f.debug_struct("RepeatN").field("count", &count).field("element", &element).finish()
    102    }
    103}
    104
    105#[stable(feature = "iter_repeat_n", since = "1.82.0")]
    106impl<A: Clone> Iterator for RepeatN<A> {
    107    type Item = A;
    108
    109    #[inline]
    110    fn next(&mut self) -> Option<A> {
    111        let inner = self.inner.as_mut()?;
    112        let count = inner.count.get();
    113
    114        if let Some(decremented) = NonZero::<usize>::new(count - 1) {
    115            // Order of these is important for optimization
    116            let tmp = inner.element.clone();
    117            inner.count = decremented;
    118            return Some(tmp);
    119        }
    120
    121        return self.take_element();
    122    }
    123
    124    #[inline]
    125    fn size_hint(&self) -> (usize, Option<usize>) {
    126        let len = self.len();
    127        (len, Some(len))
    128    }
    129
    130    #[inline]
    131    fn advance_by(&mut self, skip: usize) -> Result<(), NonZero<usize>> {
    132        let Some(inner) = self.inner.as_mut() else {
    133            return NonZero::<usize>::new(skip).map(Err).unwrap_or(Ok(()));
    134        };
    135
    136        let len = inner.count.get();
    137
    138        if let Some(new_len) = len.checked_sub(skip).and_then(NonZero::<usize>::new) {
    139            inner.count = new_len;
    140            return Ok(());
    141        }
    142
    143        self.inner = None;
    144        return NonZero::<usize>::new(skip - len).map(Err).unwrap_or(Ok(()));
    145    }
    146
    147    #[inline]
    148    fn last(mut self) -> Option<A> {
    149        self.take_element()
    150    }
    151
    152    #[inline]
    153    fn count(self) -> usize {
    154        self.len()
    155    }
    156}
    157
    158#[stable(feature = "iter_repeat_n", since = "1.82.0")]
    159impl<A: Clone> ExactSizeIterator for RepeatN<A> {
    160    fn len(&self) -> usize {
    161        self.inner.as_ref().map(|inner| inner.count.get()).unwrap_or(0)
    162    }
    163}
    164
    165#[stable(feature = "iter_repeat_n", since = "1.82.0")]
    166impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
    167    #[inline]
    168    fn next_back(&mut self) -> Option<A> {
    169        self.next()
    170    }
    171
    172    #[inline]
    173    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
    174        self.advance_by(n)
    175    }
    176
    177    #[inline]
    178    fn nth_back(&mut self, n: usize) -> Option<A> {
    179        self.nth(n)
    180    }
    181
    182    #[inline]
    183    fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
    184    where
    185        F: FnMut(B, A) -> R,
    186        R: Try<Output = B>,
    187    {
    188        self.try_fold(init, f)
    189    }
    190
    191    #[inline]
    192    fn rfold<B, F>(self, init: B, f: F) -> B
    193    where
    194        F: FnMut(B, A) -> B,
    195    {
    196        self.fold(init, f)
    197    }
    198}
    199
    200#[stable(feature = "iter_repeat_n", since = "1.82.0")]
    201impl<A: Clone> FusedIterator for RepeatN<A> {}
    202
    203#[unstable(feature = "trusted_len", issue = "37572")]
    204unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}
    205#[stable(feature = "iter_repeat_n", since = "1.82.0")]
    206impl<A: Clone> UncheckedIterator for RepeatN<A> {}