ibc-test-framework 0.32.2

Framework for writing integration tests for IBC relayers
use std::collections::{btree_map, BTreeMap};

#[derive(Clone, Debug)]
pub struct TwoDimMap<T> {
    pub map: BTreeMap<usize, BTreeMap<usize, T>>,
}

impl<T> Default for TwoDimMap<T> {
    fn default() -> Self {
        TwoDimMap {
            map: BTreeMap::new(),
        }
    }
}

impl<T> TwoDimMap<T> {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn get(&self, (x, y): (usize, usize)) -> Option<&T> {
        self.map.get(&x).and_then(|inner| inner.get(&y))
    }

    pub fn insert(&mut self, (x, y): (usize, usize), value: T) -> Option<T> {
        if let Some(existing_values) = self.map.get_mut(&x) {
            existing_values.insert(y, value)
        } else {
            let mut new_values = BTreeMap::new();
            new_values.insert(y, value);
            self.map.insert(x, new_values);
            None
        }
    }

    pub fn iter(&self) -> Iter<T> {
        Iter {
            outer_iter: self.map.iter(),
            inner_iter: None,
            outer_key: 0,
        }
    }
}

pub struct Iter<'a, T> {
    outer_iter: btree_map::Iter<'a, usize, BTreeMap<usize, T>>,
    inner_iter: Option<btree_map::Iter<'a, usize, T>>,
    outer_key: usize,
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = (usize, usize, &'a T);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(inner_iter) = &mut self.inner_iter {
                if let Some((inner_key, inner_value)) = inner_iter.next() {
                    return Some((self.outer_key, *inner_key, inner_value));
                }
            }

            if let Some((outer_key, inner_map)) = self.outer_iter.next() {
                self.outer_key = *outer_key;
                self.inner_iter = Some(inner_map.iter());
            } else {
                return None;
            }
        }
    }
}

impl<T> From<BTreeMap<usize, BTreeMap<usize, T>>> for TwoDimMap<T> {
    fn from(map: BTreeMap<usize, BTreeMap<usize, T>>) -> Self {
        Self { map }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use super::*;

    #[test]
    fn test_two_dim_hash_map_iterator() {
        let mut outer_hashmap = BTreeMap::new();
        let mut inner_hashmap1 = BTreeMap::new();
        let mut inner_hashmap2 = BTreeMap::new();
        let mut inner_hashmap3 = BTreeMap::new();

        inner_hashmap1.insert(1, "a");

        inner_hashmap2.insert(2, "c");
        inner_hashmap2.insert(0, "b");

        inner_hashmap3.insert(1, "d");

        outer_hashmap.insert(0, inner_hashmap1);
        outer_hashmap.insert(1, inner_hashmap2);
        outer_hashmap.insert(2, inner_hashmap3);

        let two_dim_hash_map = TwoDimMap::from(outer_hashmap);
        let mut two_dim_hash_map_iter = two_dim_hash_map.iter();

        assert_eq!(two_dim_hash_map_iter.next(), Some((0, 1, &"a")));

        assert_eq!(two_dim_hash_map_iter.next(), Some((1, 0, &"b")));
        assert_eq!(two_dim_hash_map_iter.next(), Some((1, 2, &"c")));

        assert_eq!(two_dim_hash_map_iter.next(), Some((2, 1, &"d")));

        assert_eq!(two_dim_hash_map_iter.next(), None);
    }
}