Documentation
use std::collections::{BTreeSet, HashMap};
use std::hash::Hash;

#[derive(Clone, Debug)]
pub struct LatestMap<Key: Eq + Hash + Clone + Ord, Value> {
    pub(crate) data: HashMap<Key, Value>,
    data_index: BTreeSet<Key>,
}

impl<Key: Eq + Hash + Clone + Ord, Value> Default for LatestMap<Key, Value> {
    fn default() -> Self {
        Self {
            data: HashMap::new(),
            data_index: BTreeSet::new(),
        }
    }
}

impl<Key: Eq + Hash + Clone + Ord, Value> LatestMap<Key, Value> {
    pub fn insert(&mut self, key: Key, value: Value) {
        self.data.insert(key.clone(), value);
        self.data_index.insert(key);
    }
    pub fn get_latest<'a>(&'a self, key: &'a Key) -> Option<(&'a Key, &'a Value)> {
        let target_key = if self.data_index.contains(key) {
            key
        } else {
            let sorted_keys: Vec<&Key> = self.data_index.iter().collect();
            match sorted_keys.binary_search(&key) {
                Ok(_) => key,
                Err(gt_index) => {
                    if gt_index == 0 {
                        return None;
                    }
                    sorted_keys[gt_index - 1]
                }
            }
        };

        self.data.get(target_key).map(|value| (target_key, value))
    }
    pub fn get_last_with_key(&self) -> Option<(&Key, &Value)> {
        let sorted_keys: Vec<&Key> = self.data_index.iter().collect();
        sorted_keys
            .last()
            .and_then(|&key| self.data.get(key).map(move |v| (key, v)))
    }

    pub fn get_mut(&mut self, key: &Key) -> Option<&mut Value> {
        self.data.get_mut(key)
    }
    pub fn contains_key(&self, key: &Key) -> bool {
        self.data_index.contains(key)
    }

    pub fn pop_latest(&mut self, key: &Key) -> Option<(Key, Value)> {
        let target_key = if self.data_index.contains(key) {
            key.clone()
        } else {
            let sorted_keys: Vec<&Key> = self.data_index.iter().collect();
            match sorted_keys.binary_search(&key) {
                Ok(_) => key.clone(),
                Err(gt_index) => {
                    if gt_index == 0 {
                        return None;
                    }
                    sorted_keys[gt_index - 1].clone()
                }
            }
        };
        self.data_index.remove(&target_key);
        self.data.remove(&target_key).map(|value| (target_key, value))
    }
}

#[cfg(test)]
mod test {
    use crate::LatestMap;

    #[test]
    fn should_insert() {
        let mut map: LatestMap<i32, i32> = LatestMap::default();
        map.insert(1, 2);
        assert!(map.data_index.contains(&1));
        assert_eq!(map.data.get(&1), Some(&2));
    }

    #[test]
    fn should_contains_key() {
        let mut map: LatestMap<i32, i32> = LatestMap::default();
        map.insert(1, 2);
        assert!(map.contains_key(&1));
        assert!(!map.contains_key(&2));
    }

    #[test]
    fn should_get_mut() {
        let mut map: LatestMap<i32, i32> = LatestMap::default();
        map.insert(1, 2);
        let value = map.get_mut(&1).unwrap();
        *value = 3;

        assert_eq!(map.data.get(&1), Some(&3));
    }

    #[test]
    fn should_get_latest() {
        let mut map: LatestMap<i32, i32> = LatestMap::default();
        map.insert(1, 2);
        map.insert(10, 20);
        map.insert(20, 40);
        map.insert(50, 100);

        assert_eq!(map.get_latest(&0), None);
        assert_eq!(map.get_latest(&1), Some((&1, &2)));
        assert_eq!(map.get_latest(&3), Some((&1, &2)));
        assert_eq!(map.get_latest(&20), Some((&20, &40)));
        assert_eq!(map.get_latest(&24), Some((&20, &40)));
        assert_eq!(map.get_latest(&1000), Some((&50, &100)));
    }

    #[test]
    fn should_work_given_map_is_empty() {
        let map: LatestMap<i32, i32> = LatestMap::default();
        assert_eq!(map.get_latest(&0), None);
        assert_eq!(map.get_latest(&1), None);
        assert_eq!(map.get_latest(&3), None);
    }

    #[test]
    fn should_work_pop_latest() {
        let mut map: LatestMap<i32, i32> = LatestMap::default();
        map.insert(1, 2);
        map.insert(10, 20);
        map.insert(20, 40);
        map.insert(50, 100);

        assert_eq!(map.get_latest(&0), None);
        assert_eq!(map.data.len(), 4);
        assert_eq!(map.get_latest(&1), Some((&1, &2)));
        assert_eq!(map.data.len(), 4);
        assert_eq!(map.pop_latest(&3), Some((1, 2)));
        assert_eq!(map.data.len(), 3);
        assert_eq!(map.data_index.len(), 3);
    }
}