wachter_storage/

type_map.rs

use std::{any::TypeId, collections::HashMap, error::Error, fmt::Display};

// Why do we need this 'static here?
// Because if this is to represent any type, then it might also need to represent references.
// So, we could conceivable want to create a `TypeMap<&Foo, Bar>`.
// In those cases, the borrow checker can't assert that references to Foo used as keys will live
// for long enough, so we use 'static.
// Incidentally, this also forbids reference types that are not 'static, so essentially all
// accidental reference types. You have to really try to make a reference type a key or value.
// pub trait Type: 'static + Sync + Send {}
pub trait Type: 'static {}
// This implements `Type` for every type.
impl<T: 'static> Type for T {}

#[derive(Debug, PartialEq)]
pub enum TypeQueryError {
    // Key, ExpectedVal
    Missing(TypeId, TypeId),
    // Key, ExpectedVal, ActualVal
    WrongType(TypeId, TypeId, TypeId),
}
impl Error for TypeQueryError {}

impl Display for TypeQueryError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TypeQueryError::Missing(key, expected) => f.write_str(&format!(
                "TypeQueryError::Empty(key={:?}, expected={:?})",
                &key, &expected
            )),
            TypeQueryError::WrongType(key, expected, actual) => f.write_str(&format!(
                "TypeQueryError::Empty(key={:?}, expected={:?}, actual={:?})",
                &key, &expected, &actual
            )),
        }
    }
}

pub type TypeQueryResult<T> = Result<T, TypeQueryError>;

// Like anymap, but stores a vector of somethings instead of somethings.
pub type RawMap = HashMap<TypeId, Box<dyn std::any::Any>>;

#[derive(Debug)]
pub struct TypeMap {
    raw: RawMap,
}

// TODO Add removals
// TODO Add unchecked queries for performance.
impl TypeMap {
    pub fn new() -> Self {
        Self {
            raw: HashMap::new(),
        }
    }
    pub fn contains_key<Key: Type>(&self) -> bool {
        self.raw.contains_key(&TypeId::of::<Key>())
    }

    // TODO Handle overrides gracefully.
    pub fn insert_raw<Value: Type>(&mut self, key_type: TypeId, val: Value) {
        self.raw.insert(key_type, Box::new(val));
    }

    pub fn insert<Key: Type, Value: Type>(&mut self, val: Value) {
        self.insert_raw(TypeId::of::<Key>(), val);
    }

    #[inline]
    pub fn get_raw<'value, Val: Type>(&'value self, key: TypeId) -> TypeQueryResult<&Val> {
        self.raw
            .get(&key)
            .ok_or(TypeQueryError::Missing(key, TypeId::of::<Val>()))
            .and_then(|elem| {
                elem.downcast_ref::<Val>().ok_or(TypeQueryError::WrongType(
                    key,
                    TypeId::of::<Val>(),
                    elem.type_id(),
                ))
            })
    }

    #[inline]
    pub fn get<'a, Key: Type, Val: Type>(&'a self) -> TypeQueryResult<&Val> {
        let key = TypeId::of::<Key>();
        self.get_raw::<Val>(key)
    }

    #[inline]
    pub fn get_mut_raw<'value, Val: Type>(
        &'value mut self,
        key: TypeId,
    ) -> TypeQueryResult<&mut Val> {
        self.raw
            .get_mut(&key)
            .ok_or(TypeQueryError::Missing(key, TypeId::of::<Val>()))
            .and_then(|elem| {
                let etid = elem.type_id();
                elem.downcast_mut::<Val>().ok_or(TypeQueryError::WrongType(
                    key,
                    TypeId::of::<Val>(),
                    etid,
                ))
            })
    }

    #[inline]
    pub fn get_mut<'a, Key: Type, Val: Type>(&'a mut self) -> TypeQueryResult<&mut Val> {
        let key = TypeId::of::<Key>();
        self.get_mut_raw::<Val>(key)
    }
}

#[cfg(test)]
mod test {
    use std::cell::RefCell;

    use super::*;

    #[test]
    fn test_map_of_u8() {
        let mut map = TypeMap::new();
        let data: u8 = 3;
        map.insert::<u8, u8>(data);
        assert_eq!(map.get_raw::<u8>(TypeId::of::<u8>()), Ok(&data));
        assert_eq!(map.get::<u8, u8>(), Ok(&data));
    }

    #[derive(PartialEq, Eq, Debug, Clone)]
    struct TestPoint {
        x: u8,
        y: u8,
    }

    #[test]
    fn test_map_of_struct() {
        let mut map = TypeMap::new();
        let data: TestPoint = TestPoint { x: 8, y: 8 };
        map.insert::<TestPoint, TestPoint>(data.clone());
    }

    #[test]
    fn test_map_of_vec_u8() {
        let mut map = TypeMap::new();
        let data: Vec<u8> = vec![2, 3];
        map.insert::<u8, Vec<u8>>(data.clone());
        // assert_eq!(map.get::<u8, Vec<u8>>(), Ok(Rc::new(data)));
        assert_eq!(map.get::<u8, Vec<u8>>(), Ok(&data));
    }

    #[test]
    fn test_mixed_map() {
        let mut map = TypeMap::new();
        let data_u8 = 8;
        map.insert::<u8, u8>(data_u8);
        map.insert::<u16, u8>(data_u8);
        assert_eq!(map.get::<u16, u8>(), Ok(&data_u8));
    }

    #[test]
    fn test_mutation() {
        let mut map = TypeMap::new();
        let data_u8 = 8;
        map.insert::<u8, RefCell<u8>>(RefCell::new(data_u8));
        let reference = map.get::<u8, RefCell<u8>>().expect("Error getting value");
        *reference.borrow_mut() = 9;
        assert_eq!(map.get::<u8, RefCell<u8>>(), Ok(&RefCell::new(9)));
    }

    #[test]
    fn test_missing_value() {
        let mut map = TypeMap::new();
        let data: Vec<u8> = vec![2, 3];
        map.insert::<u8, Vec<u8>>(data.clone());
        let typeof_u8 = TypeId::of::<u8>();
        let typeof_u16 = TypeId::of::<u16>();
        assert_eq!(
            map.get::<u16, u8>(),
            Err(TypeQueryError::Missing(typeof_u16, typeof_u8))
        );
    }

    #[test]
    fn test_wrong_type() {
        let mut map = TypeMap::new();
        let data: Vec<u8> = vec![2, 3];
        map.insert::<u8, Vec<u8>>(data.clone());
        let typeof_u8 = TypeId::of::<u8>();
        let typeof_vec_u8 = TypeId::of::<Vec<u8>>();
        assert_eq!(
            map.get::<u8, u8>(),
            Err(TypeQueryError::WrongType(
                typeof_u8,
                typeof_u8,
                typeof_vec_u8
            ))
        );
    }
}