glassy 0.0.3

//! Concatenating arrays in streaming operation.
//!
//! Most of the formats supported by this crate require some degree of chunking
//! (input bytes are collected in fixed-size groups before being used).  This
//! creates some problems for streaming implementations -- if a partial group is
//! provided in one run, it needs to be carried over to the next.  This module
//! offers some functionality for dealing with this idiomatically.

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

use super::Group;

/// A buffer of carry-over data.
///
/// This buffer can be used to carry over a partial chunk of data from one
/// iteration to the next.  This is useful when exposing streaming interfaces
/// for formats like Base32 or Base64, where inputs are processed in chunks.
pub struct Carry<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> Carry<T, N> {
    const N_IS_NONZERO: () = assert!(N > 0, "'Carry<T, 0>' is invalid!");

    /// Construct a new [`Carry`].
    pub fn new() -> Self {
        let _ = Self::N_IS_NONZERO;

        Self {
            buf: array::from_fn(|_| MaybeUninit::uninit()),
            len: 0,
        }
    }

    /// Add an element.
    pub fn push(&mut self, elem: T) -> Option<[T; N]> {
        // SAFETY: 'buf' has 'N > len' elements.
        unsafe { self.buf.get_unchecked_mut(self.len) }.write(elem);
        self.len += 1;

        // Restore the 'len < N' property.
        (self.len == N).then(|| {
            let buf = mem::take(self).into_raw_parts().0;
            // SAFETY: 'buf[0 .. len == N]' is initialized.
            buf.map(|e| unsafe { MaybeUninit::assume_init(e) })
        })
    }

    /// Add elements from another [`Carry`].
    pub fn append(&mut self, elems: Self) -> Option<[T; N]> {
        let mut iter = elems.into_iter();
        // Fill 'self' with as many elements as possible.
        let result = self.extend(&mut iter);
        // Having been emptied, another overflow cannot happen.
        let excess = self.extend(&mut iter);
        debug_assert!(excess.is_none());
        debug_assert!(iter.next().is_none());
        result
    }

    /// Add elements from a [`Group`].
    pub fn append_group(&mut self, elems: Group<T, N>) -> Option<[T; N]> {
        let mut iter = elems.into_iter();
        // Fill 'self' with as many elements as possible.
        let result = self.extend(&mut iter);
        // Having been emptied, another overflow cannot happen.
        let excess = self.extend(&mut iter);
        debug_assert!(excess.is_none());
        debug_assert!(iter.next().is_none());
        result
    }

    /// Add elements from an array.
    pub fn append_array(&mut self, elems: [T; N]) -> Option<[T; N]> {
        let mut iter = elems.into_iter();
        // Fill 'self' with as many elements as possible.
        let result = self.extend(&mut iter);
        // Having been emptied, another overflow cannot happen.
        let excess = self.extend(&mut iter);
        debug_assert!(excess.is_none());
        debug_assert!(iter.next().is_none());
        result
    }

    /// Add an elements from an iterator.
    ///
    /// When the [`Carry`] becomes full, the array is output, and no more
    /// elements are extracted.  The remaining elements are still available.
    pub fn extend<I>(&mut self, elems: &mut I) -> Option<[T; N]>
    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).
            d.write(elems.next()?);
            self.len += 1;
        }

        // Restore the 'len < N' property.
        let buf = mem::take(self).into_raw_parts().0;
        // SAFETY: 'buf[0 .. len == N]' is initialized.
        Some(buf.map(|e| unsafe { MaybeUninit::assume_init(e) }))
    }

    /// Copy an elements from an iterator.
    ///
    /// This operates like [`Carry::extend()`], except that the given iterator
    /// provides references to elements which are then copied.
    ///
    /// When the [`Carry`] becomes full, the array is output, and no more
    /// elements are extracted.  The remaining elements are still available.
    pub fn extend_copying<'a, I>(&mut self, elems: &mut I) -> Option<[T; N]>
    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).
            d.write(*elems.next()?);
            self.len += 1;
        }

        // Restore the 'len < N' property.
        let buf = mem::replace(self, Self::new()).buf;
        // SAFETY: 'buf[0 .. len == N]' is initialized.
        Some(buf.map(|e| unsafe { MaybeUninit::assume_init(e) }))
    }

    /// 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 Carry<T, N> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T, const N: usize> Drop for Carry<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 Carry<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 Carry<T, N> {}

impl<T, const N: usize> IntoIterator for Carry<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 Carry<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 Carry<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 Carry<T, N> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.deref(), f)
    }
}

/// The result of [`Carry::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 [`Carry::drain()`].
pub struct Drain<'a, T, const N: usize> {
    /// The containing [`Carry`]'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> {}