common_stdx/
dualhash.rs

1use std::collections::{HashMap, HashSet};
2
3#[derive(Debug, PartialEq, Eq)]
4pub struct DualHashMap<K1, K2, V>
5where
6    K1: std::hash::Hash + Eq + Send + Clone,
7    K2: std::hash::Hash + Eq + Send + Clone,
8{
9    data_map: HashMap<K1, V>,
10    sec_key_map: HashMap<K2, HashSet<K1>>,
11}
12
13impl<K1, K2, V> DualHashMap<K1, K2, V>
14where
15    K1: std::hash::Hash + Eq + Send + Clone,
16    K2: std::hash::Hash + Eq + Send + Clone,
17{
18    pub fn new() -> Self {
19        DualHashMap {
20            data_map: HashMap::new(),
21            sec_key_map: HashMap::new(),
22        }
23    }
24
25    pub fn insert(&mut self, k1: K1, k2: K2, value: V) {
26        for set in self.sec_key_map.values_mut() {
27            set.remove(&k1);
28        }
29        self.sec_key_map.retain(|_, set| !set.is_empty());
30
31        self.data_map.insert(k1.clone(), value);
32        self.sec_key_map.entry(k2).or_default().insert(k1);
33    }
34
35    pub fn get_by_k1(&self, k1: &K1) -> Option<&V> {
36        self.data_map.get(k1)
37    }
38
39    pub fn get_by_k2(&self, k2: &K2) -> impl Iterator<Item = &V> {
40        self.sec_key_map
41            .get(k2)
42            .into_iter()
43            .flat_map(|k1s| k1s.iter())
44            .filter_map(|k1| self.data_map.get(k1))
45    }
46
47    pub fn get_mut_by_k1(&mut self, k1: &K1) -> Option<&mut V> {
48        self.data_map.get_mut(k1)
49    }
50
51    pub fn for_each_mut_by_k2<F>(&mut self, k2: &K2, mut f: F)
52    where
53        F: FnMut(&mut V),
54    {
55        if let Some(k1s) = self.sec_key_map.get(k2) {
56            for k1 in k1s.iter() {
57                if let Some(v) = self.data_map.get_mut(k1) {
58                    f(v);
59                }
60            }
61        }
62    }
63
64    pub fn get_values_mut(&mut self) -> impl Iterator<Item = &mut V> {
65        self.data_map.values_mut()
66    }
67    pub fn get_values(&self) -> impl Iterator<Item = &V> {
68        self.data_map.values()
69    }
70    pub fn get_keys1(&self) -> impl Iterator<Item = &K1> {
71        self.data_map.keys()
72    }
73    pub fn get_keys2(&self) -> impl Iterator<Item = &K2> {
74        self.sec_key_map.keys()
75    }
76
77    pub fn remove_by_k1(&mut self, k1: &K1) {
78        if self.data_map.remove(k1).is_some() {
79            for (_, k1_set) in self.sec_key_map.iter_mut() {
80                k1_set.remove(k1);
81            }
82            self.sec_key_map.retain(|_, set| !set.is_empty());
83        }
84    }
85
86    pub fn remove_by_k2(&mut self, k2: &K2) {
87        if let Some(k1s) = self.sec_key_map.remove(k2) {
88            for k1 in k1s {
89                self.data_map.remove(&k1);
90            }
91        }
92    }
93
94    pub fn contains_k1(&self, k1: &K1) -> bool {
95        self.data_map.contains_key(k1)
96    }
97
98    pub fn contains_k2(&self, k2: &K2) -> bool {
99        self.sec_key_map.contains_key(k2)
100    }
101
102    pub fn get_k1s_by_k2(&self, k2: &K2) -> Option<&HashSet<K1>> {
103        self.sec_key_map.get(k2)
104    }
105
106    pub fn get_k2_by_k1(&self, k1: &K1) -> Option<&K2> {
107        for (k2, k1_set) in &self.sec_key_map {
108            if k1_set.contains(k1) {
109                return Some(k2);
110            }
111        }
112        None
113    }
114
115    pub fn len(&self) -> usize {
116        self.data_map.len()
117    }
118
119    pub fn is_empty(&self) -> bool {
120        self.data_map.is_empty()
121    }
122}
123
124impl<K1, K2, V> Default for DualHashMap<K1, K2, V>
125where
126    K1: std::hash::Hash + Eq + Send + Clone,
127    K2: std::hash::Hash + Eq + Send + Clone,
128{
129    fn default() -> Self {
130        Self::new()
131    }
132}
common_stdx/dualhash.rs

common_stdx/
dualhash.rs
