PkmnLib_rs/src/static_data/libraries/type_library.rs

use anyhow_ext::Result;
use atomig::Atom;
use hashbrown::HashMap;
use parking_lot::{RwLock, RwLockReadGuard};
use std::fmt::{Debug, Display};

use crate::{PkmnError, StringKey};

/// A unique key that can be used to store a reference to a type. Opaque reference to a byte
/// internally.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Default, Hash, Atom)]
#[repr(C)]
pub struct TypeIdentifier {
    /// The unique internal value.
    val: u8,
}

impl From<u8> for TypeIdentifier {
    fn from(val: u8) -> Self {
        Self { val }
    }
}

impl From<TypeIdentifier> for u8 {
    fn from(id: TypeIdentifier) -> Self {
        id.val
    }
}

impl Display for TypeIdentifier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "TypeId({})", self.val)
    }
}

/// All data related to types and effectiveness.
pub trait TypeLibrary: Debug {
    /// Gets the type identifier for a type with a name.
    fn get_type_id(&self, key: &StringKey) -> Option<TypeIdentifier>;

    /// Gets the type name from the type identifier.
    fn get_type_name(&self, t: TypeIdentifier) -> Option<StringKey>;

    /// Gets the effectiveness for a single attacking type against a single defending type.
    fn get_single_effectiveness(&self, attacking: TypeIdentifier, defending: TypeIdentifier) -> Result<f32>;

    /// Gets the effectiveness for a single attacking type against an amount of defending types.
    /// This is equivalent to running [`get_single_effectiveness`] on each defending type, and
    /// multiplying the results with each other.
    fn get_effectiveness(&self, attacking: TypeIdentifier, defending: &[TypeIdentifier]) -> Result<f32>;

    /// Registers a new type in the library.
    fn register_type(&self, name: &StringKey) -> TypeIdentifier;

    /// Sets the effectiveness for an attacking type against a defending type.
    fn set_effectiveness(&self, attacking: TypeIdentifier, defending: TypeIdentifier, effectiveness: f32)
        -> Result<()>;
}

/// All data related to types and effectiveness.
#[derive(Debug)]
pub struct TypeLibraryImpl {
    /// A list of types
    types: RwLock<HashMap<StringKey, TypeIdentifier>>,
    /// The effectiveness of the different types against each other.
    effectiveness: RwLock<Vec<Vec<f32>>>,
}

impl TypeLibraryImpl {
    /// Instantiates a new type library with a specific capacity.
    pub fn new(capacity: usize) -> Self {
        Self {
            types: RwLock::new(HashMap::with_capacity(capacity)),
            effectiveness: RwLock::new(vec![]),
        }
    }
}

impl TypeLibraryImpl {
    fn get_single_effectiveness_with_lock(
        lock: &RwLockReadGuard<Vec<Vec<f32>>>,
        attacking: TypeIdentifier,
        defending: TypeIdentifier,
    ) -> Result<f32> {
        Ok(*lock
            .get((attacking.val - 1) as usize)
            .ok_or(PkmnError::InvalidTypeIdentifier { type_id: attacking })?
            .get((defending.val - 1) as usize)
            .ok_or(PkmnError::InvalidTypeIdentifier { type_id: defending })?)
    }
}

impl TypeLibrary for TypeLibraryImpl {
    /// Gets the type identifier for a type with a name.
    fn get_type_id(&self, key: &StringKey) -> Option<TypeIdentifier> {
        self.types.read().get(key).cloned()
    }

    /// Gets the type name from the type identifier.
    fn get_type_name(&self, t: TypeIdentifier) -> Option<StringKey> {
        let types = self.types.read();
        for kv in types.iter() {
            if *kv.1 == t {
                return Some(kv.0.clone());
            }
        }
        None
    }

    /// Gets the effectiveness for a single attacking type against a single defending type.
    fn get_single_effectiveness(&self, attacking: TypeIdentifier, defending: TypeIdentifier) -> Result<f32> {
        Self::get_single_effectiveness_with_lock(&self.effectiveness.read(), attacking, defending)
    }

    /// Gets the effectiveness for a single attacking type against an amount of defending types.
    /// This is equivalent to running [`get_single_effectiveness`] on each defending type, and
    /// multiplying the results with each other.
    fn get_effectiveness(&self, attacking: TypeIdentifier, defending: &[TypeIdentifier]) -> Result<f32> {
        let mut e = 1.0;
        let lock = self.effectiveness.read();
        for def in defending {
            e *= Self::get_single_effectiveness_with_lock(&lock, attacking, *def)?;
        }
        Ok(e)
    }

    /// Registers a new type in the library.
    fn register_type(&self, name: &StringKey) -> TypeIdentifier {
        let mut types_write_lock = self.types.write();
        let mut effectiveness_write_lock = self.effectiveness.write();

        let id = TypeIdentifier {
            val: (types_write_lock.len() + 1) as u8,
        };
        types_write_lock.insert(name.clone(), id);
        effectiveness_write_lock.resize((id.val) as usize, vec![]);
        for effectiveness in &mut effectiveness_write_lock.iter_mut() {
            effectiveness.resize((id.val) as usize, 1.0)
        }
        id
    }

    /// Sets the effectiveness for an attacking type against a defending type.
    fn set_effectiveness(
        &self,
        attacking: TypeIdentifier,
        defending: TypeIdentifier,
        effectiveness: f32,
    ) -> Result<()> {
        *self
            .effectiveness
            .write()
            .get_mut((attacking.val - 1) as usize)
            .ok_or(PkmnError::InvalidTypeIdentifier { type_id: attacking })?
            .get_mut((defending.val - 1) as usize)
            .ok_or(PkmnError::InvalidTypeIdentifier { type_id: defending })? = effectiveness;
        Ok(())
    }
}

#[cfg(test)]
#[allow(clippy::indexing_slicing)]
#[allow(clippy::unwrap_used)]
pub mod tests {
    use assert_approx_eq::assert_approx_eq;

    use super::*;
    use crate::static_data::libraries::type_library::TypeLibrary;

    pub fn build() -> TypeLibraryImpl {
        let mut lib = TypeLibraryImpl::new(2);

        // Borrow as mut so we can insert
        let w = &mut lib;
        let t0 = w.register_type(&"foo".into());
        let t1 = w.register_type(&"bar".into());
        // Drops borrow as mut

        w.set_effectiveness(t0, t1, 0.5).unwrap();
        w.set_effectiveness(t1, t0, 2.0).unwrap();

        lib
    }

    #[test]
    fn add_two_types_retrieve_them() {
        let mut lib = TypeLibraryImpl::new(2);

        // Borrow as mut so we can insert
        let w = &mut lib;
        let t0 = w.register_type(&"foo".into());
        let t1 = w.register_type(&"bar".into());
        // Drops borrow as mut

        // Borrow as read so we can read
        let r = &lib;
        assert_eq!(r.get_type_id(&"foo".into()).unwrap(), t0);
        assert_eq!(r.get_type_id(&"bar".into()).unwrap(), t1);
    }

    #[test]
    fn add_two_types_set_effectiveness_retrieve() {
        let mut lib = TypeLibraryImpl::new(2);

        // Borrow as mut so we can insert
        let w = &mut lib;
        let t0 = w.register_type(&"foo".into());
        let t1 = w.register_type(&"bar".into());
        w.set_effectiveness(t0, t1, 0.5).unwrap();
        w.set_effectiveness(t1, t0, 2.0).unwrap();
        // Drops borrow as mut

        // Borrow as read so we can read
        let r = &lib;
        assert_approx_eq!(r.get_single_effectiveness(t0, t1).unwrap(), 0.5);
        assert_approx_eq!(r.get_single_effectiveness(t1, t0).unwrap(), 2.0);
    }

    #[test]
    fn add_two_types_get_aggregate_effectiveness() {
        let mut lib = TypeLibraryImpl::new(2);

        // Borrow as mut so we can insert
        let w = &mut lib;
        let t0 = w.register_type(&"foo".into());
        let t1 = w.register_type(&"bar".into());
        w.set_effectiveness(t0, t1, 0.5).unwrap();
        w.set_effectiveness(t1, t0, 2.0).unwrap();
        // Drops borrow as mut

        // Borrow as read so we can read
        let r = &lib;
        assert_approx_eq!(r.get_effectiveness(t0, &[t1, t1]).unwrap(), 0.25);
        assert_approx_eq!(r.get_effectiveness(t1, &[t0, t0]).unwrap(), 4.0);
    }

    #[test]
    fn add_two_types_get_type_name() {
        let mut lib = TypeLibraryImpl::new(2);

        // Borrow as mut so we can insert
        let w = &mut lib;
        let t0 = w.register_type(&"foo".into());
        let t1 = w.register_type(&"bar".into());
        // Drops borrow as mut

        // Borrow as read so we can read
        let r = &lib;
        assert_eq!(r.get_type_name(t0).unwrap(), "foo".into());
        assert_eq!(r.get_type_name(t1).unwrap(), "bar".into());
    }
}