1use drain_while::DrainWhileable;
2use std::collections::HashMap;
3
4#[derive(Default, Debug, Clone)]
5pub struct OpMap<T>
6where
7    T: Clone + std::default::Default,
8{
9    pub map: HashMap<String, Vec<T>>,
10}
11
12impl<T> OpMap<T>
13where
14    T: Clone + std::default::Default,
15{
16    pub fn join(&mut self, mut other: OpMap<T>, sort_fn: fn(&T, &T) -> std::cmp::Ordering) {
17        for (k, v) in other.map.iter_mut() {
18            v.sort_by(sort_fn);
19            if self.map.contains_key(k) {
20                self.map
21                    .entry(k.to_owned())
22                    .or_default()
23                    .extend(v.to_owned());
24            } else {
25                self.map.insert(k.to_owned(), v.to_owned());
26            }
27        }
28    }
29
30    pub fn insert(&mut self, k: &str, v: T) {
31        self.map.entry(k.to_owned()).or_default().push(v);
32    }
33
34    pub fn set(&mut self, k: &str, v: Vec<T>) {
35        self.map.insert(k.to_string(), v);
36    }
37
38    pub fn get(&mut self, k: &str) -> Option<&Vec<T>> {
39        self.map.get(k)
40    }
41
42    pub fn drain<P>(&mut self, k: &str, predicate: P) -> Vec<T>
43    where
44        P: Fn(&T) -> bool,
45    {
46        if self.map.contains_key(k) {
47            let v = self.map.get_mut(k).unwrap();
48            v.drain_while(predicate).collect()
49        } else {
50            vec![]
51        }
52    }
53
54    pub fn from_vec<P>(v: Vec<T>, key_from_t: P) -> OpMap<T>
55    where
56        P: Fn(&T) -> &str,
57    {
58        let mut result: OpMap<T> = OpMap::default();
59        v.iter().for_each(|t| {
60            let k = key_from_t(t);
61            result.insert(k, t.to_owned());
62        });
63
64        result
65    }
66}
67
68#[cfg(test)]
69mod test {
70    use super::*;
71    use hamcrest2::prelude::*;
72
73    #[derive(Default, Debug, Clone, PartialEq)]
74    struct TestStruct {
75        t: f32,
76    }
77
78    #[test]
79    fn test_insert() {
80        let mut ops = OpMap::default();
81        let a = TestStruct { t: 0.0 };
82        let b = TestStruct { t: 1.0 };
83        let c = TestStruct { t: 2.0 };
84        let d = TestStruct { t: 3.0 };
85
86        ops.insert("x", a.clone());
87        ops.insert("y", b.clone());
88        ops.insert("x", c.clone());
89        ops.insert("z", d.clone());
90
91        assert_that!(ops.map.get("x").unwrap(), contains(vec![a, c]).exactly());
92        assert_that!(ops.map.get("y").unwrap(), contains(vec![b]).exactly());
93        assert_that!(ops.map.get("z").unwrap(), contains(vec![d]).exactly());
94    }
95
96    #[test]
97    fn test_join() {
98        let mut ops1 = OpMap::default();
99        let mut ops2 = OpMap::default();
100        let a = TestStruct { t: 0.0 };
101        let b = TestStruct { t: 1.0 };
102        let c = TestStruct { t: 2.0 };
103        let d = TestStruct { t: 3.0 };
104        let e = TestStruct { t: 4.0 };
105        let f = TestStruct { t: 5.0 };
106
107        ops1.insert("x", a.clone());
108        ops1.insert("y", b.clone());
109
110        ops2.insert("x", d.clone());
111        ops2.insert("x", c.clone());
112        ops2.insert("y", e.clone());
113        ops2.insert("z", f.clone());
114
115        ops1.join(ops2, |a, b| a.t.partial_cmp(&b.t).unwrap());
116
117        assert_that!(
118            ops1.map.get("x").unwrap(),
119            contains(vec![a, c, d]).in_order()
120        );
121        assert_that!(ops1.map.get("y").unwrap(), contains(vec![b, e]).in_order());
122        assert_that!(ops1.map.get("z").unwrap(), contains(vec![f]).in_order());
123    }
124}