irid-std 0.3.1

use core::{alloc::Layout, marker::PhantomData, mem::swap, ops::{Deref, DerefMut, Index, IndexMut}, ptr::{copy, drop_in_place, null_mut, slice_from_raw_parts, slice_from_raw_parts_mut}};

use irid_syscall::{io::FileDescriptor, memory::{MemoryFlags, memory_map, memory_unmap}};

pub struct PageVec<T> {
    allocation: *mut u8,
    allocation_length: usize,
    length: usize,
    _data: PhantomData<T>
}

unsafe impl <T: Send> Send for PageVec<T> {
}

unsafe impl <T: Sync> Sync for PageVec<T> {
}

impl <T> PageVec<T> {
    pub const fn new() -> Self {
        Self {
            allocation: null_mut(),
            allocation_length: 0,
            length: 0,
            _data: PhantomData
        }
    }

    pub fn clear(&mut self) {
        let layout = Layout::new::<T>().pad_to_align();
        let size = layout.size();
        
        for i in 0..self.length {
            let ptr = unsafe {self.allocation.byte_offset((i * size) as isize)} as *mut T;
            unsafe { drop_in_place(ptr) };
        }

        unsafe {memory_unmap(self.allocation, self.allocation_length)}.expect("failed to free page vector");

        self.allocation = null_mut();
        self.allocation_length = 0;
        self.length = 0;
    }

    pub fn push(&mut self, value: T) {
        let layout = Layout::new::<T>().pad_to_align();
        
        let index = layout.size() * self.length;

        if index + layout.size() > self.allocation_length {
            self.grow();
        }

        let ptr = unsafe {self.allocation.byte_offset(index as isize)} as *mut T;
        unsafe {ptr.write(value)};

        self.length += 1;
    }
    
    pub fn pop(&mut self) -> Option<T> {
        if self.length == 0 {
            return None;
        }

        let layout = Layout::new::<T>().pad_to_align();
        
        let index = layout.size() * (self.length - 1);

        let ptr = unsafe {self.allocation.byte_offset(index as isize)} as *mut T;

        self.length -= 1;

        Some(unsafe { ptr.read() })
    }
    
    pub fn remove(&mut self, index: usize) -> Option<T> {
        if self.length == 0 || index >= self.length {
            return None;
        }

        let layout = Layout::new::<T>().pad_to_align();
        
        let index = layout.size() * index;

        let ptr = unsafe {self.allocation.byte_offset(index as isize)} as *mut T;
        self.length -= 1;
          

        unsafe {
            let value = ptr.read();
           
            let copy_length = self.allocation_length - index - layout.size();
            copy(ptr.byte_offset(layout.size() as isize) as *const u8, ptr as *mut u8, copy_length);

            Some(value)
        }
    }
    
    pub fn shrink_to_fit(&mut self) {
        let layout = Layout::new::<T>().pad_to_align();
        let mut target_length = self.length * layout.size();

        unsafe {
            let mut allocation2 = memory_map(null_mut(), target_length, MemoryFlags::READABLE | MemoryFlags::WRITEABLE | MemoryFlags::ZERO, FileDescriptor::MAX, 0)
                .expect("failed to shrink page vec");

            let old: &[u8] = slice_from_raw_parts(self.allocation, target_length).as_ref_unchecked();
            let new: &mut [u8] = slice_from_raw_parts_mut(allocation2, target_length).as_mut_unchecked();

            new.copy_from_slice(old);

            swap(&mut self.allocation, &mut allocation2);
            swap(&mut self.allocation_length, &mut target_length);

            memory_unmap(allocation2, target_length).expect("failed to free old page vec memory")
        }
    }

    fn grow(&mut self) {
        unsafe {
            let mut new_length = if self.allocation_length == 0 {
                let layout = Layout::new::<T>().pad_to_align();

                layout.size()
            } else {
                self.allocation_length * 2
            };

            let mut allocation2 = memory_map(null_mut(), new_length, MemoryFlags::WRITEABLE | MemoryFlags::READABLE | MemoryFlags::ZERO, FileDescriptor::MAX, 0)
                .expect("failed to allocate space for page vector");

            if !self.allocation.is_null() {
                let old: &[u8] = slice_from_raw_parts(self.allocation, self.allocation_length).as_ref_unchecked();
                let new: &mut [u8] = slice_from_raw_parts_mut(allocation2, self.allocation_length).as_mut_unchecked();

                new.copy_from_slice(old);
                
                swap(&mut self.allocation, &mut allocation2);
                swap(&mut self.allocation_length, &mut new_length);

                memory_unmap(allocation2, new_length).expect("failed to free old page vec memory");
            } else {
                swap(&mut self.allocation, &mut allocation2);
            }
        }
    }
}

impl <T> Index<usize> for PageVec<T> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        if index < self.length {
            let layout = Layout::new::<T>().pad_to_align();
            let offset = index * layout.size();

            let ptr = unsafe {self.allocation.byte_offset(offset as isize)} as *const T;

            unsafe { ptr.as_ref_unchecked() } 
        } else {
            panic!("Index {index} out of bounds for page vec of length {}", self.length);
        }
    }
}

impl <T> IndexMut<usize> for PageVec<T> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        if index < self.length {
            let layout = Layout::new::<T>().pad_to_align();
            let offset = index * layout.size();

            let ptr = unsafe {self.allocation.byte_offset(offset as isize)} as *mut T;

            unsafe { ptr.as_mut_unchecked() } 
        } else {
            panic!("Index {index} out of bounds for page vec of length {}", self.length);
        }
    }
}

impl <T> Drop for PageVec<T> {
    fn drop(&mut self) {
        let layout = Layout::new::<T>().pad_to_align();
        
        for i in 0..self.length {
            let ptr = unsafe {self.allocation.byte_offset((i * layout.size()) as isize)} as *mut T;
            unsafe { drop_in_place(ptr) };
        }

        unsafe {memory_unmap(self.allocation, self.allocation_length)}.expect("failed to free page vector")
    }
}

impl <T> Deref for PageVec<T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        unsafe { slice_from_raw_parts(self.allocation as *const T, self.length).as_ref_unchecked() }
    }
}

impl <T> DerefMut for PageVec<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { slice_from_raw_parts_mut(self.allocation as *mut T, self.length).as_mut_unchecked() }
    }
}