ps-alloc 0.1.0-9

use std::ptr::null_mut;

use crate::{
    free::free_block,
    header::{initialize_block, layout_of, total_size, validate, AllocationHeader},
    AllocationError, DeallocationError, ReallocationError,
};

/// Reallocates memory allocated by [`crate::alloc`].
///
/// Mirrors C's `realloc`:
///
/// - `realloc(null, new_size)` is equivalent to [`crate::alloc`]`(new_size)`.
/// - `realloc(ptr, 0)` frees `ptr` and returns a null pointer.
/// - Otherwise, the allocation is resized, in place when possible, and its contents are
///   preserved up to the lesser of the old and new sizes. A `new_size` that rounds up to
///   the allocation's current total size returns `ptr` unchanged.
///
/// On success, the old pointer must be considered invalid: it must not be dereferenced
/// or passed to [`crate::free`] or this function again. Only the returned pointer may be
/// used.
///
/// # Errors
///
/// The old pointer remains valid whenever an error is returned.
///
/// - `AllocationError` is returned if `ptr` is null and the allocation fails, or if
///   `new_size` overflows the size computation or does not form a valid layout.
/// - `DeallocationError` is returned if the size stored in the header does not form a
///   valid layout, which signals memory corruption.
/// - `ImproperAlignment` is returned if the pointer provided was not aligned properly.
/// - `MarkerFree` is returned if `ptr` was previously freed or reallocated.
/// - `MarkerCorrupted` is returned if the marker is neither free nor used, which
///   generally signals memory corruption, or the passing of a pointer not allocated by
///   [`crate::alloc`].
/// - `NewAllocationFailed` is returned if the allocator failed to resize the allocation.
///
/// # Safety
///
/// See [`crate::free`] for safety information. Additionally, best-effort detection of
/// stale pointers after a successful `realloc` is only possible when the block has
/// moved; after an in-place resize, the old pointer aliases the live allocation and
/// misuse cannot be detected.
pub unsafe fn realloc(ptr: *mut u8, new_size: usize) -> Result<*mut u8, ReallocationError> {
    if ptr.is_null() {
        return Ok(crate::alloc(new_size)?);
    }

    // SAFETY: the caller's obligations for `realloc` include those of [`validate`].
    let header_ptr = unsafe { validate(ptr) }?;

    if new_size == 0 {
        // This is equivalent to `free` and returns a nullptr.

        // SAFETY: `validate` succeeded, so `header_ptr` satisfies `free_block`'s contract.
        unsafe { free_block(header_ptr) }?;

        return Ok(null_mut());
    }

    // SAFETY: `validate` succeeded, so `header_ptr` points to a live, initialized header.
    let old_total = unsafe { &*header_ptr }.get_size();

    let new_total = total_size(new_size)?;

    if new_total == old_total {
        // the allocation already has the requested capacity
        return Ok(ptr);
    }

    // reconstructing the old layout can only fail if the header was corrupted
    let old_layout = layout_of(old_total).map_err(DeallocationError::from)?;

    // pre-validate the new layout so the raw realloc call below cannot be handed an
    // invalid size
    layout_of(new_total).map_err(AllocationError::from)?;

    // SAFETY: as above; the freed marker is written before the allocator is invoked so
    // that, if the block moves, the abandoned memory retains the freed marker for
    // best-effort stale-pointer detection.
    unsafe { &mut *header_ptr }.mark_as_free();

    // SAFETY: `header_ptr` is the base of an allocation created with `old_layout`, and
    // `new_total` is nonzero and forms a valid layout, as checked above.
    let new_base = unsafe { std::alloc::realloc(header_ptr.cast(), old_layout, new_total) };

    if new_base.is_null() {
        // the old block is untouched on failure; restore its header
        // SAFETY: `header_ptr` still points to the live header of the old block.
        unsafe { header_ptr.write(AllocationHeader::new(old_total)) };

        return Err(ReallocationError::NewAllocationFailed);
    }

    // SAFETY: `new_base` is non-null, sufficiently aligned, and valid for writes of
    // `new_total` bytes.
    Ok(unsafe { initialize_block(new_base, new_total) })
}