jflib-rt 0.1.0

// mem -- the global allocator and the memory intrinsics the compiler calls
//
// Small blocks (up to 2 KB) come from per-size-class free lists carved out of 64 KB
// chunks; larger ones are whole page runs from the OS (os::pages_alloc), like mem_alloc
// in ASMLIB.  Medium runs (up to 1 MB, rounded to a power of two of pages) are kept in a
// cache when freed (at most CACHE_MAX bytes) and handed out again: worker threads that
// allocate scratch per item then do not pay a VirtualAlloc / mmap, page faults and a free
// each time.
//
// Threads: the small blocks live in SHARDS heaps with a spin lock each; a thread takes the one its
// stack address leads to, so threads that allocate all the time (a simulation on worker threads does
// so for every frame) rarely meet.  Any heap may free any block, as blocks of a class are all alike.
// A thread that only frees would fill its heap while another one only allocates: so before a heap asks
// for a new chunk it takes what the others hold of that class, as far as their locks are free.  The
// cache of page runs has a lock of its own.

use core::alloc::{GlobalAlloc, Layout};
use core::ffi::c_void;
use core::sync::atomic::{AtomicBool, Ordering};

use crate::os;

const CLASSES: usize = 8; // 16, 32, ..., 2048 bytes
const CHUNK: usize = 64 * 1024;
const PAGE: usize = 4096;
const MID: usize = 9; // cached page runs: PAGE << k bytes, k = 0..8 (4 KB .. 1 MB)
const CACHE_MAX: usize = 32 << 20; // bytes of page runs the cache keeps at most
const SHARDS: usize = 64; // heaps of small blocks (a power of two: few threads share one by chance)

struct Heap {
    free: [*mut u8; CLASSES],
    chunk: *mut u8,
    chunk_left: usize,
}

/// a heap of small blocks on its own cache line
#[repr(align(64))]
struct Shard {
    lock: AtomicBool,
    heap: core::cell::UnsafeCell<Heap>,
}

struct Runs {
    mid: [*mut u8; MID], // cached runs (header word 1 = the next run)
    mid_bytes: usize,
}

pub struct Alloc {
    shards: [Shard; SHARDS],
    runs_lock: AtomicBool,
    runs: core::cell::UnsafeCell<Runs>,
}

unsafe impl Sync for Alloc {}

#[cfg_attr(not(feature = "std"), global_allocator)]
static ALLOC: Alloc = Alloc::new();

impl Alloc {
    /// an empty allocator; there is one per program (the `#[global_allocator]`)
    pub const fn new() -> Alloc {
        Alloc {
            shards: [const {
                Shard {
                    lock: AtomicBool::new(false),
                    heap: core::cell::UnsafeCell::new(Heap { free: [core::ptr::null_mut(); CLASSES], chunk: core::ptr::null_mut(), chunk_left: 0 }),
                }
            }; SHARDS],
            runs_lock: AtomicBool::new(false),
            runs: core::cell::UnsafeCell::new(Runs { mid: [core::ptr::null_mut(); MID], mid_bytes: 0 }),
        }
    }
}

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

/// the cache bucket of a page run of `total` bytes (rounded up to PAGE << k), None = too large
fn mid_of(total: usize) -> Option<usize> {
    let pages = total.div_ceil(PAGE);
    let k = (usize::BITS - (pages - 1).leading_zeros()) as usize; // ceil(log2(pages))
    if k < MID { Some(k) } else { None }
}

fn class_of(size: usize) -> Option<usize> {
    let mut c = 0;
    let mut cap = 16;
    while c < CLASSES {
        if size <= cap {
            return Some(c);
        }
        c += 1;
        cap <<= 1;
    }
    None
}

fn lock(l: &AtomicBool) {
    while l.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed).is_err() {
        core::hint::spin_loop();
    }
}

fn unlock(l: &AtomicBool) {
    l.store(false, Ordering::Release);
}

impl Alloc {
    /// the heap of this thread: by where its stack is, in steps of 1 MiB (a thread that crosses one
    /// only changes heaps, and threads with smaller stacks may share one)
    #[inline]
    fn shard(&self) -> usize {
        let probe = 0u8;
        let a = core::ptr::addr_of!(probe) as usize >> 20;
        a.wrapping_mul(0x9E37_79B9_7F4A_7C15_u64 as usize) >> (usize::BITS - SHARDS.trailing_zeros())
    }

    /// what another heap holds of a class, as a list (null: nothing, or their locks are taken); the
    /// caller holds the lock of heap `me`, and no lock is waited for here
    fn steal(&self, me: usize, c: usize) -> *mut u8 {
        for (i, s) in self.shards.iter().enumerate() {
            if i == me || s.lock.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed).is_err() {
                continue;
            }
            // SAFETY: the lock gives exclusive access to that heap's state.
            let list = unsafe { core::mem::replace(&mut (*s.heap.get()).free[c], core::ptr::null_mut()) };
            unlock(&s.lock);
            if !list.is_null() {
                return list;
            }
        }
        core::ptr::null_mut()
    }
}

impl Alloc {
    /// the block and whether its bytes may be non-zero (reused memory)
    unsafe fn alloc_inner(&self, layout: Layout) -> (*mut u8, bool) {
        let size = layout.size().max(layout.align());
        match class_of(size) {
            Some(c) if layout.align() <= 16 => {
                let me = self.shard();
                let s = &self.shards[me];
                lock(&s.lock);
                // SAFETY: the lock gives exclusive access to the heap state.
                let h = unsafe { &mut *s.heap.get() };
                let bytes = 16usize << c;
                let mut p = h.free[c];
                // no block and no room for one: first what the others freed
                if p.is_null() && h.chunk_left < bytes {
                    p = self.steal(me, c);
                }
                let r = if !p.is_null() {
                    // SAFETY: a free block stores the next pointer in its first word.
                    h.free[c] = unsafe { *(p as *mut *mut u8) };
                    p
                } else {
                    if h.chunk_left < bytes {
                        // the rest of the old chunk goes to the free lists
                        while h.chunk_left >= 16 {
                            let mut k = CLASSES - 1;
                            while 16usize << k > h.chunk_left {
                                k -= 1;
                            }
                            // SAFETY: the block lies inside the old chunk.
                            unsafe { *(h.chunk as *mut *mut u8) = h.free[k] };
                            h.free[k] = h.chunk;
                            // SAFETY: the block lies inside the old chunk.
                            h.chunk = unsafe { h.chunk.add(16 << k) };
                            h.chunk_left -= 16 << k;
                        }
                        // SAFETY: plain page allocation.
                        h.chunk = unsafe { os::pages_alloc(CHUNK) };
                        h.chunk_left = if h.chunk.is_null() { 0 } else { CHUNK };
                    }
                    if h.chunk_left < bytes {
                        core::ptr::null_mut()
                    } else {
                        let q = h.chunk;
                        // SAFETY: q + bytes stays inside the chunk.
                        h.chunk = unsafe { h.chunk.add(bytes) };
                        h.chunk_left -= bytes;
                        q
                    }
                };
                unlock(&s.lock);
                (r, true)
            }
            _ if layout.align() > 64 => {
                // over-aligned: the pointer is rounded up inside the run; the run's base and
                // byte count sit right before it
                let Some(total) = size
                    .checked_add(layout.align() + 64 + PAGE - 1)
                    .map(|t| t & !(PAGE - 1))
                else {
                    return (core::ptr::null_mut(), false);
                };
                // SAFETY: plain page allocation.
                let base = unsafe { os::pages_alloc(total) };
                if base.is_null() {
                    return (base, false);
                }
                let user = (base as usize + 64).next_multiple_of(layout.align()) as *mut u8;
                // SAFETY: user - 16 >= base and user + size <= base + total.
                unsafe {
                    *(user as *mut usize).sub(1) = base as usize;
                    *(user as *mut usize).sub(2) = total;
                }
                (user, false)
            }
            _ => {
                // page run with a 64-byte header holding the byte count
                let mut total = (size + 64 + PAGE - 1) & !(PAGE - 1);
                if let Some(k) = mid_of(total) {
                    total = PAGE << k;
                    lock(&self.runs_lock);
                    // SAFETY: exclusive access under the lock; a cached run keeps its byte
                    // count in header word 0 and the next cached run in word 1.
                    let p = unsafe {
                        let h = &mut *self.runs.get();
                        let p = h.mid[k];
                        if !p.is_null() {
                            h.mid[k] = *(p as *mut *mut u8).add(1);
                            h.mid_bytes -= total;
                        }
                        p
                    };
                    unlock(&self.runs_lock);
                    if !p.is_null() {
                        // SAFETY: the run is at least PAGE bytes; its header is word 0.
                        return (unsafe { p.add(64) }, true);
                    }
                }
                // SAFETY: plain page allocation.
                let p = unsafe { os::pages_alloc(total) };
                if p.is_null() {
                    return (p, false);
                }
                // SAFETY: the header lies inside the run.
                unsafe {
                    *(p as *mut usize) = total;
                    (p.add(64), false)
                }
            }
        }
    }
}

unsafe impl GlobalAlloc for Alloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // SAFETY: forwards to alloc_inner.
        unsafe { self.alloc_inner(layout).0 }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let size = layout.size().max(layout.align());
        match class_of(size) {
            Some(c) if layout.align() <= 16 => {
                let s = &self.shards[self.shard()];
                lock(&s.lock);
                // SAFETY: exclusive access under the lock; the block is at least 16 bytes.
                unsafe {
                    let h = &mut *s.heap.get();
                    *(ptr as *mut *mut u8) = h.free[c];
                    h.free[c] = ptr;
                }
                unlock(&s.lock);
            }
            _ if layout.align() > 64 => unsafe {
                // SAFETY: ptr came from the over-aligned path of alloc.
                let base = *(ptr as *mut usize).sub(1) as *mut u8;
                let total = *(ptr as *mut usize).sub(2);
                os::pages_free(base, total);
            },
            _ => unsafe {
                // SAFETY: ptr came from the page path of alloc.
                let base = ptr.sub(64);
                let total = *(base as *mut usize);
                if let Some(k) = mid_of(total).filter(|&k| PAGE << k == total) {
                    lock(&self.runs_lock);
                    let h = &mut *self.runs.get();
                    let keep = h.mid_bytes + total <= CACHE_MAX;
                    if keep {
                        // SAFETY: exclusive access under the lock; header word 1 links the run.
                        *(base as *mut *mut u8).add(1) = h.mid[k];
                        h.mid[k] = base;
                        h.mid_bytes += total;
                    }
                    unlock(&self.runs_lock);
                    if keep {
                        return;
                    }
                }
                os::pages_free(base, total);
            },
        }
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // SAFETY: forwards to alloc_inner; small blocks and cached page runs are reused
        // memory, so clear them.  Fresh page runs come from the OS (VirtualAlloc / mmap: zero
        // pages) and need no clearing.
        let (p, reused) = unsafe { self.alloc_inner(layout) };
        if !p.is_null() && reused {
            unsafe { memset(p as *mut c_void, 0, layout.size()) };
        }
        p
    }
}

// ---- memory intrinsics (no C runtime): rep movsb / stosb, so the compiler cannot turn
// them back into calls to themselves
#[cfg_attr(not(feature = "std"), unsafe(no_mangle))]
pub unsafe extern "C" fn memcpy(dst: *mut c_void, src: *const c_void, n: usize) -> *mut c_void {
    // SAFETY: the caller guarantees valid, non-overlapping ranges.
    unsafe {
        core::arch::asm!("rep movsb", inout("rcx") n => _, inout("rdi") dst => _,
                         inout("rsi") src => _, options(nostack, preserves_flags));
    }
    dst
}

#[cfg_attr(not(feature = "std"), unsafe(no_mangle))]
pub unsafe extern "C" fn memmove(dst: *mut c_void, src: *const c_void, n: usize) -> *mut c_void {
    let (d, s) = (dst as *mut u8, src as *const u8);
    if n == 0 || (d as usize) <= (s as usize) || (d as usize) >= (s as usize) + n {
        // SAFETY: forward copy is correct when dst does not start inside src.
        unsafe { memcpy(dst, src, n) };
    } else {
        // SAFETY: backward copy with the direction flag, cleared again afterwards.
        unsafe {
            core::arch::asm!("std", "rep movsb", "cld",
                             inout("rcx") n => _, inout("rdi") d.add(n - 1) => _,
                             inout("rsi") s.add(n - 1) => _, options(nostack));
        }
    }
    dst
}

#[cfg_attr(not(feature = "std"), unsafe(no_mangle))]
pub unsafe extern "C" fn memset(dst: *mut c_void, c: i32, n: usize) -> *mut c_void {
    // SAFETY: the caller guarantees a valid range.
    unsafe {
        core::arch::asm!("rep stosb", inout("rcx") n => _, inout("rdi") dst => _,
                         in("al") c as u8, options(nostack, preserves_flags));
    }
    dst
}

// 8 bytes through an asm load: the optimizer cannot turn the compare loop into a call to
// memcmp / bcmp (itself)
#[inline(always)]
unsafe fn load8(p: *const u8) -> u64 {
    let v: u64;
    // SAFETY: the caller guarantees 8 readable bytes at p (any alignment).
    unsafe {
        core::arch::asm!("mov {v}, qword ptr [{p}]", v = out(reg) v, p = in(reg) p,
                         options(pure, readonly, nostack, preserves_flags));
    }
    v
}

#[cfg_attr(not(feature = "std"), unsafe(no_mangle))]
pub unsafe extern "C" fn memcmp(a: *const c_void, b: *const c_void, n: usize) -> i32 {
    let (a, b) = (a as *const u8, b as *const u8);
    let mut i = 0;
    // 8 bytes at a time (little endian: the lowest differing byte is the first one)
    while i + 8 <= n {
        // SAFETY: i + 8 <= n, so both ranges hold these 8 bytes.
        let (x, y) = unsafe { (load8(a.add(i)), load8(b.add(i))) };
        if x != y {
            let sh = (x ^ y).trailing_zeros() & !7;
            return ((x >> sh) as u8) as i32 - ((y >> sh) as u8) as i32;
        }
        i += 8;
    }
    while i < n {
        // SAFETY: the caller guarantees valid ranges of n bytes; volatile reads keep the
        // optimizer from turning this loop back into a memcmp call.
        let (x, y) = unsafe { (a.add(i).read_volatile(), b.add(i).read_volatile()) };
        if x != y {
            return x as i32 - y as i32;
        }
        i += 1;
    }
    0
}

#[cfg_attr(not(feature = "std"), unsafe(no_mangle))]
pub unsafe extern "C" fn bcmp(a: *const c_void, b: *const c_void, n: usize) -> i32 {
    // SAFETY: same contract as memcmp.
    unsafe { memcmp(a, b, n) }
}

#[cfg_attr(not(feature = "std"), unsafe(no_mangle))]
pub unsafe extern "C" fn strlen(s: *const core::ffi::c_char) -> usize {
    let s = s as *const u8;
    let mut n = 0;
    // SAFETY: the caller passes a NUL-terminated string; volatile reads keep the loop a loop.
    while unsafe { s.add(n).read_volatile() } != 0 {
        n += 1;
    }
    n
}