unreal-niggery-rs/src/t_array.rs

use std::fmt::{Display, Formatter};
use std::sync::OnceLock;

use widestring::WideCStr;

use crate::{Add, Container};
use crate::f_string::FString;

const T_ARRAY_HEADER_SIZE: usize = 0x10;
const DATA_PTR_OFFSET: usize = 0x00;
const LEN_OFFSET: usize = 0x08;
const CAP_OFFSET: usize = 0x0C;

type TArrayResize = fn(ptr: usize, len: u32);
static T_ARRAY_RESIZE_GROW: OnceLock<usize> = OnceLock::new();

pub struct TArray {
    pub(crate) ptr: *const u8,
    pub(crate) data_ptr: u64,
    pub(crate) len: u32,
    pub(crate) cap: u32,
}

impl TArray {
    pub fn set_t_array_resize_grow_ptr(t_array_resize_grow: usize) {
        T_ARRAY_RESIZE_GROW.set(t_array_resize_grow).unwrap();
    }

    pub fn read(ptr: usize) -> Self {
        let data = unsafe { std::slice::from_raw_parts(ptr as *const u8, T_ARRAY_HEADER_SIZE) };
        Self {
            ptr: ptr as *const u8,
            data_ptr: u64::from_le_bytes(data[DATA_PTR_OFFSET..LEN_OFFSET].try_into().unwrap()),
            len: u32::from_le_bytes(data[LEN_OFFSET..CAP_OFFSET].try_into().unwrap()),
            cap: u32::from_le_bytes(data[CAP_OFFSET..T_ARRAY_HEADER_SIZE].try_into().unwrap()),
        }
    }

    pub fn capacity(&self) -> u32 {
        self.cap
    }

    pub fn clear(&mut self, len: u32) -> Self {
        if self.data_ptr != 0 && self.len > 0 {
            let data = unsafe {
                std::slice::from_raw_parts_mut(self.data_ptr as *mut u8, self.len as usize)
            };
            data.fill(0);
            self.len = 0;
        }
    }

    #[inline(always)]
    pub(crate) fn resize_grow(&mut self, count: Option<u32>) -> u32 {
        let count = count.unwrap_or(1);
        let old_len = self.len;
        self.len = old_len + count;
        let data = unsafe {
            std::slice::from_raw_parts_mut(self.ptr as *mut u8, T_ARRAY_HEADER_SIZE)
        };
        self.write_len_internal(data);
        if self.cap < old_len + count {
            unsafe {
                std::mem::transmute::<usize, TArrayResize>(*t_array_resize_grow())(
                    self.ptr as usize,
                    old_len,
                );
            }
            // ResizeGrow will change capacity and might change data_ptr, so force refresh
            self.data_ptr = u64::from_le_bytes(data[DATA_PTR_OFFSET..LEN_OFFSET].try_into().unwrap());
            self.len = u32::from_le_bytes(data[LEN_OFFSET..CAP_OFFSET].try_into().unwrap());
            self.cap = u32::from_le_bytes(data[CAP_OFFSET..T_ARRAY_HEADER_SIZE].try_into().unwrap());
        }
        old_len
    }

    #[inline(always)]
    pub(crate) fn release_header_stack(&mut self) {
        let data = unsafe {
            std::slice::from_raw_parts_mut(self.ptr as *mut u8, T_ARRAY_HEADER_SIZE)
        };
        data.fill(0);
        self.data_ptr = u64::from_le_bytes(data[DATA_PTR_OFFSET..LEN_OFFSET].try_into().unwrap());
        self.len = u32::from_le_bytes(data[LEN_OFFSET..CAP_OFFSET].try_into().unwrap());
        self.cap = u32::from_le_bytes(data[CAP_OFFSET..T_ARRAY_HEADER_SIZE].try_into().unwrap());
    }

    #[inline(always)]
    fn write(&self) {
        let data = unsafe {
            std::slice::from_raw_parts_mut(self.ptr as *mut u8, T_ARRAY_HEADER_SIZE)
        };
        self.write_data_ptr_internal(data);
        self.write_len_internal(data);
        self.write_cap_internal(data);
    }

    #[inline(always)]
    fn write_data_ptr_internal(&self, data: &mut [u8]) {
        let bytes = self.data_ptr.to_le_bytes();
        data[DATA_PTR_OFFSET..LEN_OFFSET].copy_from_slice(&bytes);
    }

    #[inline(always)]
    fn write_len_internal(&self, data: &mut [u8]) {
        let bytes = self.len.to_le_bytes();
        data[LEN_OFFSET..CAP_OFFSET].copy_from_slice(&bytes);
    }

    #[inline(always)]
    fn write_cap_internal(&self, data: &mut [u8]) {
        let bytes = self.cap.to_le_bytes();
        data[CAP_OFFSET..T_ARRAY_HEADER_SIZE].copy_from_slice(&bytes);
    }
}

impl<'a> Container<Vec<&'a WideCStr>> for TArray {
    fn read_container(&self) -> Vec<&'a WideCStr> {
        let mut result = Vec::with_capacity(self.len as usize);
        for i in 0..self.len {
            result.push(
                FString(
                    Self::read(self.data_ptr as usize + (i as usize * T_ARRAY_HEADER_SIZE))
                ).read_container()
            );
        }
        result
    }
}

impl Add for TArray {
    fn add(&mut self, other: &mut Self) -> &mut Self {
        let old_len = self.resize_grow(None);
        Self {
            ptr: (self.data_ptr as usize + (old_len as usize * T_ARRAY_HEADER_SIZE)) as *const u8,
            data_ptr: other.data_ptr,
            len: other.len,
            cap: other.cap,
        }.write();
        // Clear the of the original container
        other.release_header_stack();
        self
    }
}

impl Display for TArray {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let elements = self.read_container();
        writeln!(f, "[")?;
        for element in elements {
            writeln!(f, "\t- {}", element.to_string_lossy())?;
        }
        write!(f, "]")
    }
}

#[inline(always)]
fn t_array_resize_grow() -> &'static usize {
    T_ARRAY_RESIZE_GROW.get().unwrap()
}