#![allow(unused_unsafe)]
#![allow(unsafe_code)]

//! Contains implementations for rust core that have not been stabilized
//!
//! Functions in this are expected to be properly peer reviewed by the community
//!
//! Any modifications done are purely to make the code compatible with bincode

use core::mem::MaybeUninit;
use core::mem::{
    self,
};

struct Guard<'a, T, const N: usize> {
    array_mut: &'a mut [MaybeUninit<T>; N],
    initialized: usize,
}

impl<T, const N: usize> Drop for Guard<'_, T, N> {
    #[inline(always)]
    fn drop(&mut self) {
        debug_assert!(self.initialized <= N);

        // SAFETY: this slice will contain only initialized objects.
        unsafe {
            core::ptr::drop_in_place(slice_assume_init_mut(
                self.array_mut.get_unchecked_mut(..self.initialized),
            ));
        }
    }
}

/// Pulls `N` items from `iter` and returns them as an array. If the iterator
/// yields fewer than `N` items, `None` is returned and all already yielded
/// items are dropped.
///
/// Since the iterator is passed as a mutable reference and this function calls
/// `next` at most `N` times, the iterator can still be used afterwards to
/// retrieve the remaining items.
///
/// If `iter.next()` panicks, all items already yielded by the iterator are
/// dropped.
#[allow(clippy::while_let_on_iterator)]
#[inline]
pub fn collect_into_array<E, I, T, const N: usize>(iter: &mut I) -> Option<Result<[T; N], E>>
where
    I: Iterator<Item = Result<T, E>>,
{
    if N == 0 {
        // SAFETY: An empty array is always inhabited and has no validity invariants.
        return unsafe { Some(Ok(mem::zeroed())) };
    }

    let mut array = uninit_array::<T, N>();
    let mut guard = Guard {
        array_mut: &mut array,
        initialized: 0,
    };

    loop {
        let next_item = iter.next();
        if let Some(item_rslt) = next_item {
            let item = match item_rslt {
                | Err(err) => {
                    return Some(Err(err));
                },
                | Ok(elem) => elem,
            };

            // SAFETY: `guard.initialized` starts at 0, is increased by one in the
            // loop and the loop is aborted once it reaches N (which is
            // `array.len()`).
            unsafe {
                guard
                    .array_mut
                    .get_unchecked_mut(guard.initialized)
                    .write(item);
            }
            guard.initialized += 1;

            // Check if the whole array was initialized.
            if guard.initialized == N {
                mem::forget(guard);

                // SAFETY: the condition above asserts that all elements are
                // initialized.
                let out = unsafe { array_assume_init(&array) };
                return Some(Ok(out));
            }
        } else {
            break;
        }
    }

    // This is only reached if the iterator is exhausted before
    // `guard.initialized` reaches `N`. Also note that `guard` is dropped here,
    // dropping all already initialized elements.
    None
}

/// Assuming all the elements are initialized, get a mutable slice to them.
///
/// # Safety
///
/// It is up to the caller to guarantee that the `MaybeUninit<T>` elements
/// really are in an initialized state.
/// Calling this when the content is not yet fully initialized causes undefined behavior.
///
/// See [`assume_init_mut`\] for more details and examples.
///
/// [`assume_init_mut`\]: `MaybeUninit::assume_init_mut`
// #[unstable(feature = "maybe_uninit_slice", issue = "63569")]
// #[rustc_const_unstable(feature = "const_maybe_uninit_assume_init", issue = "none")]
#[inline(always)]
pub unsafe fn slice_assume_init_mut<T>(slice: &mut [MaybeUninit<T>]) -> &mut [T] {
    // SAFETY: similar to safety notes for `slice_get_ref`, but we have a
    // mutable reference which is also guaranteed to be valid for writes.
    unsafe { &mut *(core::ptr::from_mut::<[MaybeUninit<T>]>(slice) as *mut [T]) }
}

/// Create a new array of `MaybeUninit<T>` items, in an uninitialized state.
///
/// Note: in a future Rust version this method may become unnecessary
/// when Rust allows
/// [inline const expressions](https://github.com/rust-lang/rust/issues/76001).
/// The example below could then use `let mut buf = [const { MaybeUninit::<u8>::uninit() }; 32];`.
///
/// # Examples
///
/// ```
/// use std::mem::MaybeUninit;
///
/// unsafe extern "C" {
///     fn read_into_buffer(
///         ptr: *mut u8,
///         max_len: usize,
///     ) -> usize;
/// }
///
/// /// Returns a (possibly smaller) slice of data that was actually read
/// fn read(buf: &mut [MaybeUninit<u8>]) -> &[u8] {
///     unsafe {
///         // Mocking the C function
///         let len = 0; // read_into_buffer(buf.as_mut_ptr() as *mut u8, buf.len());
///         core::slice::from_raw_parts(buf.as_ptr() as *const u8, len)
///     }
/// }
///
/// let mut buf: [MaybeUninit<u8>; 32] = unsafe { MaybeUninit::uninit().assume_init() };
/// let data = read(&mut buf);
/// ```
// #[unstable(feature = "maybe_uninit_uninit_array", issue = "none")]
// #[rustc_const_unstable(feature = "maybe_uninit_uninit_array", issue = "none")]
#[inline(always)]
fn uninit_array<T, const LEN: usize>() -> [MaybeUninit<T>; LEN] {
    // SAFETY: An uninitialized `[MaybeUninit<_>; LEN]` is valid.
    unsafe { MaybeUninit::<[MaybeUninit<T>; LEN]>::uninit().assume_init() }
}

/// Extracts the values from an array of `MaybeUninit` containers.
///
/// # Safety
///
/// It is up to the caller to guarantee that all elements of the array are
/// in an initialized state.
///
/// # Examples
///
/// ```
/// use std::mem::MaybeUninit;
///
/// let mut array: [MaybeUninit<i32>; 3] = unsafe { MaybeUninit::uninit().assume_init() };
/// array[0].write(0);
/// array[1].write(1);
/// array[2].write(2);
///
/// // SAFETY: Now safe as we initialised all elements
/// let array: [i32; 3] = unsafe { core::ptr::read(&array as *const _ as *const [i32; 3]) };
///
/// assert_eq!(array, [0, 1, 2]);
/// ```
// #[unstable(feature = "maybe_uninit_array_assume_init", issue = "80908")]
#[inline(always)]
pub unsafe fn array_assume_init<T, const N: usize>(array: &[MaybeUninit<T>; N]) -> [T; N] {
    // SAFETY:
    // * The caller guarantees that all elements of the array are initialized
    // * `MaybeUninit<T>` and T are guaranteed to have the same layout
    // * `MaybeUninit` does not drop, so there are no double-frees
    // And thus the conversion is safe
    unsafe {
        // intrinsics::assert_inhabited::<[T; N]>();
        core::ptr::from_ref(array).cast::<[T; N]>().read()
    }
}
