orthotope 0.1.4

use core::ptr::NonNull;
use std::collections::HashMap;
use std::collections::hash_map::Entry;

use parking_lot::Mutex;

use crate::error::FreeError;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
/// Live-registry entry for one large allocation.
pub(crate) struct LargeAllocationRecord {
    pub(crate) block_addr: usize,
    pub(crate) block_size: usize,
    pub(crate) requested_size: usize,
    pub(crate) usable_size: usize,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct FreeLargeBlock {
    pub(crate) block_addr: usize,
    pub(crate) block_size: usize,
}

impl FreeLargeBlock {
    #[must_use]
    pub(crate) const fn usable_size(self) -> usize {
        self.block_size - crate::header::HEADER_SIZE
    }

    #[must_use]
    pub(crate) fn block_start(self) -> NonNull<u8> {
        let ptr = self.block_addr as *mut u8;
        NonNull::new(ptr).unwrap_or_else(|| panic!("stored free large block address was null"))
    }
}

#[derive(Default)]
struct LargeObjectState {
    live: HashMap<usize, LargeAllocationRecord>,
    free: Vec<FreeLargeBlock>,
}

/// Tracks live allocations larger than the largest small class.
pub(crate) struct LargeObjectAllocator {
    state: Mutex<LargeObjectState>,
}

impl LargeObjectAllocator {
    #[must_use]
    /// Creates an empty large-allocation registry.
    pub(crate) fn new() -> Self {
        Self {
            state: Mutex::new(LargeObjectState::default()),
        }
    }

    #[must_use]
    pub(crate) fn take_reusable_block(&self, minimum_block_size: usize) -> Option<FreeLargeBlock> {
        let mut state = self.state.lock();
        let best_fit_index = state
            .free
            .iter()
            .enumerate()
            .filter(|(_, block)| block.block_size >= minimum_block_size)
            .min_by_key(|(_, block)| block.block_size)
            .map(|(index, _)| index)?;

        Some(state.free.swap_remove(best_fit_index))
    }

    /// Records one newly allocated large block as live.
    pub(crate) fn record_live_allocation(
        &self,
        user_ptr: NonNull<u8>,
        block_start: NonNull<u8>,
        block_size: usize,
        requested_size: usize,
        usable_size: usize,
    ) {
        debug_assert!(requested_size > 0);
        debug_assert!(usable_size >= requested_size);

        let record = LargeAllocationRecord {
            block_addr: block_start.as_ptr().addr(),
            block_size,
            requested_size,
            usable_size,
        };

        match self.state.lock().live.entry(user_ptr.as_ptr().addr()) {
            Entry::Vacant(entry) => {
                entry.insert(record);
            }
            Entry::Occupied(_) => {
                panic!("large allocation registered twice for the same user pointer");
            }
        }
    }

    /// Validates and releases the live-tracking record for a large allocation.
    ///
    /// # Errors
    ///
    /// Returns [`FreeError::AlreadyFreedOrUnknownLarge`] if `user_ptr` is not a
    /// currently tracked large allocation, or [`FreeError::CorruptHeader`] when the
    /// decoded header disagrees with the live registry entry.
    pub(crate) fn validate_and_release_live_allocation(
        &self,
        user_ptr: NonNull<u8>,
        requested_size: usize,
        usable_size: usize,
    ) -> Result<(), FreeError> {
        let mut state = self.state.lock();
        let Some(record) = state.live.get(&user_ptr.as_ptr().addr()).copied() else {
            return Err(FreeError::AlreadyFreedOrUnknownLarge);
        };

        if record.requested_size != requested_size || record.usable_size != usable_size {
            return Err(FreeError::CorruptHeader);
        }

        let removed = state.live.remove(&user_ptr.as_ptr().addr());
        state.free.push(FreeLargeBlock {
            block_addr: record.block_addr,
            block_size: record.block_size,
        });
        drop(state);
        debug_assert!(removed.is_some(), "validated live record must still exist");
        Ok(())
    }

    #[cfg(test)]
    #[must_use]
    pub(crate) fn live_len(&self) -> usize {
        self.state.lock().live.len()
    }

    #[cfg(test)]
    #[must_use]
    pub(crate) fn free_len(&self) -> usize {
        self.state.lock().free.len()
    }
}

#[cfg(test)]
mod tests {
    use super::{FreeLargeBlock, LargeAllocationRecord, LargeObjectAllocator};
    use crate::error::FreeError;
    use crate::header::{HEADER_SIZE, user_ptr_from_block_start};
    use core::alloc::Layout;
    use core::ptr::NonNull;
    use std::alloc::{alloc, dealloc};

    struct TestBlock {
        ptr: NonNull<u8>,
        size: usize,
    }

    impl TestBlock {
        fn new(size: usize) -> Self {
            let layout = match Layout::from_size_align(size, 64) {
                Ok(layout) => layout,
                Err(error) => panic!("expected valid test layout: {error}"),
            };
            // SAFETY: `layout` has non-zero size and valid alignment for the test allocation.
            let ptr = unsafe { alloc(layout) };
            let Some(ptr) = NonNull::new(ptr) else {
                panic!("expected heap allocation for test block to succeed");
            };

            Self { ptr, size }
        }

        fn block_start(&self) -> NonNull<u8> {
            self.ptr
        }
    }

    impl Drop for TestBlock {
        fn drop(&mut self) {
            let layout = match Layout::from_size_align(self.size, 64) {
                Ok(layout) => layout,
                Err(error) => panic!("expected valid test layout during drop: {error}"),
            };

            // SAFETY: `self.ptr` was allocated with `alloc(layout)` in `TestBlock::new`
            // and has not been deallocated yet.
            unsafe {
                dealloc(self.ptr.as_ptr(), layout);
            }
        }
    }

    fn live_record(
        allocator: &LargeObjectAllocator,
        block: &TestBlock,
        requested_size: usize,
        usable_size: usize,
    ) -> (NonNull<u8>, LargeAllocationRecord) {
        let block_start = block.block_start();
        let user_ptr = user_ptr_from_block_start(block_start);
        let record = LargeAllocationRecord {
            block_addr: block_start.as_ptr().addr(),
            block_size: block.size,
            requested_size,
            usable_size,
        };

        allocator.record_live_allocation(
            user_ptr,
            block_start,
            record.block_size,
            record.requested_size,
            record.usable_size,
        );

        (user_ptr, record)
    }

    #[test]
    fn matching_release_succeeds_and_removes_live_record() {
        let allocator = LargeObjectAllocator::new();
        let block_size = HEADER_SIZE + 16_777_217;
        let block = TestBlock::new(block_size);
        let usable_size = block_size - HEADER_SIZE;
        let (user_ptr, expected) = live_record(&allocator, &block, 16_777_217, usable_size);

        let released = allocator.validate_and_release_live_allocation(
            user_ptr,
            expected.requested_size,
            expected.usable_size,
        );

        assert_eq!(released, Ok(()));
        assert_eq!(allocator.live_len(), 0);
        assert_eq!(allocator.free_len(), 1);
    }

    #[test]
    fn releasing_unknown_pointer_fails() {
        let allocator = LargeObjectAllocator::new();
        let block = TestBlock::new(HEADER_SIZE + 16_777_217);
        let user_ptr = user_ptr_from_block_start(block.block_start());

        let result =
            allocator.validate_and_release_live_allocation(user_ptr, 16_777_217, 16_777_217);

        assert_eq!(result, Err(FreeError::AlreadyFreedOrUnknownLarge));
    }

    #[test]
    fn releasing_same_pointer_twice_detects_duplicate_free() {
        let allocator = LargeObjectAllocator::new();
        let block_size = HEADER_SIZE + 16_777_280;
        let block = TestBlock::new(block_size);
        let usable_size = block_size - HEADER_SIZE;
        let (user_ptr, expected) = live_record(&allocator, &block, 16_777_217, usable_size);

        let first = allocator.validate_and_release_live_allocation(
            user_ptr,
            expected.requested_size,
            expected.usable_size,
        );
        let second = allocator.validate_and_release_live_allocation(
            user_ptr,
            expected.requested_size,
            expected.usable_size,
        );

        assert_eq!(first, Ok(()));
        assert_eq!(second, Err(FreeError::AlreadyFreedOrUnknownLarge));
    }

    #[test]
    fn mismatched_header_sizes_fail_without_consuming_live_record() {
        let allocator = LargeObjectAllocator::new();
        let block_size = HEADER_SIZE + 16_777_280;
        let block = TestBlock::new(block_size);
        let usable_size = block_size - HEADER_SIZE;
        let (user_ptr, expected) = live_record(&allocator, &block, 16_777_217, usable_size);

        let mismatch = allocator.validate_and_release_live_allocation(
            user_ptr,
            expected.requested_size + 1,
            expected.usable_size,
        );
        assert_eq!(mismatch, Err(FreeError::CorruptHeader));
        assert_eq!(allocator.live_len(), 1);

        let release = allocator.validate_and_release_live_allocation(
            user_ptr,
            expected.requested_size,
            expected.usable_size,
        );
        assert_eq!(release, Ok(()));
        assert_eq!(allocator.live_len(), 0);
        assert_eq!(allocator.free_len(), 1);
    }

    #[test]
    fn records_are_isolated_per_pointer() {
        let allocator = LargeObjectAllocator::new();
        let first_block_size = HEADER_SIZE + 16_777_344;
        let second_block_size = HEADER_SIZE + 20_000_000;
        let first_block = TestBlock::new(first_block_size);
        let second_block = TestBlock::new(second_block_size);
        let first_usable_size = first_block_size - HEADER_SIZE;
        let second_usable_size = second_block_size - HEADER_SIZE;
        let (first_ptr, first_expected) =
            live_record(&allocator, &first_block, 16_777_280, first_usable_size);
        let (second_ptr, second_expected) =
            live_record(&allocator, &second_block, 20_000_000, second_usable_size);

        let first = allocator.validate_and_release_live_allocation(
            first_ptr,
            first_expected.requested_size,
            first_expected.usable_size,
        );

        assert_eq!(first, Ok(()));
        assert_eq!(allocator.live_len(), 1);

        let second = allocator.validate_and_release_live_allocation(
            second_ptr,
            second_expected.requested_size,
            second_expected.usable_size,
        );
        assert_eq!(second, Ok(()));
        assert_eq!(allocator.live_len(), 0);
        assert_eq!(allocator.free_len(), 2);
    }

    #[test]
    fn reusable_block_prefers_smallest_sufficient_free_span() {
        let allocator = LargeObjectAllocator::new();
        let first = TestBlock::new(HEADER_SIZE + 16_777_280);
        let second = TestBlock::new(HEADER_SIZE + 20_000_000);
        let third = TestBlock::new(HEADER_SIZE + 18_000_000);

        let first_block = FreeLargeBlock {
            block_addr: first.block_start().as_ptr().addr(),
            block_size: first.size,
        };
        let second_block = FreeLargeBlock {
            block_addr: second.block_start().as_ptr().addr(),
            block_size: second.size,
        };
        let third_block = FreeLargeBlock {
            block_addr: third.block_start().as_ptr().addr(),
            block_size: third.size,
        };

        let user_ptr = user_ptr_from_block_start(first.block_start());
        allocator.record_live_allocation(
            user_ptr,
            first.block_start(),
            first.size,
            16_777_217,
            first.size - HEADER_SIZE,
        );
        let _ = allocator.validate_and_release_live_allocation(
            user_ptr,
            16_777_217,
            first.size - HEADER_SIZE,
        );

        let user_ptr = user_ptr_from_block_start(second.block_start());
        allocator.record_live_allocation(
            user_ptr,
            second.block_start(),
            second.size,
            19_999_937,
            second.size - HEADER_SIZE,
        );
        let _ = allocator.validate_and_release_live_allocation(
            user_ptr,
            19_999_937,
            second.size - HEADER_SIZE,
        );

        let user_ptr = user_ptr_from_block_start(third.block_start());
        allocator.record_live_allocation(
            user_ptr,
            third.block_start(),
            third.size,
            17_999_937,
            third.size - HEADER_SIZE,
        );
        let _ = allocator.validate_and_release_live_allocation(
            user_ptr,
            17_999_937,
            third.size - HEADER_SIZE,
        );

        let reused = allocator
            .take_reusable_block(HEADER_SIZE + 17_000_000)
            .unwrap_or_else(|| panic!("expected a reusable large block"));

        assert_eq!(reused, third_block);
        assert_ne!(reused, first_block);
        assert_ne!(reused, second_block);
        assert!(reused.usable_size() >= 17_000_000);
    }

    #[test]
    #[should_panic(expected = "large allocation registered twice for the same user pointer")]
    fn duplicate_registration_panics_in_all_builds() {
        let allocator = LargeObjectAllocator::new();
        let block_size = HEADER_SIZE + 16_777_280;
        let block = TestBlock::new(block_size);
        let usable_size = block_size - HEADER_SIZE;
        let (user_ptr, expected) = live_record(&allocator, &block, 16_777_217, usable_size);

        allocator.record_live_allocation(
            user_ptr,
            NonNull::new(expected.block_addr as *mut u8)
                .unwrap_or_else(|| panic!("expected live record block address to remain non-null")),
            expected.block_size,
            expected.requested_size,
            expected.usable_size,
        );
    }
}