Rust Buffered Iterator

buffered.rs

use std::mem::MaybeUninit;

pub struct Buffered<I: Iterator, const N: usize> {
    iter: I,
    count: usize,
    next: usize,
    buf: [MaybeUninit<I::Item>; N],
}

impl<I: Iterator, const N: usize> Buffered<I, N> {
    pub fn new(iter: I) -> Self {
        Buffered { 
            iter,
            count: 0,
            next: 0,
            buf: [const { MaybeUninit::uninit() }; N],
        }
    }

    pub fn peek(&mut self) -> Option<&I::Item> {
        self.ensure_buffered_and_peek()
            .map(|x| unsafe { x.assume_init_ref() })
    }

    fn ensure_buffered_and_peek(&mut self) -> Option<&MaybeUninit<I::Item>> {
        if self.next < self.count {
            unsafe { Some(self.buf.get_unchecked(self.next)) }
        } else {
            self.count = 0;
            self.next = 0;

            while self.count < N {
                let Some(x) = self.iter.next() else { break };
                self.buf[self.count] = MaybeUninit::new(x);
                self.count += 1;
            }

            if self.next < self.count {
                unsafe { Some(self.buf.get_unchecked(self.next)) }
            } else {
                None
            }
        }
    }

    fn get_buffered_mut(&mut self) -> &mut [MaybeUninit<I::Item>] {
        unsafe { self.buf.get_unchecked_mut(self.next..self.count) }
    }
}

impl<I: Iterator, const N: usize> Iterator for Buffered<I, N> {
    type Item = I::Item;

    fn next(&mut self) -> Option<I::Item> {
        self.ensure_buffered_and_peek()
            .map(|x| unsafe { x.assume_init_read() })
            .inspect(|_| { self.next += 1 })
    }
}

impl<I: Iterator, const N: usize> Drop for Buffered<I, N> {
    fn drop(&mut self) {
        for x in self.get_buffered_mut() {
            unsafe { x.assume_init_drop() }
        }
    }
}

pub trait Bufferable where Self: Iterator + Sized {
    fn buffered<const N: usize>(self) -> Buffered<Self, N> {
        Buffered::<Self, N>::new(self)
    }
}

impl<T> Bufferable for T where T: Iterator {}