ark-module 0.17.0-pre.18

use std::alloc::GlobalAlloc;
use std::alloc::Layout;
use std::alloc::System;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use ark_api_ffi::entrypoints::core::AllocStats;

/// Statistics memory tracker, enables host to get stats of allocations in the module. Hooks into `[global_allocator]`.
///
/// It only tracks cumulative statistics to keep updating them simple. Other statistics
/// can be derived from these by subtracting allocated and free'd, and saving previous stats,
/// for example.
pub struct StatsAllocator {
    cumul_alloc_count: AtomicUsize,
    cumul_alloc_size: AtomicUsize,
    cumul_free_count: AtomicUsize,
    cumul_free_size: AtomicUsize,

    // In the format for the host to read.
    stats: AllocStats,
}

#[global_allocator]
pub static mut GLOBAL_ALLOCATOR: StatsAllocator = StatsAllocator {
    cumul_alloc_count: AtomicUsize::new(0),
    cumul_alloc_size: AtomicUsize::new(0),
    cumul_free_count: AtomicUsize::new(0),
    cumul_free_size: AtomicUsize::new(0),
    // C-readable version of the above.
    stats: AllocStats {
        cumul_alloc_count: 0,
        cumul_alloc_size: 0,
        cumul_free_count: 0,
        cumul_free_size: 0,
    },
};

#[allow(dead_code)]
impl StatsAllocator {
    pub fn copy_stats() {
        // SAFETY: Sound to access static as we do not use threads in Wasm
        unsafe {
            GLOBAL_ALLOCATOR.stats = AllocStats {
                cumul_alloc_count: GLOBAL_ALLOCATOR.cumul_alloc_count.load(Ordering::SeqCst) as u64,
                cumul_alloc_size: GLOBAL_ALLOCATOR.cumul_alloc_size.load(Ordering::SeqCst) as u64,
                cumul_free_count: GLOBAL_ALLOCATOR.cumul_free_count.load(Ordering::SeqCst) as u64,
                cumul_free_size: GLOBAL_ALLOCATOR.cumul_free_size.load(Ordering::SeqCst) as u64,
            }
        }
    }
}

// SAFETY: allocator trait is unsafe, our implementation is not
unsafe impl GlobalAlloc for StatsAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let _ = self.cumul_alloc_count.fetch_add(1, Ordering::SeqCst);
        let _ = self
            .cumul_alloc_size
            .fetch_add(layout.size(), Ordering::SeqCst);
        // SAFETY: Nested allocator call, safety contract is same as on the unsafe function itself
        unsafe { System.alloc(layout) }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let _ = self.cumul_free_count.fetch_add(1, Ordering::SeqCst);
        let _ = self
            .cumul_free_size
            .fetch_add(layout.size(), Ordering::SeqCst);
        // SAFETY: Nested allocator call, safety contract is same as on the unsafe function itself
        unsafe {
            System.dealloc(ptr, layout);
        }
    }
}

/// Allocate a raw buffer in the module
///
/// This is called from the host to be able to allocate and write memory
/// in a module that it then can pass in as a pointer when calling other
/// entrypoint functions in the module
///
/// # Safety
///
/// Explicit memory allocation is inherently unsafe, this should only be called by
/// the host inside specific low-level FFI implementations
#[no_mangle]
pub unsafe fn ark_alloc(size: usize, alignment: usize) -> *mut u8 {
    let layout = std::alloc::Layout::from_size_align(size, alignment).unwrap();
    // SAFETY: Nested allocator call, safety contract is same as on the unsafe function itself
    unsafe { std::alloc::alloc(layout) }
}

/// Free a raw buffer in the module
/// This is intended to be used on memory allocated with `ark_alloc`
///
/// # Safety
///
/// Explicit memory allocation is inherently unsafe, this should only be called by
/// the host inside specific low-level FFI implementations
#[no_mangle]
pub unsafe fn ark_free(ptr: *mut u8, size: usize, alignment: usize) {
    let layout = std::alloc::Layout::from_size_align(size, alignment).unwrap();
    // SAFETY: Nested allocator call, safety contract is same as on the unsafe function itself
    unsafe {
        std::alloc::dealloc(ptr, layout);
    }
}

/// Returns a pointer to host-readable allocator statistics.
#[no_mangle]
#[doc(hidden)]
pub fn ark_get_alloc_stats() -> u32 {
    StatsAllocator::copy_stats();
    // SAFETY: Access of global static can race, but our Wasm code is single-threaded so not a problem
    unsafe { &GLOBAL_ALLOCATOR.stats as *const AllocStats }.cast::<u8>() as usize as u32
}