wmaker/wutil-rs/src/memory.rs

//! Custom implementations of malloc/free/realloc.
//!
//! These are intended for use by C functions that need to allocate. Window
//! Maker originally provided [`wmalloc`], [`wfree`], and [`wrealloc`] for
//! customizable handling of memory exhaustion (to save workspace state before
//! aborting) and to allow optional use of the Boehm GC library. It also tracked
//! reference counts, via [`wretain`] and [`wrelease`].
//!
//! If everything gets rewritten in Rust, we won't need this module anymore. For
//! now, it helps to move our allocations into Rust so that it is more
//! straightforward to store Rust objects in heap memory that was allocated from
//! C. (Rust may have stricter requirements for heap-allocated segments than are
//! provided by arbitrary C allocators).
//!
//! TODO: We may want to restore handling of OOM errors. This would require
//! installing a customized Rust allocator, which isn't something you can do yet
//! in stable Rust. And, unless our rewrite ends up taking up obscenely more
//! memory than the baseline Window Maker code, it isn't really necessary in
//! this day and age.

use std::{alloc, ffi::{c_char, CStr}, mem, ptr::{self, NonNull}};

/// Tracks the layout and reference count of an allocated chunk of memory.
#[derive(Clone, Copy)]
struct Header {
    ptr: NonNull<u8>,
    payload_size: usize,
    layout: alloc::Layout,
    refcount: u16,
}

impl Header {
    /// Recovers the `Header` for the allocated memory chunk `b`.
    ///
    /// ## Safety
    ///
    /// Callers must ensure that `b` is a live allocation from [`wmalloc`] or [`wrealloc`].
    unsafe fn for_alloc_bytes(b: *mut u8) -> *mut Header {
        unsafe {
            b.sub(mem::size_of::<Header>())
                .cast::<Header>()
        }
    }
}

/// Allocates at least `size` bytes and returns a pointer to them.
///
/// Returns null if `size` is 0.
pub fn alloc_bytes(size: usize) -> *mut u8 {
    if size == 0 {
        return ptr::null_mut();
    }
    let Ok(header_layout) = alloc::Layout::from_size_align(mem::size_of::<Header>(), 8) else {
        return ptr::null_mut();
    };

    let Ok(layout) = alloc::Layout::from_size_align(size, 8) else {
        return ptr::null_mut();
    };
    let Ok((full_layout, result_offset)) = header_layout.extend(layout) else {
        return ptr::null_mut();
    };

    let full_segment = unsafe { alloc::alloc_zeroed(full_layout) };
    if full_segment.is_null() {
        return ptr::null_mut();
    }
    let result = unsafe { full_segment.add(result_offset) };
    if result.is_null() {
        return ptr::null_mut();
    }

    unsafe {
        let header = result.sub(mem::size_of::<Header>()).cast::<Header>();
        header.write_unaligned(Header {
            ptr: NonNull::new_unchecked(full_segment),
            payload_size: size,
            layout: full_layout,
            refcount: 0,
        });
    }

    result
}

/// Allocates a segment with [`alloc_bytes`] and fills it with the contents of
/// `s`. The resulting string should be free'd by passing it to [`free_bytes`].
pub fn alloc_string(s: &CStr) -> *mut c_char {
    let len = s.count_bytes() + 1;
    let result = alloc_bytes(len).cast::<c_char>();
    unsafe { ptr::copy_nonoverlapping(s.as_ptr().cast::<c_char>(), result, len); }
    result
}

/// Frees the bytes pointed to by `b`.
///
/// ## Safety
///
/// Callers must ensure that `b` is a live allocation from [`wmalloc`] or [`wrealloc`].
pub unsafe fn free_bytes(b: *mut u8) {
    if b.is_null() {
        return;
    }
    unsafe {
        let header = &*Header::for_alloc_bytes(b);
        alloc::dealloc(header.ptr.as_ptr(), header.layout);
    }
}

/// Functions to be called from C.
pub mod ffi {
    use super::{alloc_bytes, free_bytes, Header};

    use std::{ffi::c_void, ptr};

    /// Allocates `size` bytes. Returns null if `size` is 0. Data will be
    /// initialized to 0.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn wmalloc(size: usize) -> *mut c_void {
        alloc_bytes(size).cast::<c_void>()
    }

    /// Frees `ptr`, which must have come from [`wmalloc`] or [`wrealloc`].
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn wfree(ptr: *mut c_void) {
        unsafe { free_bytes(ptr.cast::<u8>()); }
    }

    /// Resizes `ptr` to be at least `newsize` bytes in size, returning the
    /// start of the new segment. If `newsize` is larger than `ptr`'s segment,
    /// data in the new space will be initialized but have an arbitrary value.
    ///
    /// ## Safety
    ///
    /// If `ptr` is non-null, callers must ensure that it came from from
    /// [`wmalloc`] or [`wrealloc`].
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn wrealloc(ptr: *mut c_void, newsize: usize) -> *mut c_void {
        if ptr.is_null() {
            unsafe {
                return wmalloc(newsize);
            }
        }
        unsafe {
            let result = wmalloc(newsize);
            let result_header = ptr::read_unaligned(Header::for_alloc_bytes(result.cast()));
            let ptr_header = ptr::read_unaligned(Header::for_alloc_bytes(ptr.cast()));
            let copy_size = usize::min(
                ptr_header.payload_size,
                result_header.payload_size,
            );
            ptr::copy_nonoverlapping(ptr, result, copy_size);
            wfree(ptr);
            result
        }
    }

    /// Bumps the refcount for `ptr`.
    ///
    /// ## Safety
    ///
    /// Callers must ensure that `b` is a live allocation from [`wmalloc`] or [`wrealloc`].
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn wretain(ptr: *mut c_void) -> *mut c_void {
        if ptr.is_null() {
            return ptr::null_mut();
        }
        unsafe {
            let header = Header::for_alloc_bytes(ptr.cast::<u8>());
            (*header).refcount += 1;
        }
        ptr
    }

    /// Decrements the refcount for `ptr`. If this brings the refcount to 0,
    /// frees `ptr`.
    ///
    /// ## Safety
    ///
    /// Callers must ensure that `ptr` is a live allocation from [`wmalloc`] or [`wrealloc`].
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn wrelease(ptr: *mut c_void) {
        let ptr = ptr.cast::<u8>();
        if ptr.is_null() {
            return;
        }
        let header = unsafe { &mut *Header::for_alloc_bytes(ptr) };
        match header.refcount {
            0 | 1 => unsafe { free_bytes(ptr) },
            _ => header.refcount -= 1,
        }
    }
}

#[cfg(test)]
mod test {
    use super::{alloc_bytes, alloc_string, ffi::{wfree, wmalloc, wrealloc}, free_bytes, Header};

    use std::{ffi::CStr, mem, os::raw::c_void, ptr, slice};

    #[test]
    fn recover_header() {
        unsafe {
            let x = alloc_bytes(mem::size_of::<i64>());
            let header = Header::for_alloc_bytes(x);
            assert_eq!(header.cast::<u8>().add(mem::size_of::<Header>()), x);
            // This may be allocator-dependent, but it's a reasonable sanity check for now.
            assert!((*header).ptr.as_ptr() <= header.cast::<u8>());
        }
    }
    
    #[test]
    fn alloc_zero_returns_null() {
        assert!(alloc_bytes(0).is_null());
    }

    #[test]
    fn free_null() {
        unsafe { free_bytes(ptr::null_mut()); }
    }

    #[test]
    fn realloc_null() {
        unsafe { assert!(wrealloc(ptr::null_mut(), 0).is_null()); }
    }

    #[test]
    fn alloc_free_nonzero() {
        let x = alloc_bytes(mem::size_of::<i64>()).cast::<i64>();
        assert!(!x.is_null());
        unsafe { *x = 42; }
        assert_eq!(unsafe { *x }, 42);
        unsafe { free_bytes(x.cast::<u8>()); }
    }

    #[test]
    fn multiple_allocs() {
        unsafe {
            let x = wmalloc(mem::size_of::<i64>());
            *x.cast::<i64>() = 30;
            let y = wmalloc(mem::size_of::<i32>());
            *y.cast::<i32>() = 5;
            let z = wmalloc(48);
            *z.cast::<f32>() = 1.0;
            wfree(x);
            wfree(y);
            wfree(z);
        }
    }

    #[test]
    fn realloc_nonzero() {
        let x = alloc_bytes(mem::size_of::<i64>()).cast::<c_void>();
        assert!(!x.is_null());
        let y = unsafe { wrealloc(x, mem::size_of::<i32>()).cast::<i32>() };
        assert!(!y.is_null());
        unsafe { *y = 17; }
        assert_eq!(unsafe { *y }, 17);
        unsafe { free_bytes(y.cast::<u8>()); }
    }

    #[test]
    fn realloc_retains_data() {
        let x: *mut u8 = unsafe { wmalloc(10).cast() };
        unsafe {
            let xs = slice::from_raw_parts_mut(&mut *x, 10);
            // We know that xs should be zeroed.
            assert_eq!(xs, &[0u8; 10]);
            for i in 0u8..10 {
                xs[i as usize] = i;
            }
            assert_eq!(xs, (0..10).collect::<Vec::<u8>>());
        }
    }

    #[test]
    fn alloc_free_string() {
        let s = alloc_string(c"hello");
        assert!(!s.is_null());
        assert_eq!(unsafe { CStr::from_ptr(s) }, c"hello");
        unsafe { free_bytes(s.cast::<u8>()); }
    }
}