memkit 0.1.0-beta.1

//! Deferred memory reclamation.

use std::alloc::Layout;
use crossbeam_queue::SegQueue;

/// An item in the deferred queue.
struct DeferredItem {
    ptr: *mut u8,
    layout: Layout,
    drop_fn: Option<unsafe fn(*mut u8)>,
}

unsafe impl Send for DeferredItem {}
unsafe impl Sync for DeferredItem {}

/// A thread-safe queue for deferred memory reclamation.
///
/// This queue allows objects to be marked for deletion by one thread and
/// reclaimed later (e.g., at the end of a frame) by another thread or a
/// system-level reclamation pass.
pub struct MkDeferredQueue {
    queue: SegQueue<DeferredItem>,
}

impl MkDeferredQueue {
    /// Create a new deferred queue.
    pub fn new() -> Self {
        Self {
            queue: SegQueue::new(),
        }
    }

    /// Push an allocation to the deferred queue.
    ///
    /// # Safety
    ///
    /// The pointer must be valid and allocated with the given layout.
    /// If `drop_fn` is provided, it will be called before the memory is freed.
    pub unsafe fn push(&self, ptr: *mut u8, layout: Layout, drop_fn: Option<unsafe fn(*mut u8)>) {
        self.queue.push(DeferredItem { ptr, layout, drop_fn });
    }

    /// Reclaim all items currently in the queue.
    ///
    /// Calls the provided `free_fn` for each item after calling its `drop_fn` if present.
    /// Returns the number of items reclaimed.
    pub fn reclaim<F>(&self, mut free_fn: F) -> usize
    where
        F: FnMut(*mut u8, Layout),
    {
        let mut count = 0;
        while let Some(item) = self.queue.pop() {
            unsafe {
                if let Some(drop_fn) = item.drop_fn {
                    drop_fn(item.ptr);
                }
                free_fn(item.ptr, item.layout);
            }
            count += 1;
        }
        count
    }

    /// Check if the queue is empty.
    pub fn is_empty(&self) -> bool {
        self.queue.is_empty()
    }

    /// Get the approximate number of items in the queue.
    pub fn len(&self) -> usize {
        self.queue.len()
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::Arc;

    #[test]
    fn test_deferred_queue_basic() {
        let queue = MkDeferredQueue::new();
        let dropped = Arc::new(AtomicBool::new(false));
        let dropped_clone = Arc::clone(&dropped);

        unsafe {
            let ptr = std::alloc::alloc(Layout::new::<u32>());
            
            // We need a wrapper to capture the Arc in the drop_fn
            // but drop_fn is a plain unsafe fn(*mut u8).
            // This is a test, so we'll use a global or a bit of hackery.
            // For simplicity, let's just test that the free_fn is called correctly.
            
            queue.push(ptr, Layout::new::<u32>(), None);
        }

        let mut freed = false;
        queue.reclaim(|ptr, layout| {
            unsafe { std::alloc::dealloc(ptr, layout) };
            freed = true;
        });

        assert!(freed);
        assert!(queue.is_empty());
    }
}