// SPDX-FileCopyrightText: 2023 Rivos Inc.
//
// SPDX-License-Identifier: Apache-2.0

use core::alloc::{AllocError, Allocator, Layout};
use core::ptr::NonNull;
use core::slice;
use riscv_pages::{InternalClean, InternalDirty, PageAddr, PageSize, RawAddr, SequentialPages};
use sync::Mutex;

struct HypAllocInner {
    mem: NonNull<[u8]>,
    end: usize,
    page_size: PageSize,
    sealed: bool,
}

impl HypAllocInner {
    /// Returns the length of the underlying heap storage.
    fn capacity(&self) -> usize {
        // SAFETY: We construct a valid reference only for the unsafe block, returning a copy of
        // the slice length.
        unsafe { self.mem.as_ref().len() }
    }

    fn allocate(&mut self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        if self.sealed {
            return Err(AllocError);
        }

        let align = layout.align();
        let size = layout.size();
        let capacity = self.capacity();

        let mem: *mut u8 = self.mem.as_ptr().cast();
        // Won't wrap since self.end is known not to overflow by construction.
        let end_ptr = mem.wrapping_add(self.end);

        let align_offset = end_ptr.align_offset(align);
        let block_start_offset = self.end.checked_add(align_offset).ok_or(AllocError)?;
        if block_start_offset > capacity {
            return Err(AllocError);
        }

        let block_end_offset = block_start_offset.checked_add(size).ok_or(AllocError)?;
        if block_end_offset > capacity {
            return Err(AllocError);
        }
        self.end = block_end_offset;

        let block_start_ptr = mem.wrapping_add(block_start_offset);
        // SAFETY: No other mutable reference to this slice exists.
        let block_slice = unsafe { slice::from_raw_parts_mut(block_start_ptr, size) };

        // SAFETY: block_slice is a non-null pointer to a valid slice.
        Ok(unsafe { NonNull::new_unchecked(block_slice) })
    }

    fn seal(&mut self) {
        self.sealed = true;
    }
}

/// A simple thread-safe bump-pointer allocator backed by a fixed-length contiguous range of Pages.
/// Implements the `Allocator` trait so that it may be used with standard containers supporting the
/// allocator API.
pub struct HypAlloc {
    inner: Mutex<HypAllocInner>,
}

impl HypAlloc {
    /// Creates an allocator from a range of pages which will be used as storage for the allocator.
    pub fn from_pages(pages: SequentialPages<InternalClean>) -> Self {
        let inner = HypAllocInner {
            mem: NonNull::slice_from_raw_parts(
                NonNull::new(pages.base().bits() as *mut u8).unwrap(),
                pages.length_bytes().try_into().unwrap(),
            ),
            end: 0,
            page_size: pages.page_size(),
            sealed: false,
        };
        HypAlloc {
            inner: Mutex::new(inner),
        }
    }

    /// Destroys the allocator and returns the pages which were used as storage.
    pub fn to_pages(self) -> SequentialPages<InternalDirty> {
        let inner = self.inner.lock();
        let base = inner.mem.as_mut_ptr() as u64;
        let num_pages = (inner.mem.len() as u64) / inner.page_size as u64;
        unsafe {
            // Safe since this allocator must own this range of pages.
            SequentialPages::from_mem_range(
                PageAddr::with_alignment(RawAddr::supervisor(base), inner.page_size).unwrap(),
                inner.page_size,
                num_pages,
            )
            .unwrap()
        }
    }

    /// Seals the allocator, preventing any further allocations.
    pub fn seal(&self) {
        self.inner.lock().seal();
    }
}

// Safety: HypAlloc is synchronized internally via a mutex, so it is safe to share a reference to it
// between threads.
unsafe impl Sync for HypAlloc {}
// Safety: HypAlloc uniquely owns the memory it manages, so we can safely transfer it between threads.
unsafe impl Send for HypAlloc {}

unsafe impl Allocator for HypAlloc {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        let align = layout.align();
        let size = layout.size();
        if size == 0 {
            // SAFETY: align is always nonzero.
            let aligned_ptr = unsafe { NonNull::new_unchecked(align as *mut u8) };

            // SAFETY: a zero-length slice is safe to construct with a dangling pointer.
            let empty_slice = unsafe { slice::from_raw_parts_mut(aligned_ptr.as_ptr(), 0) };

            // SAFETY: empty_slice is a dangling, aligned non-null pointer.
            return Ok(unsafe { NonNull::new_unchecked(empty_slice) });
        }

        let mut inner = self.inner.lock();
        inner.allocate(layout)
    }

    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {}
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::boxed::Box;
    use alloc::vec::Vec;

    fn stub_heap() -> HypAlloc {
        const ONE_MEG: usize = 1024 * 1024;
        const MEM_ALIGN: usize = 2 * ONE_MEG;
        const MEM_SIZE: usize = 256 * ONE_MEG;
        let backing_mem = vec![0u8; MEM_SIZE + MEM_ALIGN];
        let aligned_pointer = unsafe {
            // Not safe - just a test
            backing_mem
                .as_ptr()
                .add(backing_mem.as_ptr().align_offset(MEM_ALIGN))
        };
        let start_page = PageAddr::new(RawAddr::supervisor(aligned_pointer as u64)).unwrap();
        let num_pages = (MEM_SIZE as u64) / PageSize::Size4k as u64;
        let pages = unsafe {
            // Not safe - just a test
            SequentialPages::from_mem_range(start_page, PageSize::Size4k, num_pages).unwrap()
        };
        // Leak the backing ram so it doesn't get freed
        std::mem::forget(backing_mem);
        HypAlloc::from_pages(pages)
    }

    #[test]
    fn basic_alloc() {
        let alloc = stub_heap();
        {
            let mut vec = Vec::new_in(&alloc);
            vec.push(1);
            vec.push(5);
            vec.push(10);
            assert_eq!(vec.len(), 3);

            let five = Box::new_in(5, &alloc);
            assert_eq!(*five, 5);
        }
        let _ = alloc.to_pages();
    }

    #[test]
    fn sealing() {
        let alloc = stub_heap();
        let mut vec_a = Vec::new_in(&alloc);
        assert!(vec_a.try_reserve(10).is_ok());
        vec_a.push(42);
        alloc.seal();
        let mut vec_b: Vec<u32, &HypAlloc> = Vec::new_in(&alloc);
        assert!(vec_b.try_reserve(10).is_err());
    }
}
