Skip to main content

turbo_tasks/
keyed.rs

1use std::{
2    borrow::Borrow,
3    collections::{HashMap, HashSet},
4    hash::{BuildHasher, Hash},
5};
6
7use indexmap::{Equivalent, IndexMap, IndexSet};
8use smallvec::SmallVec;
9
10pub trait KeyedEq {
11    type Key;
12
13    fn different_keys<'l>(&'l self, other: &'l Self) -> SmallVec<[&'l Self::Key; 2]>;
14}
15
16pub trait KeyedAccess<Q: ?Sized> {
17    type Value;
18
19    fn get(&self, key: &Q) -> Option<&Self::Value>;
20    fn contains_key(&self, key: &Q) -> bool {
21        self.get(key).is_some()
22    }
23}
24
25impl<K: Eq + Hash, V: PartialEq, H: BuildHasher> KeyedEq for HashMap<K, V, H> {
26    type Key = K;
27
28    fn different_keys<'l>(&'l self, other: &'l Self) -> SmallVec<[&'l Self::Key; 2]> {
29        let mut different_keys = SmallVec::new();
30
31        for (key, value) in self.iter() {
32            if let Some(other_value) = HashMap::get(other, key) {
33                if value != other_value {
34                    different_keys.push(key);
35                }
36            } else {
37                different_keys.push(key);
38            }
39        }
40
41        if other.len() != self.len() || !different_keys.is_empty() {
42            for key in other.keys() {
43                if !self.contains_key(key) {
44                    different_keys.push(key);
45                }
46            }
47        }
48
49        different_keys
50    }
51}
52
53impl<K: Eq + Hash + Borrow<Q>, V, H: BuildHasher, Q: Eq + Hash + ?Sized> KeyedAccess<Q>
54    for HashMap<K, V, H>
55{
56    type Value = V;
57
58    fn get(&self, key: &Q) -> Option<&Self::Value> {
59        HashMap::get(self, key)
60    }
61
62    fn contains_key(&self, key: &Q) -> bool {
63        HashMap::contains_key(self, key)
64    }
65}
66
67impl<K: Eq + Hash, V: PartialEq, H: BuildHasher> KeyedEq for IndexMap<K, V, H> {
68    type Key = K;
69
70    fn different_keys<'l>(&'l self, other: &'l Self) -> SmallVec<[&'l Self::Key; 2]> {
71        let mut different_keys = SmallVec::new();
72
73        for (key, value) in self.iter() {
74            if let Some(other_value) = other.get(key) {
75                if value != other_value {
76                    different_keys.push(key);
77                }
78            } else {
79                different_keys.push(key);
80            }
81        }
82
83        if other.len() != self.len() || !different_keys.is_empty() {
84            for key in other.keys() {
85                if !self.contains_key(key) {
86                    different_keys.push(key);
87                }
88            }
89        }
90
91        different_keys
92    }
93}
94
95impl<K, V, H: BuildHasher, Q: Equivalent<K> + Hash + ?Sized> KeyedAccess<Q> for IndexMap<K, V, H> {
96    type Value = V;
97
98    fn get(&self, key: &Q) -> Option<&Self::Value> {
99        self.get(key)
100    }
101
102    fn contains_key(&self, key: &Q) -> bool {
103        self.contains_key(key)
104    }
105}
106
107impl<K: Eq + Hash, H: BuildHasher> KeyedEq for HashSet<K, H> {
108    type Key = K;
109
110    fn different_keys<'l>(&'l self, other: &'l Self) -> SmallVec<[&'l Self::Key; 2]> {
111        let mut different_keys = SmallVec::new();
112
113        for key in self.iter() {
114            if !other.contains(key) {
115                different_keys.push(key);
116            }
117        }
118
119        if other.len() != self.len() || !different_keys.is_empty() {
120            for key in other.iter() {
121                if !self.contains(key) {
122                    different_keys.push(key);
123                }
124            }
125        }
126
127        different_keys
128    }
129}
130
131impl<K: Eq + Hash + Borrow<Q>, Q: Eq + Hash + ?Sized, H: BuildHasher> KeyedAccess<Q>
132    for HashSet<K, H>
133{
134    type Value = ();
135
136    fn get(&self, key: &Q) -> Option<&Self::Value> {
137        if self.contains(key) { Some(&()) } else { None }
138    }
139
140    fn contains_key(&self, key: &Q) -> bool {
141        self.contains(key)
142    }
143}
144
145impl<K: Eq + Hash, H: BuildHasher> KeyedEq for IndexSet<K, H> {
146    type Key = K;
147
148    fn different_keys<'l>(&'l self, other: &'l Self) -> SmallVec<[&'l Self::Key; 2]> {
149        let mut different_keys = SmallVec::new();
150
151        for key in self.iter() {
152            if !other.contains(key) {
153                different_keys.push(key);
154            }
155        }
156
157        if other.len() != self.len() || !different_keys.is_empty() {
158            for key in other.iter() {
159                if !self.contains(key) {
160                    different_keys.push(key);
161                }
162            }
163        }
164
165        different_keys
166    }
167}
168
169impl<K, Q: Equivalent<K> + Hash + ?Sized, H: BuildHasher> KeyedAccess<Q> for IndexSet<K, H> {
170    type Value = ();
171
172    fn get(&self, key: &Q) -> Option<&Self::Value> {
173        if self.contains(key) { Some(&()) } else { None }
174    }
175
176    fn contains_key(&self, key: &Q) -> bool {
177        self.contains(key)
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    fn abc() -> [(&'static str, i32); 3] {
186        [("a", 1), ("b", 2), ("c", 3)]
187    }
188
189    fn abd() -> [(&'static str, i32); 3] {
190        [("a", 1), ("b", 22), ("d", 4)]
191    }
192
193    fn ab() -> [(&'static str, i32); 2] {
194        [("a", 1), ("b", 2)]
195    }
196
197    fn abcde() -> [(&'static str, i32); 5] {
198        [("a", 1), ("b", 2), ("c", 3), ("d", 4), ("e", 5)]
199    }
200
201    fn edcba() -> [(&'static str, i32); 5] {
202        [("e", 5), ("d", 4), ("c", 3), ("b", 2), ("a", 1)]
203    }
204
205    fn assert_diff<T: KeyedEq>(a: &T, b: &T, expected: &[&T::Key])
206    where
207        T::Key: std::fmt::Debug + PartialEq,
208    {
209        let diffs = a.different_keys(b);
210        assert_eq!(diffs.len(), expected.len());
211        for key in expected {
212            assert!(diffs.contains(key));
213        }
214    }
215
216    #[test]
217    fn test_hash_map_diff() {
218        let map1 = HashMap::from(abc());
219        let map2 = HashMap::from(abd());
220        assert_diff(&map1, &map2, &[&"b", &"c", &"d"]);
221    }
222
223    #[test]
224    fn test_index_map_diff() {
225        let map1 = IndexMap::from(abc());
226        let map2 = IndexMap::from(abd());
227        assert_diff(&map1, &map2, &[&"b", &"c", &"d"]);
228    }
229
230    #[test]
231    fn test_hash_map_equal() {
232        let map1 = HashMap::from(abcde());
233        let map2 = HashMap::from(edcba());
234        assert_diff(&map1, &map2, &[]);
235    }
236
237    #[test]
238    fn test_index_map_equal() {
239        let map1 = IndexMap::from(abcde());
240        let map2 = IndexMap::from(edcba());
241        assert_diff(&map1, &map2, &[]);
242    }
243
244    #[test]
245    fn test_hash_map_add_key() {
246        let map1 = HashMap::from(ab());
247        let map2 = HashMap::from(abc());
248        assert_diff(&map1, &map2, &[&"c"]);
249    }
250
251    #[test]
252    fn test_index_map_add_key() {
253        let map1 = IndexMap::from(ab());
254        let map2 = IndexMap::from(abc());
255        assert_diff(&map1, &map2, &[&"c"]);
256    }
257
258    #[test]
259    fn test_hash_map_remove_key() {
260        let map1 = HashMap::from(abc());
261        let map2 = HashMap::from(ab());
262        assert_diff(&map1, &map2, &[&"c"]);
263    }
264
265    #[test]
266    fn test_index_map_remove_key() {
267        let map1 = IndexMap::from(abc());
268        let map2 = IndexMap::from(ab());
269        assert_diff(&map1, &map2, &[&"c"]);
270    }
271}