dkmap.rs

use anyhow::{bail, Result};
use serde::{Deserialize, Deserializer, Serialize};
use std::cmp::Eq;
use std::collections::HashMap;
use std::hash::Hash;

#[derive(Debug, Serialize, Clone)]
pub struct DKMap<K1: Hash + Eq, K2: Hash + Eq, Values> {
    #[serde(skip)]
    key_mapping_1: HashMap<K1, usize>,
    #[serde(skip)]
    key_mapping_2: HashMap<K2, usize>,

    values: Vec<Values>,
}

#[derive(Debug, Deserialize, Clone)]
struct DKMapDeserializeProxy<Values> {
    values: Vec<Values>,
}

impl<K1, K2, RemainingValues> DKMap<K1, K2, (K1, K2, RemainingValues)>
where
    K1: Hash + Eq + Copy,
    K2: Hash + Eq + Clone,
    RemainingValues: Copy,
{
    pub fn get_k1(&self, key: &K1) -> Option<(K1, K2, RemainingValues)> {
        if let Some(index) = self.key_mapping_1.get(key) {
            Some(self.values[*index].clone())
        } else {
            None
        }
    }

    pub fn get_k2(&self, key: &K2) -> Option<(K1, K2, RemainingValues)> {
        if let Some(index) = self.key_mapping_2.get(key) {
            Some(self.values[*index].clone())
        } else {
            None
        }
    }

    pub fn update_from_k1(
        &mut self,
        key: &K1,
        remaining_values: RemainingValues,
    ) -> Result<(K1, K2, RemainingValues)> {
        if let Some(index) = self.key_mapping_1.get(key) {
            let ret = self.values[*index].clone();
            self.values[*index] = (ret.0, ret.1.clone(), remaining_values);
            Ok(ret)
        } else {
            bail!("Called DKMap::update_from_k1() on a non existing key")
        }
    }

    pub fn update_from_k2(
        &mut self,
        key: &K2,
        remaining_values: RemainingValues,
    ) -> Result<(K1, K2, RemainingValues)> {
        if let Some(index) = self.key_mapping_2.get(key) {
            let ret = self.values[*index].clone();
            self.values[*index] = (ret.0, ret.1.clone(), remaining_values);
            Ok(ret)
        } else {
            bail!("Called DKMap::update_from_k2() on a non existing key")
        }
    }

    pub fn insert(
        &mut self,
        value: (K1, K2, RemainingValues),
    ) -> Option<(K1, K2, RemainingValues)> {
        let k1 = value.0;
        let k2 = value.1;
        let remaining_values = value.2;

        // let's arbitrarily do it through K1
        if let Some(index) = self.key_mapping_1.get(&k1) {
            let old_values = self.values[*index].clone();
            self.values[*index] = (k1, k2, remaining_values);
            Some(old_values)
        } else {
            self.values.push((k1, k2.clone(), remaining_values));
            let index = self.values.len() - 1;
            self.key_mapping_1.insert(k1, index);
            self.key_mapping_2.insert(k2, index);
            None
        }
    }
}

impl<K1, K2, RemainingValues> FromIterator<(K1, K2, RemainingValues)>
    for DKMap<K1, K2, (K1, K2, RemainingValues)>
where
    K1: Hash + Eq + Copy,
    K2: Hash + Eq + Clone,
    RemainingValues: Copy,
{
    fn from_iter<T: IntoIterator<Item = (K1, K2, RemainingValues)>>(iter: T) -> Self {
        let mut len = 0;
        let iter = iter.into_iter().map(move |e| {
            len = len + 1;
            e
        });

        let mut ret = Self {
            key_mapping_1: HashMap::new(),
            key_mapping_2: HashMap::new(),
            values: Vec::with_capacity(len),
        };

        for (key1, key2, remaining_values) in iter.into_iter() {
            ret.insert((key1, key2, remaining_values));
        }

        ret
    }
}

impl<K1, K2, RemainingValues> IntoIterator for DKMap<K1, K2, (K1, K2, RemainingValues)>
where
    K1: Hash + Eq + Copy,
    K2: Hash + Eq + Clone,
    RemainingValues: Copy,
{
    type Item = (K1, K2, RemainingValues);
    type IntoIter = std::vec::IntoIter<Self::Item>;

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

// This is the trait that informs Serde how to deserialize MyMap.
impl<'de, K1, K2, RemainingValues> Deserialize<'de> for DKMap<K1, K2, (K1, K2, RemainingValues)>
where
    K1: Hash + Eq + Copy + Deserialize<'de>,
    K2: Hash + Eq + Clone + Deserialize<'de>,
    RemainingValues: Copy + Deserialize<'de>,
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let proxy = DKMapDeserializeProxy::<(K1, K2, RemainingValues)>::deserialize(deserializer)?;

        let mut key_mapping_1 = HashMap::<K1, usize>::new();
        let mut key_mapping_2 = HashMap::<K2, usize>::new();

        for (index, (k1, k2, _)) in proxy.values.clone().into_iter().enumerate() {
            key_mapping_1.insert(k1, index);
            key_mapping_2.insert(k2, index);
        }

        Ok(DKMap {
            key_mapping_1,
            key_mapping_2,
            values: proxy.values,
        })
    }
}