Gen7ScriptsRs/pkmn_lib_interface/src/handling/temporary.rs

use crate::{ExternRef, ExternalReferenceType};
use alloc::boxed::Box;
use core::cell::Cell;
use core::mem::forget;

struct TemporaryData<T> {
    use_count: Cell<usize>,
    is_deleted: bool,
    value: T,
}

#[cfg(not(feature = "mock_data"))]
pub struct Temporary<T> {
    value: *mut TemporaryData<T>,
}

#[cfg(not(feature = "mock_data"))]
static mut TEMPORARIES: alloc::collections::BTreeMap<u32, *const u8> =
    alloc::collections::BTreeMap::new();

#[cfg(not(feature = "mock_data"))]
impl<T> Temporary<T>
where
    T: ExternalReferenceType,
    T: Clone,
{
    #[inline]
    pub fn from_reference(reference: ExternRef<T>) -> Self {
        let temporaries = unsafe { &mut TEMPORARIES };
        let existing = temporaries.get(&reference.get_internal_index());
        unsafe {
            if let Some(v) = existing {
                let rc = (*v as *mut TemporaryData<T>).as_mut().unwrap();
                if !rc.is_deleted {
                    *rc.use_count.get_mut() += 1;
                    return Self {
                        value: rc as *mut TemporaryData<T>,
                    };
                }
            }
        }

        let value = reference.get_value().unwrap();
        let mut reference_counter = TemporaryData {
            use_count: Cell::new(1),
            is_deleted: false,
            value,
        };
        let ptr = &mut reference_counter as *mut TemporaryData<T>;
        forget(reference_counter);
        temporaries.insert(reference.get_internal_index(), ptr as *const u8);
        Self { value: ptr }
    }

    #[inline]
    pub fn value(&self) -> T {
        unsafe { self.value.as_ref().unwrap().value.clone() }
    }

    #[inline]
    pub fn value_ref(&self) -> &T {
        unsafe { &self.value.as_ref().unwrap().value }
    }

    pub(crate) fn mark_as_deleted(reference: ExternRef<T>) {
        let temporaries = unsafe { &mut TEMPORARIES };
        let existing = temporaries.get(&reference.get_internal_index());
        unsafe {
            if let Some(v) = existing {
                crate::utils::print_raw(b"Dropping temporary");
                let rc = (*v as *mut TemporaryData<T>).as_mut().unwrap();
                rc.is_deleted = true;
                if rc.use_count.get() == 0 {
                    drop(Box::from(*v as *mut TemporaryData<T>))
                }
                temporaries.remove(&reference.get_internal_index());
            }
        }
    }
}

#[cfg(not(feature = "mock_data"))]
impl<T> Clone for Temporary<T> {
    fn clone(&self) -> Self {
        *unsafe { self.value.as_mut() }.unwrap().use_count.get_mut() += 1;
        Self { value: self.value }
    }
}

#[cfg(not(feature = "mock_data"))]
impl<T> Drop for Temporary<T> {
    fn drop(&mut self) {
        let data = unsafe { self.value.as_mut() }.unwrap();
        let v = data.use_count.get_mut();
        *v -= 1;
        if *v == 0 && data.is_deleted {
            drop(Box::from(self.value))
        }
    }
}

#[cfg(feature = "mock_data")]
pub struct Temporary<T: Clone> {
    value: T,
}

#[cfg(feature = "mock_data")]
impl<T: Clone> Temporary<T> {
    #[inline]
    pub fn value(&self) -> T {
        self.value.clone()
    }

    #[inline]
    pub fn value_ref(&self) -> &T {
        &self.value
    }
}