platform-mem 0.3.0

use std::{
    alloc::Layout,
    mem::MaybeUninit,
    ops::{Bound, Range, RangeBounds},
};

/// Converts a `RangeBounds` to a `Range` with bounds checking (stable alternative to `slice::range`)
fn range_bounds_to_range<R: RangeBounds<usize>>(range: R, len: usize) -> Range<usize> {
    let start = match range.start_bound() {
        Bound::Included(&n) => n,
        Bound::Excluded(&n) => n.checked_add(1).expect("range start overflow"),
        Bound::Unbounded => 0,
    };
    let end = match range.end_bound() {
        Bound::Included(&n) => n.checked_add(1).expect("range end overflow"),
        Bound::Excluded(&n) => n,
        Bound::Unbounded => len,
    };
    assert!(start <= end, "range start ({start}) > end ({end})");
    assert!(end <= len, "range end ({end}) > length ({len})");
    start..end
}

/// Errors that can occur during memory operations.
#[derive(thiserror::Error, Debug)]
#[non_exhaustive]
pub enum Error {
    /// Error due to the computed capacity exceeding the maximum
    /// (usually `isize::MAX` bytes).
    ///
    /// ## Examples
    ///
    /// grow more than `isize::MAX` bytes:
    ///
    /// ```
    /// # use allocator_api2::alloc::Global;
    /// # use platform_mem::{Error, Alloc, RawMem};
    /// let mut mem = Alloc::new(Global);
    /// assert!(matches!(mem.grow_filled(usize::MAX, 0u64), Err(Error::CapacityOverflow)));
    /// ```
    #[error("exceeding the capacity maximum")]
    CapacityOverflow,

    /// Cannot grow because the requested size exceeds available capacity.
    #[error("can't grow {to_grow} elements, only available {available}")]
    OverGrow {
        /// Number of elements requested.
        to_grow: usize,
        /// Number of elements available.
        available: usize,
    },

    /// The memory allocator returned an error
    #[error("memory allocation of {layout:?} failed")]
    AllocError {
        /// The layout of allocation request that failed
        layout: Layout,

        #[doc(hidden)]
        non_exhaustive: (),
    },

    /// An I/O or system-level error.
    #[error(transparent)]
    System(#[from] std::io::Error),
}

/// Alias for `Result<T, Error>`.
pub type Result<T> = std::result::Result<T, Error>;

/// Unified trait for growable, shrinkable memory regions.
///
/// Implementors manage a contiguous region of initialized `Item` elements.
/// The region can be extended with [`grow`](Self::grow) variants and shortened
/// with [`shrink`](Self::shrink).
pub trait RawMem {
    /// The element type stored in this memory.
    type Item;

    /// Returns a shared slice of all currently initialized elements.
    fn allocated(&self) -> &[Self::Item];
    /// Returns a mutable slice of all currently initialized elements.
    fn allocated_mut(&mut self) -> &mut [Self::Item];

    /// Low-level growth: extends the memory by `cap` elements.
    ///
    /// The `fill` closure receives `(inited, (initialized_slice, uninitialized_slice))`
    /// where `inited` is the count of elements already initialized by the backend.
    ///
    /// # Safety
    /// Caller must guarantee that `fill` fully initializes the uninitialized portion.
    ///
    /// ### Incorrect usage
    /// ```no_run
    /// # use allocator_api2::alloc::Global;
    /// # use std::mem::MaybeUninit;
    /// # use platform_mem::Result;
    /// use platform_mem::{Alloc, RawMem};
    ///
    /// let mut alloc = Alloc::new(Global);
    /// unsafe {
    ///     alloc.grow(10, |_init, (_, _uninit): (_, &mut [MaybeUninit<u64>])| {
    ///         // `RawMem` relies on the fact that we initialize memory
    ///         // even if they are primitives
    ///     })?;
    /// }
    /// # Result::Ok(())
    /// ```
    unsafe fn grow(
        &mut self,
        cap: usize,
        fill: impl FnOnce(usize, (&mut [Self::Item], &mut [MaybeUninit<Self::Item>])),
    ) -> Result<&mut [Self::Item]>;

    /// Removes the last `cap` elements, dropping them.
    fn shrink(&mut self, cap: usize) -> Result<()>;

    /// Returns an optional hint about the total capacity, if known.
    fn size_hint(&self) -> Option<usize> {
        None
    }

    /// [`grow`] which assumes that the memory is already initialized
    ///
    /// # Safety
    ///
    /// When calling this method, you have to ensure that one of the following is true:
    ///
    /// * memory already initialized as [`Item`]
    ///
    /// * memory is initialized bytes and [`Item`] can be represented as bytes
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use platform_mem::Result;
    /// use platform_mem::{FileMapped, RawMem};
    ///
    /// let mut file = FileMapped::from_path("..")?;
    /// // file is always represents as initialized bytes
    /// // and usize is transparent as bytes
    /// let _: &mut [usize] = unsafe { file.grow_assumed(10)? };
    /// # Result::Ok(())
    /// ```
    ///
    /// [`grow`]: Self::grow
    /// [`Item`]: Self::Item
    unsafe fn grow_assumed(&mut self, cap: usize) -> Result<&mut [Self::Item]> {
        unsafe {
            self.grow(cap, |inited, (_, uninit)| {
                // fixme: maybe change it to `assert_eq!`
                debug_assert_eq!(
                    inited,
                    uninit.len(),
                    "grown memory must be initialized, \
                 usually allocators-like provide uninitialized memory, \
                 which is only safe for writing"
                )
            })
        }
    }

    /// # Safety
    /// [`Item`](Self::Item) must satisfy the [initialization invariant][inv] for
    /// [`core::mem::zeroed`].
    ///
    ///  [inv]: MaybeUninit#initialization-invariant
    ///
    /// # Examples
    /// Correct usage of this function: initializing an integral-like types with zeroes:
    /// ```
    /// # use platform_mem::Error;
    /// use platform_mem::{Global, RawMem};
    ///
    /// let mut alloc = Global::new();
    /// let zeroes: &mut [(u8, u16)] = unsafe {
    ///     alloc.grow_zeroed(10)?
    /// };
    ///
    /// assert_eq!(zeroes, [(0, 0); 10]);
    /// # Ok::<_, Error>(())
    /// ```
    ///
    /// Incorrect usage of this function: initializing a reference with zero:
    /// ```no_run
    /// # use platform_mem::Error;
    /// use platform_mem::{Global, RawMem};
    ///
    /// let mut alloc = Global::new();
    /// let zeroes: &mut [&'static str] = unsafe {
    ///     alloc.grow_zeroed(10)? // Undefined behavior!
    /// };
    ///
    /// # Ok::<_, Error>(())
    /// ```
    ///
    unsafe fn grow_zeroed(&mut self, cap: usize) -> Result<&mut [Self::Item]> {
        unsafe {
            self.grow(cap, |_, (_, uninit)| {
                uninit.as_mut_ptr().write_bytes(0u8, uninit.len());
            })
        }
    }

    /// # Safety
    /// [`Item`](Self::Item) must satisfy the [initialization invariant](MaybeUninit#initialization-invariant) for zeroed memory.
    unsafe fn grow_zeroed_exact(&mut self, cap: usize) -> Result<&mut [Self::Item]> {
        unsafe {
            self.grow(cap, |inited, (_, uninit)| {
                uninit.get_unchecked_mut(inited..).as_mut_ptr().write_bytes(0u8, uninit.len());
            })
        }
    }

    /// Grows by `addition` elements, initializing each with the closure `f`.
    fn grow_with(
        &mut self,
        addition: usize,
        f: impl FnMut() -> Self::Item,
    ) -> Result<&mut [Self::Item]> {
        unsafe {
            self.grow(addition, |_, (_, uninit)| {
                uninit::fill_with(uninit, f);
            })
        }
    }

    /// # Safety
    /// The caller must ensure that `addition` accounts for already-initialized elements in the underlying buffer.
    unsafe fn grow_with_exact(
        &mut self,
        addition: usize,
        f: impl FnMut() -> Self::Item,
    ) -> Result<&mut [Self::Item]> {
        unsafe {
            self.grow(addition, |inited, (_, uninit)| {
                uninit::fill_with(&mut uninit[inited..], f);
            })
        }
    }

    /// Grows by `cap` elements, each cloned from `value`.
    fn grow_filled(&mut self, cap: usize, value: Self::Item) -> Result<&mut [Self::Item]>
    where
        Self::Item: Clone,
    {
        unsafe {
            self.grow(cap, |_, (_, uninit)| {
                uninit::fill(uninit, value);
            })
        }
    }

    /// # Safety
    /// The caller must ensure that `cap` accounts for already-initialized elements in the underlying buffer.
    unsafe fn grow_filled_exact(
        &mut self,
        cap: usize,
        value: Self::Item,
    ) -> Result<&mut [Self::Item]>
    where
        Self::Item: Clone,
    {
        unsafe {
            self.grow(cap, |inited, (_, uninit)| {
                uninit::fill(&mut uninit[inited..], value);
            })
        }
    }

    /// Grows by cloning a sub-range of the currently allocated data.
    fn grow_within<R: RangeBounds<usize>>(&mut self, range: R) -> Result<&mut [Self::Item]>
    where
        Self::Item: Clone,
    {
        let Range { start, end } = range_bounds_to_range(range, self.allocated().len());
        unsafe {
            self.grow(end - start, |_, (within, uninit)| {
                uninit::write_clone_of_slice(uninit, &within[start..end]);
            })
        }
    }

    /// Grows by cloning all elements from `src`.
    fn grow_from_slice(&mut self, src: &[Self::Item]) -> Result<&mut [Self::Item]>
    where
        Self::Item: Clone,
    {
        unsafe {
            self.grow(src.len(), |_, (_, uninit)| {
                uninit::write_clone_of_slice(uninit, src);
            })
        }
    }
}

/// A callable trait for fill functions, usable as a trait object.
/// This is a stable alternative to implementing `FnMut` manually.
pub trait FillFn<T> {
    /// Invoke the fill function with the given initialized count and memory slices.
    fn call(&mut self, inited: usize, slices: (&mut [T], &mut [MaybeUninit<T>]));
}

/// Implements `FillFn` for any `FnMut` closure.
impl<T, F> FillFn<T> for F
where
    F: FnMut(usize, (&mut [T], &mut [MaybeUninit<T>])),
{
    fn call(&mut self, inited: usize, slices: (&mut [T], &mut [MaybeUninit<T>])) {
        self(inited, slices)
    }
}

/// A wrapper that allows calling a `FnOnce` through the `FillFn` interface.
/// Used to pass `FnOnce` closures to erased trait objects that expect `FillFn`.
struct CallOnce<F> {
    inner: Option<F>,
}

impl<F> CallOnce<F> {
    fn new(f: F) -> Self {
        Self { inner: Some(f) }
    }
}

impl<T, F> FillFn<T> for CallOnce<F>
where
    F: FnOnce(usize, (&mut [T], &mut [MaybeUninit<T>])),
{
    fn call(&mut self, inited: usize, slices: (&mut [T], &mut [MaybeUninit<T>])) {
        let f = self.inner.take().expect("CallOnce::call called more than once");
        f(inited, slices);
    }
}

/// # Safety
/// Implementors must uphold the same memory-safety invariants as [`RawMem`].
pub unsafe trait ErasedMem {
    /// The element type.
    type Item;

    /// Returns a shared slice of initialized elements (type-erased).
    fn erased_allocated(&self) -> &[Self::Item];
    /// Returns a mutable slice of initialized elements (type-erased).
    fn erased_allocated_mut(&mut self) -> &mut [Self::Item];

    /// # Safety
    /// The caller must guarantee that `fill` fully initializes the uninitialized portion.
    unsafe fn erased_grow(
        &mut self,
        cap: usize,
        fill: &mut dyn FillFn<Self::Item>,
    ) -> Result<&mut [Self::Item]>;

    /// Removes the last `cap` elements (type-erased).
    fn erased_shrink(&mut self, cap: usize) -> Result<()>;

    /// Returns an optional capacity hint (type-erased).
    fn erased_size_hint(&self) -> Option<usize> {
        None
    }
}

macro_rules! impl_erased {
    ($ty:ty => $($imp:tt)+) => {
        impl $($imp)+ {
            type Item = $ty;

            fn allocated(&self) -> &[Self::Item] {
                (**self).erased_allocated()
            }

            fn allocated_mut(&mut self) -> &mut [Self::Item] {
                (**self).erased_allocated_mut()
            }

            unsafe fn grow(
                &mut self,
                cap: usize,
                fill: impl FnOnce(usize, (&mut [Self::Item], &mut [MaybeUninit<Self::Item>])),
            ) -> Result<&mut [Self::Item]> { unsafe {
                (**self).erased_grow(cap, &mut CallOnce::new(fill))
            }}

            fn shrink(&mut self, cap: usize) -> Result<()> {
                (**self).erased_shrink(cap)
            }

            fn size_hint(&self) -> Option<usize> {
                (**self).erased_size_hint()
            }
        }
    };
}

impl_erased!(All::Item => <'a, All: ?Sized + RawMem> RawMem for &'a mut All);

impl_erased!(I => <'a, I> RawMem for Box<dyn ErasedMem<Item = I> + 'a>);
impl_erased!(I => <'a, I> RawMem for Box<dyn ErasedMem<Item = I> + Sync + 'a>);
impl_erased!(I => <'a, I> RawMem for Box<dyn ErasedMem<Item = I> + Sync + Send + 'a>);

unsafe impl<All: RawMem + ?Sized> ErasedMem for All {
    type Item = All::Item;

    fn erased_allocated(&self) -> &[Self::Item] {
        self.allocated()
    }

    fn erased_allocated_mut(&mut self) -> &mut [Self::Item] {
        self.allocated_mut()
    }

    unsafe fn erased_grow(
        &mut self,
        cap: usize,
        fill: &mut dyn FillFn<Self::Item>,
    ) -> Result<&mut [Self::Item]> {
        unsafe { self.grow(cap, |inited, slices| fill.call(inited, slices)) }
    }

    fn erased_shrink(&mut self, cap: usize) -> Result<()> {
        self.shrink(cap)
    }

    fn erased_size_hint(&self) -> Option<usize> {
        self.size_hint()
    }
}

/// Utilities for initializing `MaybeUninit` slices.
///
/// These are stable alternatives to nightly-only `MaybeUninit` slice methods.
pub mod uninit {
    use std::{mem, mem::MaybeUninit, ptr, slice};

    /// Stable alternative to `MaybeUninit::slice_assume_init_mut`.
    ///
    /// # Safety
    /// All elements of `s` must be initialized.
    #[inline]
    pub unsafe fn assume_init_mut<T>(s: &mut [MaybeUninit<T>]) -> &mut [T] {
        // SAFETY: The caller guarantees all elements are initialized.
        // MaybeUninit<T> has the same layout as T.
        unsafe { slice::from_raw_parts_mut(s.as_mut_ptr().cast::<T>(), s.len()) }
    }

    /// Stable alternative to `MaybeUninit::write_slice_cloned` / `write_clone_of_slice`.
    ///
    /// # Panics
    /// Panics if `uninit.len() != src.len()`.
    pub fn write_clone_of_slice<T: Clone>(uninit: &mut [MaybeUninit<T>], src: &[T]) {
        assert_eq!(uninit.len(), src.len(), "slice lengths must match");

        for (dst, val) in uninit.iter_mut().zip(src.iter()) {
            dst.write(val.clone());
        }
    }

    /// Initializes all elements of `uninit` by cloning `val`. Panic-safe.
    pub fn fill<T: Clone>(uninit: &mut [MaybeUninit<T>], val: T) {
        let mut guard = Guard { slice: uninit, init: 0 };

        if let Some((last, elems)) = guard.slice.split_last_mut() {
            for el in elems.iter_mut() {
                el.write(val.clone());
                guard.init += 1;
            }
            last.write(val);
            guard.init += 1;
        }

        mem::forget(guard);
    }

    /// Initializes all elements of `uninit` using the closure `fill`. Panic-safe.
    pub fn fill_with<T>(uninit: &mut [MaybeUninit<T>], mut fill: impl FnMut() -> T) {
        let mut guard = Guard { slice: uninit, init: 0 };

        for el in guard.slice.iter_mut() {
            el.write(fill());
            guard.init += 1;
        }

        mem::forget(guard);
    }

    struct Guard<'a, T> {
        slice: &'a mut [MaybeUninit<T>],
        init: usize,
    }

    impl<T> Drop for Guard<'_, T> {
        fn drop(&mut self) {
            debug_assert!(self.init <= self.slice.len());
            // SAFETY: this raw slice will contain only initialized objects
            // that's why, it is allowed to drop it.
            unsafe {
                let inited_slice = self.slice.get_unchecked_mut(..self.init);
                ptr::drop_in_place(assume_init_mut(inited_slice));
            }
        }
    }
}