glassy 0.0.3

use core::array;
use core::fmt;
use core::mem::{self, MaybeUninit};
use core::iter::{self, FusedIterator};
use core::ops::{Deref, DerefMut};

use super::Carry;

/// An array of at most `N` elements.
pub struct Group<T, const N: usize> {
    /// The internal buffer of carry-over data.
    ///
    /// Soundness Constraints:
    /// - Elements in `0 .. len` are initialized.
    buf: [MaybeUninit<T>; N],

    /// The number of items currently in the buffer.
    ///
    /// Soundness Constraints:
    /// - `len <= N` outside the methods in this module.
    len: usize,
}

impl<T, const N: usize> Group<T, N> {
    /// Construct a new [`Group`].
    pub fn new() -> Self {
        Self {
            buf: array::from_fn(|_| MaybeUninit::uninit()),
            len: 0,
        }
    }

    /// Truncate to the specified number of elements.
    pub fn truncate(&mut self, size: usize) {
        assert!(size <= self.len());
        // SAFETY: 'buf' has 'N >= len' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(size .. self.len) };
        self.len = size;
        for e in buf {
            // SAFETY: 'buf[size .. len]' is initialized.
            unsafe { e.assume_init_drop() };
        }
    }

    /// Fill the [`Group`] until a complete array is formed.
    pub fn fill(mut self, filler: T) -> [T; N]
    where T: Copy {
        self.extend(&mut iter::repeat(filler));
        let (buf, _) = self.into_raw_parts();
        // SAFETY: 'buf[0 .. len == N]' is initialized.
        buf.map(|e| unsafe { MaybeUninit::assume_init(e) })
    }

    /// Fill the [`Group`] until a complete array is formed.
    pub fn fill_with<F>(mut self, mut filler: F) -> [T; N]
    where F: FnMut() -> T {
        self.extend(&mut iter::from_fn(|| Some((filler)())));
        let (buf, _) = self.into_raw_parts();
        // SAFETY: 'buf[0 .. len == N]' is initialized.
        buf.map(|e| unsafe { MaybeUninit::assume_init(e) })
    }

    /// Add an element, if there is space to.
    pub fn push(&mut self, elem: T) -> Result<(), T> {
        if self.len == N { return Err(elem); }

        // SAFETY: 'buf' has 'N > len' elements.
        unsafe { self.buf.get_unchecked_mut(self.len) }.write(elem);
        self.len += 1;
        Ok(())
    } 

    /// Add an elements from an iterator.
    ///
    /// When the [`Group`] becomes full, no more elements are extracted; the
    /// remaining elements are still available.
    pub fn extend<I>(&mut self, elems: &mut I)
    where I: Iterator<Item = T> {
        // SAFETY: 'buf' has 'N >= len' elements.
        for d in unsafe { self.buf.get_unchecked_mut(self.len ..) } {
            // Fill the buffer as much as possible; if the source iterator
            // empties, return immediately (since the buffer is still not full).
            let Some(elem) = elems.next() else { return };
            d.write(elem);
            self.len += 1;
        }
    }

    /// Copy elements from an iterator.
    ///
    /// This operates like [`Group::extend()`], except that the given iterator
    /// provides references to elements which are then copied.
    ///
    /// When the [`Group`] becomes full, no more elements are extracted.  The
    /// remaining elements are still available.
    pub fn extend_copying<'a, I>(&mut self, elems: &mut I)
    where T: 'a + Copy, I: Iterator<Item = &'a T> {
        // SAFETY: 'buf' has 'N > len' elements.
        for d in unsafe { self.buf.get_unchecked_mut(self.len ..) } {
            // Fill the buffer as much as possible; if the source iterator
            // empties, return immediately (since the buffer is still not full).
            let Some(elem) = elems.next() else { return };
            d.write(*elem);
            self.len += 1;
        }
    }

    /// Drain the current buffer, returning any elements in it.
    pub fn drain(&mut self) -> Drain<'_, T, N> {
        Drain {
            beg: 0,
            end: mem::replace(&mut self.len, 0),
            buf: &mut self.buf,
        }
    }

    /// Destructure this object. 
    pub fn into_raw_parts(mut self) -> ([MaybeUninit<T>; N], usize) {
        let buf = array::from_fn(|_| MaybeUninit::uninit());
        let buf = mem::replace(&mut self.buf, buf);
        let len = mem::replace(&mut self.len, 0);
        (buf, len)
    }
}

impl<T, const N: usize> Default for Group<T, N> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T, const N: usize> Drop for Group<T, N> {
    fn drop(&mut self) {
        // SAFETY: 'buf' has 'N >= len' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(0 .. self.len) };
        // SAFETY: 'buf[0 .. len]' is initialized.
        buf.iter_mut().for_each(|v| unsafe { v.assume_init_drop() });
        self.len = 0;
    }
}

impl<T: Clone, const N: usize> Clone for Group<T, N> {
    fn clone(&self) -> Self {
        let buf = array::from_fn(|i| {
            self.buf.get(i)
                // SAFETY: 'buf[0 .. len]' is initialized.
                .map(|e| unsafe { e.assume_init_ref() }.clone())
                .map_or(MaybeUninit::uninit(), MaybeUninit::new)
        });
        Self { buf, len: self.len }
    }
}

// impl<T: Copy, const N: usize> Copy for Group<T, N> {}

impl<T, const N: usize> IntoIterator for Group<T, N> {
    type Item = T;
    type IntoIter = IntoIter<T, N>;

    fn into_iter(self) -> Self::IntoIter {
        let (buf, len) = self.into_raw_parts();
        IntoIter { beg: 0, end: len, buf }
    }
}

impl<T, const N: usize> Deref for Group<T, N> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        // SAFETY: 'buf' has 'N >= len' elements.
        let buf = unsafe { self.buf.get_unchecked(0 .. self.len) };
        // SAFETY: 'buf[0 .. len]' is always initialized.
        unsafe { mem::transmute(buf) }
    }
}

impl<T, const N: usize> DerefMut for Group<T, N> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY: 'buf' has 'N >= len' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(0 .. self.len) };
        // SAFETY: 'buf[0 .. len]' is always initialized.
        unsafe { mem::transmute(buf) }
    }
}

impl<T: fmt::Debug, const N: usize> fmt::Debug for Group<T, N> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.deref(), f)
    }
}

impl<T, const N: usize> From<[T; N]> for Group<T, N> {
    fn from(array: [T; N]) -> Self {
        let buf = array.map(MaybeUninit::new);
        Self { buf, len: N }
    }
}

impl<T, const N: usize> From<Carry<T, N>> for Group<T, N> {
    fn from(carry: Carry<T, N>) -> Self {
        let (buf, len) = carry.into_raw_parts();
        Self { buf, len }
    }
}

/// The result of [`Group::into_iter()`].
pub struct IntoIter<T, const N: usize> {
    /// The start of the valid range.
    beg: usize,
    /// The end of the valid range.
    end: usize,
    /// The buffer of elements.
    buf: [MaybeUninit<T>; N],
}

impl<T, const N: usize> Drop for IntoIter<T, N> {
    fn drop(&mut self) {
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(self.beg .. self.end) };
        // SAFETY: 'buf[beg .. end]' is initialized.
        buf.iter_mut().for_each(|v| unsafe { v.assume_init_drop() });
    }
}

impl<T, const N: usize> Iterator for IntoIter<T, N> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.beg >= self.end { return None; }
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let elem = unsafe { self.buf.get_unchecked_mut(self.beg) };
        let elem = mem::replace(elem, MaybeUninit::uninit());
        self.beg += 1;
        // SAFETY: 'buf[beg .. end]' is always initialized.
        Some(unsafe { elem.assume_init() })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.end - self.beg, Some(self.end - self.beg))
    }

    fn count(self) -> usize {
        // Note: remaining elements are dropped automatically.
        self.end - self.beg
    }

    fn last(mut self) -> Option<Self::Item> {
        // Note: remaining elements are dropped automatically.
        self.next_back()
    }

    fn for_each<F: FnMut(Self::Item)>(mut self, mut f: F) {
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(self.beg .. self.end) };
        self.beg = self.end;
        buf.iter_mut().for_each(|elem| {
            let elem = mem::replace(elem, MaybeUninit::uninit());
            // SAFETY: 'buf[beg .. end]' is always initialized.
            (f)(unsafe { elem.assume_init() })
        })
    }

    fn fold<B, F: FnMut(B, Self::Item) -> B>(mut self, init: B, mut f: F) -> B {
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(self.beg .. self.end) };
        self.beg = self.end;
        buf.iter_mut().fold(init, |state, elem| {
            let elem = mem::replace(elem, MaybeUninit::uninit());
            // SAFETY: 'buf[beg .. end]' is always initialized.
            (f)(state, unsafe { elem.assume_init() })
        })
    }
}

impl<'a, T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.end >= self.beg { return None; }
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        self.end -= 1;
        let elem = unsafe { self.buf.get_unchecked_mut(self.end) };
        let elem = mem::replace(elem, MaybeUninit::uninit());
        // SAFETY: 'buf[beg .. end]' is always initialized.
        Some(unsafe { elem.assume_init() })
    }

    fn rfold<B, F: FnMut(B, Self::Item) -> B>(mut self, init: B, mut f: F) -> B {
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(self.beg .. self.end) };
        self.beg = self.end;
        buf.iter_mut().rfold(init, |state, elem| {
            let elem = mem::replace(elem, MaybeUninit::uninit());
            // SAFETY: 'buf[beg .. end]' is always initialized.
            (f)(state, unsafe { elem.assume_init() })
        })
    }
}

impl<'a, T, const N: usize> ExactSizeIterator for IntoIter<T, N> {
    fn len(&self) -> usize {
        self.end - self.beg
    }
}

impl<'a, T, const N: usize> FusedIterator for IntoIter<T, N> {}

/// The result of [`Group::drain()`].
pub struct Drain<'a, T, const N: usize> {
    /// The containing [`Group`]'s buffer.
    buf: &'a mut [MaybeUninit<T>; N],
    /// The start of the valid range.
    beg: usize,
    /// The end of the valid range.
    end: usize,
}

impl<'a, T, const N: usize> Drop for Drain<'a, T, N> {
    fn drop(&mut self) {
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(self.beg .. self.end) };
        // SAFETY: 'buf[beg .. end]' is initialized.
        buf.iter_mut().for_each(|v| unsafe { v.assume_init_drop() });
    }
}

impl<'a, T, const N: usize> Iterator for Drain<'a, T, N> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.beg >= self.end { return None; }
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let elem = unsafe { self.buf.get_unchecked_mut(self.beg) };
        let elem = mem::replace(elem, MaybeUninit::uninit());
        self.beg += 1;
        // SAFETY: 'buf[beg .. end]' is always initialized.
        Some(unsafe { elem.assume_init() })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.end - self.beg, Some(self.end - self.beg))
    }

    fn count(self) -> usize {
        // Note: remaining elements are dropped automatically.
        self.end - self.beg
    }

    fn last(mut self) -> Option<Self::Item> {
        // Note: remaining elements are dropped automatically.
        self.next_back()
    }

    fn for_each<F: FnMut(Self::Item)>(mut self, mut f: F) {
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(self.beg .. self.end) };
        self.beg = self.end;
        buf.iter_mut().for_each(|elem| {
            let elem = mem::replace(elem, MaybeUninit::uninit());
            // SAFETY: 'buf[beg .. end]' is always initialized.
            (f)(unsafe { elem.assume_init() })
        })
    }

    fn fold<B, F: FnMut(B, Self::Item) -> B>(mut self, init: B, mut f: F) -> B {
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(self.beg .. self.end) };
        self.beg = self.end;
        buf.iter_mut().fold(init, |state, elem| {
            let elem = mem::replace(elem, MaybeUninit::uninit());
            // SAFETY: 'buf[beg .. end]' is always initialized.
            (f)(state, unsafe { elem.assume_init() })
        })
    }
}

impl<'a, T, const N: usize> DoubleEndedIterator for Drain<'a, T, N> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.end >= self.beg { return None; }
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        self.end -= 1;
        let elem = unsafe { self.buf.get_unchecked_mut(self.end) };
        let elem = mem::replace(elem, MaybeUninit::uninit());
        // SAFETY: 'buf[beg .. end]' is always initialized.
        Some(unsafe { elem.assume_init() })
    }

    fn rfold<B, F: FnMut(B, Self::Item) -> B>(mut self, init: B, mut f: F) -> B {
        // SAFETY: 'buf' has 'N >= end > beg' elements.
        let buf = unsafe { self.buf.get_unchecked_mut(self.beg .. self.end) };
        self.beg = self.end;
        buf.iter_mut().rfold(init, |state, elem| {
            let elem = mem::replace(elem, MaybeUninit::uninit());
            // SAFETY: 'buf[beg .. end]' is always initialized.
            (f)(state, unsafe { elem.assume_init() })
        })
    }
}

impl<'a, T, const N: usize> ExactSizeIterator for Drain<'a, T, N> {
    fn len(&self) -> usize {
        self.end - self.beg
    }
}

impl<'a, T, const N: usize> FusedIterator for Drain<'a, T, N> {}