mi_heap.rs

#![cfg(feature = "nightly", feature(core_intrinsics))]
use libmimalloc_sys::{self as mi};
use core::ptr::NonNull;
use std::alloc::Layout;
#[cfg(feature = "nightly")]
use core::intrinsics::{likely, unlikely};
#[cfg(not(feature = "nightly"))]
#[inline(always)] fn likely(b: bool) -> bool { b }
#[cfg(not(feature = "nightly"))]
#[inline(always)] fn unlikely(b: bool) -> bool { b }

#[derive(Debug)]
pub struct MiHeap {
    ptr: NonNull<mi::mi_heap_t>,
    on_drop: BehaviorOnDrop,
}

impl PartialEq for MiHeap {
    #[inline]
    fn eq(&self, o: &MiHeap) -> bool {
        self.ptr == o.ptr
    }
}

impl Drop for MiHeap {
    fn drop(&mut self) {
        match self.on_drop {
            BehaviorOnDrop::DoNothing => {}
            BehaviorOnDrop::ReleaseUnused => {
                unsafe { mi::mi_heap_delete(self.ptr.as_ptr()) }
            }
            BehaviorOnDrop::DangerouslyFreeAllMemory => {
                unsafe { mi::mi_heap_destroy(self.ptr.as_ptr()) }
            }
        }
    }
}

impl MiHeap {
    /// Create a mimalloc heap, panicing on failure (e.g. due to out of memory)
    #[inline]
    pub fn new() -> Self {
        Self::try_new().expect("Failed to allocate a mimalloc heap!")
    }

    /// Return our internal `mi_heap_t`.
    #[inline]
    pub fn into_raw(self) -> NonNull<mi::mi_heap_t> {
        let p = self.ptr;
        core::mem::forget(self);
        p
    }

    /// Wrap a raw `mi_heap_t` in a `MiHeap`.
    #[inline]
    pub unsafe fn from_raw(p: NonNull<mi::mi_heap_t>, drop_behavior: BehaviorOnDrop) -> Self {
        Self { ptr: p, on_drop: drop_behavior }
    }

    /// Create a mimalloc heap, if possible.
    pub fn try_new() -> Option<Self> {
        let p = unsafe { NonNull::new(mi::mi_heap_new())? };
        Some(Self { ptr: p, on_drop: BehaviorOnDrop::ReleaseUnused })
    }

    /// Get the default heap the current thread.
    ///
    /// Note that when the object returned by this goes out of scope, it will
    /// not release the heap's resources by default.
    pub fn get_default() -> Self {
        let ptr = unsafe { NonNull::new(mi::mi_heap_get_default()).expect("mi_heap_get_default returned null") };
        Self { ptr, on_drop: BehaviorOnDrop::DoNothing }
    }

    /// Get the backing heap the current thread.
    pub fn get_backing() -> Self {
        let ptr = unsafe { NonNull::new(mi::mi_heap_get_backing()).expect("mi_heap_get_backing returned null") };
        Self { ptr, on_drop: BehaviorOnDrop::DoNothing }
    }

    /// Return the current drop behavior for the heap.
    ///
    /// By default, this is `ReleaseUnused` for heaps created by `MiHeap::new`
    /// (or try_new) (as that is generally what you want)  and `DoNothing` for
    /// ones returned from `MiHeap::get_default` or `MiHeap::get_backing`.
    pub fn get_drop_behavior(&self) -> BehaviorOnDrop {
        self.on_drop
    }

    /// Return unused resources to the OS.
    pub fn collect(&self, force: bool) {
        unsafe { mi::mi_heap_collect(self.ptr.as_ptr(), force); }
    }

    /// Set the drop behavior for a heap.
    ///
    /// ## Panics
    ///
    /// Panics if behavior is not `DoNothing` and `self` is the backing heap.
    ///
    /// ## Safety
    ///
    /// - It's always safe to set this to `BehaviorOnDrop::DoNothing`.
    ///
    /// - It's usually safe to set it to `BehaviorOnDrop::ReleaseUnused`, unless
    ///   you've managed to get two copies of the same heap (for example, if you
    ///   they were returned by `get_default`/`get_backing`).
    ///
    /// - It's very dangerous to set it to `DangerouslyFreeAllMemory`, and only
    ///   safe if you are certain all pointers allocated by this heap are no
    ///   longer in use.
    pub unsafe fn set_drop_behavior(&mut self, bhvr: BehaviorOnDrop) {
        if bhvr != BehaviorOnDrop::DoNothing {
            assert!(self.ptr.as_ptr() != mi::mi_heap_get_backing());
        }
        self.on_drop = bhvr;
    }

    /// Set `self` to be the thread's default heap, and safely reclaim the
    /// previous current heap's resources.
    #[must_use]
    pub unsafe fn set_default(&mut self) {
        let old_cur = mi::mi_heap_set_default(self.ptr.as_ptr());
        if !old_cur.is_null() && old_cur != self.ptr.as_ptr() && old_cur != mi::mi_heap_get_backing() {
            mi::mi_heap_delete(old_cur);
        }
    }

    /// Return true if this heap owns the given pointer.
    ///
    /// Note: This function is fairly expensive, and is a linear search over the
    /// blocks.
    pub fn check_owned(&self, p: *const ()) -> bool {
        unsafe { mi::mi_heap_check_owned(self.ptr.as_ptr(), p as *const _) }
    }

    /// The maximum size that can be used with mi::MI_SMALL_SIZE_MAX.
    /// This is currently 128 * size_of::<usize>()
    pub const ALLOC_SMALL_MAX_SIZE: usize = mi::MI_SMALL_SIZE_MAX;
    pub const DEFAULT_ALIGN: usize = 16;

    /// Allocate `v` bytes. v must be less than ALLOC_SMALL_MAX_SIZE
    #[inline]
    pub unsafe fn raw_alloc_small(&self, v: usize) -> *mut u8 {
        debug_assert!(v < Self::ALLOC_SMALL_MAX_SIZE && v != 0, "not small: {:?}", v);
        mi::mi_heap_malloc_small(self.ptr.as_ptr(), v) as *mut u8
    }

    #[inline]
    pub unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        let mut size = l.size();
        let align = l.align();
        if unlikely(size < align) {
            debug_assert!(size != 0);
            size = align;
        }
        if likely(align <= DEFAULT_ALIGN) {
            if size <= Self::ALLOC_SMALL_MAX_SIZE {
                mi::mi_heap_malloc_small(self.ptr.as_ptr(), size) as *mut u8
            } else {
                mi::mi_heap_malloc(self.ptr.as_ptr(), size) as *mut u8
            }
        } else {
            mi::mi_heap_malloc_aligned(self.ptr.as_ptr(), size, align) as *mut u8
        }
    }

    #[inline]
    pub unsafe fn free(&self, p: *mut u8, _: Layout) {
        if !p.is_null() {
            mi::mi_free(p as *mut _)
        }
    }

    #[inline]
    pub unsafe fn realloc(&self, p: *mut u8, l: Layout, new_size: usize) -> *mut u8 {
        let align = l.align();
        if likely(align <= DEFAULT_ALIGN) {
            mi::mi_heap_realloc(self.ptr.as_ptr(), p as *mut _, new_size) as *mut u8
        } else {
            mi::mi_heap_realloc_aligned(self.ptr.as_ptr(), p as *mut _, new_size, align) as *mut u8
        }
    }

    #[inline]
    pub fn is_in_heap(p: *mut u8) -> bool {
        unsafe { mi::mi_is_in_heap_region(p as *mut _) }
    }

    #[inline]
    pub unsafe fn usable_size(p: *mut u8) -> usize {
        mi::mi_usable_size(p as *mut _)
    }

    #[inline]
    pub fn good_size(p: usize) -> usize {
        unsafe { mi::mi_good_size(p) }
    }
}

#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum BehaviorOnDrop {
    /// On drop, we do nothing. This is the default for heaps which are returned
    /// by `MiHeap::get_default` (usually the right call) and
    /// `MiHeap::get_backing` (always the right call).
    DoNothing,

    /// Release unused resources back to the current thread-local heap. This is
    /// the default for heaps returned by `MiHeap::new` and `MiHeap::try_new`
    ReleaseUnused,

    /// ⚠️ Dangerous! This means we will free all pointers which had been
    /// returned by this heap, even though they may still be in use!
    DangerouslyFreeAllMemory,
}