realloc 0.1.0

use crate::{Align, Alloc, Allocator, Brand, Error, Layout};
use core::ptr::NonNull;

unsafe extern "C" {
    fn _aligned_malloc(size: usize, align: usize) -> *mut core::ffi::c_void;
    fn _aligned_realloc(
        ptr: *mut core::ffi::c_void,
        size: usize,
        align: usize,
    ) -> *mut core::ffi::c_void;
    fn _aligned_free(ptr: *mut core::ffi::c_void);
}

/// A "handle" for the platform's allocator.
///
/// On Windows, this uses `_aligned_malloc`.
///
/// Only available on feature `platform`.
pub struct PlatformAllocator(Brand);

impl PlatformAllocator {
    /// Return a handle to the platform allocator.
    ///
    /// Multiple of these can safely exist and be used concurrently, since the
    /// platform allocator is thread-safe.
    pub const fn new() -> Self {
        // SAFETY: this is the platform allocator
        Self(unsafe { crate::Brand::platform() })
    }
}

impl Default for PlatformAllocator {
    fn default() -> Self { Self::new() }
}

unsafe impl Allocator for PlatformAllocator {
    fn brand(&self) -> &Brand {
        &self.0
    }

    fn alloc_bytes(&self, layout: Layout) -> Result<Alloc<'_, u8>, Error> {
        let align = layout.align.align();

        if layout.size == 0 {
            // SAFETY:
            // - `ptr` was just "allocated" with this allocator
            // - `layout` is the layout used to "allocate" `ptr`
            // - This is not being used to double-drop an allocation
            //
            // When freed, this allocator will simply do nothing, as it did
            // during allocation.
            return Ok(unsafe {
                Alloc::new(NonNull::new(align as *mut u8).unwrap(), layout, self)
            });
        }

        // SAFETY: `size` is non-zero and `align` is a power of 2
        let Some(ptr) = NonNull::new(unsafe { _aligned_malloc(layout.size, align) }) else {
            return Err(Error::PlatformError);
        };

        // SAFETY:
        // - `ptr` was just allocated with this allocator
        // - `layout` is the layout used to allocate `ptr`
        // - This is not being used to double-drop an allocation
        Ok(unsafe { Alloc::new(ptr.cast(), layout, self) })
    }

    fn dealloc_bytes<'this>(&'this self, alloc: Alloc<'this, u8>) {
        if alloc.allocator().brand() != &self.0 || alloc.layout().size == 0 {
            return;
        }

        // SAFETY: the pointer was gotten from this allocator, thus it must also
        // be from `aligned_alloc`
        unsafe { _aligned_free(alloc.as_ptr().cast()) };
    }

    fn dangling(&self, align: Align) -> Alloc<'_, u8> {
        let layout = Layout { size: 0, align };
        unsafe {
            Alloc::new(
                NonNull::new(align.align() as *mut u8).unwrap(),
                layout,
                self,
            )
        }
    }

    fn realloc<'this>(
        &'this self,
        alloc: Alloc<'this, u8>,
        layout: Layout,
    ) -> Result<Alloc<'this, u8>, Error> {
        if alloc.allocator().brand() != &self.0 {
            return Err(Error::WrongAllocator);
        } else if alloc.as_ptr() as usize % layout.align.align() != 0
            || alloc.align() > layout.align
        {
            // `_aligned_realloc` doesn't let you change the alignment of an
            // allocation. So if the alignment isn't already suitable, it must
            // be totally reallocated.

            let new = self.alloc_bytes(layout)?;
            let size = new.size();
            // SAFETY: this ptr is valid for exactly `new.size()` bytes
            unsafe { new.ptr().cast::<u8>().copy_from(alloc.ptr().cast(), size) };

            return Ok(new);
        } else if layout.size == 0 {
            if layout.align <= alloc.align() {
                // The allocation is already at least as large, and properly
                // aligned
                return Ok(alloc);
            } else {
                alloc.dealloc();
                // SAFETY:
                // - `ptr` was just "allocated" with this allocator
                // - `layout` is the layout used to "allocate" `ptr`
                // - This is not being used to double-drop an allocation
                //
                // When freed, this allocator will simply do nothing, as it did
                // during allocation.
                return Ok(unsafe {
                    Alloc::new(
                        NonNull::new(layout.align.align() as *mut u8).unwrap(),
                        layout,
                        self,
                    )
                });
            }
        }

        let layout = Layout {
            size: layout.size,
            align: alloc.align(),
        };

        // SAFETY: `size` is non-zero and `align` is a power of 2
        let Some(ptr) = NonNull::new(unsafe {
            _aligned_realloc(alloc.as_ptr().cast(), layout.size, layout.align.align())
        }) else {
            return Err(Error::PlatformError);
        };

        // SAFETY:
        // - `ptr` was just allocated with this allocator
        // - `layout` is the layout used to allocate `ptr` (with the exception
        //   that `align` may be larger than necessary)
        // - This is not being used to double-drop an allocation
        Ok(unsafe { Alloc::new(ptr.cast(), layout, self) })
    }

    fn grow<'this>(
        &'this self,
        alloc: Alloc<'this, u8>,
        layout: Layout,
    ) -> Result<Alloc<'this, u8>, Error> {
        self.realloc(alloc, layout)
    }

    fn shrink<'this>(
        &'this self,
        alloc: Alloc<'this, u8>,
        layout: Layout,
    ) -> Result<Alloc<'this, u8>, Error> {
        self.realloc(alloc, layout)
    }
}