Compiler Explorer

Source code

#![feature(trusted_len)]
#![feature(iter_advance_by)]

use std::fmt;
use std::iter::{FusedIterator, TrustedLen};
use std::mem::{self, MaybeUninit};
use std::num::NonZero;

/// Creates a new iterator that repeats a single element a given number of times.
///
/// The `repeat_n()` function repeats a single value exactly `n` times.
///
/// This is very similar to using [`repeat()`] with [`Iterator::take()`],
/// but `repeat_n()` can return the original value, rather than always cloning.
///
/// [`repeat()`]: crate::iter::repeat
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// use std::iter;
///
/// // four of the number four:
/// let mut four_fours = iter::repeat_n(4, 4);
///
/// assert_eq!(Some(4), four_fours.next());
/// assert_eq!(Some(4), four_fours.next());
/// assert_eq!(Some(4), four_fours.next());
/// assert_eq!(Some(4), four_fours.next());
///
/// // no more fours
/// assert_eq!(None, four_fours.next());
/// ```
///
/// For non-`Copy` types,
///
/// ```
/// use std::iter;
///
/// let v: Vec<i32> = Vec::with_capacity(123);
/// let mut it = iter::repeat_n(v, 5);
///
/// for i in 0..4 {
///     // It starts by cloning things
///     let cloned = it.next().unwrap();
///     assert_eq!(cloned.len(), 0);
///     assert_eq!(cloned.capacity(), 0);
/// }
///
/// // ... but the last item is the original one
/// let last = it.next().unwrap();
/// assert_eq!(last.len(), 0);
/// assert_eq!(last.capacity(), 123);
///
/// // ... and now we're done
/// assert_eq!(None, it.next());
/// ```
#[inline]
fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
    let element = if count == 0 {
        // `element` gets dropped eagerly.
        MaybeUninit::uninit()
    } else {
        MaybeUninit::new(element)
    };

RepeatN { element, count }
}

/// An iterator that repeats an element an exact number of times.
///
/// This `struct` is created by the [`repeat_n()`] function.
/// See its documentation for more.
struct RepeatN<A> {
    count: usize,
    // Invariant: uninit iff count == 0.
    element: MaybeUninit<A>,
}

impl<A> RepeatN<A> {
    /// Returns the element if it hasn't been dropped already.
    fn element_ref(&self) -> Option<&A> {
        if self.count > 0 {
            // SAFETY: The count is non-zero, so it must be initialized.
            Some(unsafe { self.element.assume_init_ref() })
        } else {
            None
        }
    }
    /// If we haven't already dropped the element, return it in an option.
    ///
    /// Clears the count so it won't be dropped again later.
    #[inline]
    fn take_element(&mut self) -> Option<A> {
        if self.count > 0 {
            self.count = 0;
            let element = mem::replace(&mut self.element, MaybeUninit::uninit());
            // SAFETY: We just set count to zero so it won't be dropped again,
            // and it used to be non-zero so it hasn't already been dropped.
            unsafe { Some(element.assume_init()) }
        } else {
            None
        }
    }
}

impl<A: Clone> RepeatN<A> {
    #[inline]
    unsafe fn next_unchecked(&mut self) -> A {
        // SAFETY: The caller promised the iterator isn't empty
        self.count = unsafe { self.count.unchecked_sub(1) };
        if self.count == 0 {
            // SAFETY: the check above ensured that the count used to be non-zero,
            // so element hasn't been dropped yet, and we just lowered the count to
            // zero so it won't be dropped later, and thus it's okay to take it here.
            unsafe { mem::replace(&mut self.element, MaybeUninit::uninit()).assume_init() }
        } else {
            // SAFETY: the count is non-zero, so it must have not been dropped yet.
            let element = unsafe { self.element.assume_init_ref() };
            A::clone(element)
        }
    }
}

impl<A> Drop for RepeatN<A> {
    fn drop(&mut self) {
        self.take_element();
    }
}

impl<A: Clone> Iterator for RepeatN<A> {
    type Item = A;

#[inline]
    fn next(&mut self) -> Option<A> {
        if self.count > 0 {
            // SAFETY: Just checked it's not empty
            unsafe { Some(self.next_unchecked()) }
        } else {
            None
        }
    }

#[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }

#[inline]
    fn advance_by(&mut self, skip: usize) -> Result<(), NonZero<usize>> {
        let len = self.count;

if skip >= len {
            self.take_element();
        }

if skip > len {
            // SAFETY: we just checked that the difference is positive
            Err(unsafe { NonZero::new_unchecked(skip - len) })
        } else {
            self.count = len - skip;
            Ok(())
        }
    }

// trying to figure out if it helped
    /*
    #[inline]
    fn fold<B, F>(mut self, init: B, mut f: F) -> B
    where
        Self: Sized,
        F: FnMut(B, Self::Item) -> B,
    {
        if self.count == 0 {
            return init;
        }
        let element = unsafe { self.element.assume_init_ref() };

let init = (0..self.count.saturating_sub(1)).fold(init, |accum, _| {
            self.count = unsafe {
                self.count.unchecked_sub(1)
            };
            f(accum, element.clone())
        });
        f(init, unsafe { self.take_element().unwrap_unchecked() })
    }
    */

#[inline]
    fn last(mut self) -> Option<A> {
        self.take_element()
    }

#[inline]
    fn count(self) -> usize {
        self.len()
    }
}

impl<A: Clone> ExactSizeIterator for RepeatN<A> {
    fn len(&self) -> usize {
        self.count
    }
}

impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
    #[inline]
    fn next_back(&mut self) -> Option<A> {
        self.next()
    }

#[inline]
    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
        self.advance_by(n)
    }

#[inline]
    fn nth_back(&mut self, n: usize) -> Option<A> {
        self.nth(n)
    }
}

impl<A: Clone> FusedIterator for RepeatN<A> {}

unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}

pub fn thing(vec: &mut Vec<Vec<u8>>, n: usize) {
    // Making the assembly smaller
    //if vec.len() != 0 || vec.capacity() != 2048 {
    //    unsafe {
    //        std::hint::unreachable_unchecked();
    //    }
    //}

vec.extend(repeat_n(Vec::new(), n));
}

Source code

#![feature(trusted_len)]
#![feature(iter_advance_by)]

use core::fmt;
use core::iter::{FusedIterator, TrustedLen};
use core::num::NonZero;

/// Creates a new iterator that repeats a single element a given number of times.
///
/// The `repeat_n()` function repeats a single value exactly `n` times.
///
/// This is very similar to using [`repeat()`] with [`Iterator::take()`],
/// but `repeat_n()` can return the original value, rather than always cloning.
///
/// [`repeat()`]: crate::iter::repeat
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// use std::iter;
///
/// // four of the number four:
/// let mut four_fours = iter::repeat_n(4, 4);
///
/// assert_eq!(Some(4), four_fours.next());
/// assert_eq!(Some(4), four_fours.next());
/// assert_eq!(Some(4), four_fours.next());
/// assert_eq!(Some(4), four_fours.next());
///
/// // no more fours
/// assert_eq!(None, four_fours.next());
/// ```
///
/// For non-`Copy` types,
///
/// ```
/// use std::iter;
///
/// let v: Vec<i32> = Vec::with_capacity(123);
/// let mut it = iter::repeat_n(v, 5);
///
/// for i in 0..4 {
///     // It starts by cloning things
///     let cloned = it.next().unwrap();
///     assert_eq!(cloned.len(), 0);
///     assert_eq!(cloned.capacity(), 0);
/// }
///
/// // ... but the last item is the original one
/// let last = it.next().unwrap();
/// assert_eq!(last.len(), 0);
/// assert_eq!(last.capacity(), 123);
///
/// // ... and now we're done
/// assert_eq!(None, it.next());
/// ```
#[inline]
pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
    RepeatN {
        inner: RepeatNInner::new(element, count),
    }
}

#[derive(Clone)]
struct RepeatNInner<T> {
    element: T,
    count: NonZero<usize>,
}

impl<T> RepeatNInner<T> {
    fn new(element: T, count: usize) -> Option<Self> {
        let count = NonZero::<usize>::new(count)?;
        Some(Self { element, count })
    }
}

/// An iterator that repeats an element an exact number of times.
///
/// This `struct` is created by the [`repeat_n()`] function.
/// See its documentation for more.
#[derive(Clone)]
pub struct RepeatN<A> {
    inner: Option<RepeatNInner<A>>,
}

impl<A> RepeatN<A> {
    /// If we haven't already dropped the element, return it in an option.
    #[inline]
    fn take_element(&mut self) -> Option<A> {
        self.inner.take().map(|inner| inner.element)
    }
}

impl<A: fmt::Debug> fmt::Debug for RepeatN<A> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (count, element) = match self.inner.as_ref() {
            Some(inner) => (inner.count.get(), Some(&inner.element)),
            None => (0, None),
        };
        f.debug_struct("RepeatN")
            .field("count", &count)
            .field("element", &element)
            .finish()
    }
}

impl<A: Clone> Iterator for RepeatN<A> {
    type Item = A;

#[inline]
    fn next(&mut self) -> Option<A> {
        let inner = self.inner.as_mut()?;
        let count = inner.count.get();

if let Some(decremented) = NonZero::<usize>::new(count - 1) {
            // Order of these is important for optimization
            let tmp = inner.element.clone();
            inner.count = decremented;
            return Some(tmp);
        }

return self.take_element();
    }

#[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }

#[inline]
    fn advance_by(&mut self, skip: usize) -> Result<(), NonZero<usize>> {
        let Some(inner) = self.inner.as_mut() else {
            return NonZero::<usize>::new(skip).map(Err).unwrap_or(Ok(()));
        };

let len = inner.count.get();

if let Some(new_len) = len.checked_sub(skip).and_then(NonZero::<usize>::new) {
            inner.count = new_len;
            return Ok(());
        }

self.inner = None;
        return NonZero::<usize>::new(skip - len).map(Err).unwrap_or(Ok(()));
    }