unreal-sdk/src/tarray.rs

use std::{
    ops::{Index, IndexMut},
    ptr::NonNull,
    slice::SliceIndex,
};

use once_cell::sync::OnceCell;

static ALLOCATOR: OnceCell<Box<dyn std::alloc::Allocator + Send + Sync + 'static>> =
    OnceCell::new();

pub fn set_allocator<A>(ally: A)
where
    A: std::alloc::Allocator + Send + Sync + 'static,
{
    _ = ALLOCATOR.set(Box::new(ally));
}

#[repr(C)]
#[derive(Debug)]
pub struct TArray<T> {
    data: Option<NonNull<T>>,
    count: u32,
    max: u32,
}

unsafe impl<T> Send for TArray<T> where T: Send {}
unsafe impl<T> Sync for TArray<T> where T: Sync {}

pub type FString = TArray<u16>;

impl ToString for FString {
    fn to_string(&self) -> String {
        widestring::U16Str::from_slice(&self).to_string_lossy()
    }
}

impl<T> Drop for TArray<T> {
    fn drop(&mut self) {
        if let Some(ptr) = self.data {
            unsafe {
                ALLOCATOR
                    .get()
                    .expect("allocator")
                    .deallocate(ptr.cast(), std::alloc::Layout::for_value(ptr.as_ref()));
            }
        }
    }
}

impl<T> TArray<T> {
    pub fn new() -> Self {
        Self {
            data: None,
            count: 0,
            max: 0,
        }
    }

    pub unsafe fn set_len(&mut self, size: usize) {
        self.count = (size as u32).min(self.max);
    }

    pub fn reserve(&mut self, size: usize) {
        self.ensure_space(size);
    }

    fn ensure_space(&mut self, size: usize) {
        if self.data.is_none() || size >= self.max as usize {
            self.grow(size.max(self.max as usize * 2));
        }
    }

    fn grow(&mut self, new_size: usize) {
        if new_size <= self.max as usize {
            return;
        }

        if self.max == 0 || self.data.is_none() {
            if let Some(alloc) = ALLOCATOR.get() {
                match alloc.allocate(std::alloc::Layout::array::<T>(32).unwrap()) {
                    Ok(ptr) => {
                        self.data = Some(ptr.cast());
                        self.max = ptr.len() as u32;
                    }
                    Err(e) => {
                        log::error!("failed to allocate with UE4 allocator: {e}.");
                    }
                }
            }
        } else if let Some(alloc) = ALLOCATOR.get() {
            unsafe {
                match alloc.grow_zeroed(
                    self.data.unwrap().cast(),
                    std::alloc::Layout::array::<T>(self.max as usize).unwrap(),
                    std::alloc::Layout::array::<T>(new_size).unwrap(),
                ) {
                    Ok(ptr) => {
                        self.data = Some(ptr.cast());
                        self.max = ptr.len() as u32;
                    }
                    Err(e) => {
                        log::error!("failed to allocate with UE4 allocator: {e}.");
                    }
                }
            }
        }
    }

    pub fn push(&mut self, value: T) {
        self.ensure_space(self.count as usize + 1);
        self.data.map(|ptr| unsafe {
            ptr.as_ptr().offset(self.count as isize).write(value);
            self.count += 1;
        });
    }

    pub fn len(&self) -> usize {
        self.count as usize
    }

    pub fn capacity(&self) -> usize {
        self.max as usize
    }

    pub fn as_slice(&self) -> &[T] {
        match self.data {
            Some(ptr) => unsafe { core::slice::from_raw_parts(ptr.as_ptr(), self.len()) },
            None => unsafe { core::slice::from_raw_parts(NonNull::dangling().as_ptr(), 0) },
        }
    }

    pub fn as_slice_mut(&mut self) -> &mut [T] {
        match self.data {
            Some(ptr) => unsafe { core::slice::from_raw_parts_mut(ptr.as_ptr(), self.len()) },
            None => unsafe { core::slice::from_raw_parts_mut(NonNull::dangling().as_ptr(), 0) },
        }
    }
}

impl<T, I: SliceIndex<[T]>> Index<I> for TArray<T> {
    type Output = I::Output;

    fn index(&self, i: I) -> &Self::Output {
        let data = self.as_slice();

        &data[i]
    }
}

impl<T, I: SliceIndex<[T]>> IndexMut<I> for TArray<T> {
    fn index_mut(&mut self, i: I) -> &mut Self::Output {
        let data = self.as_slice_mut();

        &mut data[i]
    }
}

impl<T> std::ops::Deref for TArray<T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        self.as_slice()
    }
}

impl<T> std::ops::DerefMut for TArray<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.as_slice_mut()
    }
}

pub struct IntoIter<T> {
    //array_ref: &'a TArray<T>,
    ptr: *const T,
    end: *const T,
}

impl<'a, T> IntoIter<T> {
    pub fn new(array: &'a TArray<T>) -> Self {
        let ptr = array.data.unwrap().as_ptr();
        Self {
            ptr,
            end: unsafe { ptr.offset(array.count as isize) },
        }
    }
}

impl<'a, T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        if self.ptr == self.end {
            None
        } else {
            let old = self.ptr;
            self.ptr = unsafe { self.ptr.offset(1) };

            Some(unsafe { std::ptr::read(old) })
        }
    }
}

impl<'a, T> IntoIterator for &'a TArray<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter::new(self)
    }
}