toolshed/
map.rs

1//! Maps of keys to values that can be used with the `Arena`.
2
3use std::hash::{Hash, Hasher};
4use rustc_hash::FxHasher;
5
6use crate::cell::CopyCell;
7use crate::Arena;
8use crate::bloom::bloom;
9
10#[derive(Clone, Copy)]
11struct MapNode<'arena, K, V> {
12    pub key: K,
13    pub hash: u64,
14    pub value: CopyCell<V>,
15    pub left: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
16    pub right: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
17    pub next: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
18}
19
20impl<'arena, K, V> MapNode<'arena, K, V> {
21    pub const fn new(key: K, hash: u64, value: V) -> Self {
22        MapNode {
23            key,
24            hash,
25            value: CopyCell::new(value),
26            left: CopyCell::new(None),
27            right: CopyCell::new(None),
28            next: CopyCell::new(None),
29        }
30    }
31}
32
33/// A map of keys `K` to values `V`. The map is built as a pseudo-random
34/// binary tree with hashes of keys used for balancing the tree nodes.
35///
36/// All the nodes of the map are also linked to allow iteration in
37/// insertion order.
38#[derive(Clone, Copy)]
39pub struct Map<'arena, K, V> {
40    root: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
41    last: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
42}
43
44impl<'arena, K, V> Default for Map<'arena, K, V> {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl<'arena, K, V> Map<'arena, K, V> {
51    /// Create a new, empty `Map`.
52    pub const fn new() -> Self {
53        Map {
54            root: CopyCell::new(None),
55            last: CopyCell::new(None),
56        }
57    }
58}
59
60impl<'arena, K, V> Map<'arena, K, V> {
61    /// Get an iterator over key value pairs.
62    #[inline]
63    pub fn iter(&self) -> MapIter<'arena, K, V> {
64        MapIter {
65            next: self.root.get()
66        }
67    }
68
69    /// Returns true if the map contains no elements.
70    #[inline]
71    pub fn is_empty(&self) -> bool {
72        self.root.get().is_none()
73    }
74
75    /// Clears the map.
76    #[inline]
77    pub fn clear(&self) {
78        self.root.set(None);
79    }
80}
81
82impl<'arena, K, V> Map<'arena, K, V>
83where
84    K: Eq + Hash + Copy,
85    V: Copy,
86{
87    #[inline]
88    fn hash_key(key: &K) -> u64 {
89        let mut hasher = FxHasher::default();
90
91        key.hash(&mut hasher);
92
93        hasher.finish()
94    }
95
96    #[inline]
97    fn find_slot(&self, key: K, hash: u64) -> &CopyCell<Option<&'arena MapNode<'arena, K, V>>> {
98        let mut node = &self.root;
99
100        loop {
101            match node.get() {
102                None         => return node,
103                Some(parent) => {
104                    if hash == parent.hash && key == parent.key {
105                        return node;
106                    } else if hash < parent.hash {
107                        node = &parent.left;
108                    } else {
109                        node = &parent.right;
110                    }
111                }
112            }
113        }
114    }
115
116    /// Inserts a key-value pair into the map. If the key was previously set,
117    /// old value is returned.
118    #[inline]
119    pub fn insert(&self, arena: &'arena Arena, key: K, value: V) -> Option<V> {
120        let hash = Self::hash_key(&key);
121        let node = self.find_slot(key, hash);
122
123        match node.get() {
124            Some(node) => {
125                let old = node.value.get();
126                node.value.set(value);
127                Some(old)
128            },
129            None => {
130                let new = Some(&*arena.alloc(MapNode::new(key, hash, value)));
131
132                if let Some(last) = self.last.get() {
133                    last.next.set(new);
134                }
135
136                self.last.set(new);
137                node.set(new);
138                None
139            }
140        }
141    }
142
143    /// Returns the value corresponding to the key.
144    #[inline]
145    pub fn get_key(&self, key: K) -> Option<&K> {
146        let hash = Self::hash_key(&key);
147
148        self.find_slot(key, hash).get().map(|node| &node.key)
149    }
150
151    /// Returns the value corresponding to the key.
152    #[inline]
153    pub fn get(&self, key: K) -> Option<V> {
154        let hash = Self::hash_key(&key);
155
156        self.find_slot(key, hash).get().map(|node| node.value.get())
157    }
158
159    /// Returns true if the map contains a value for the specified key.
160    #[inline]
161    pub fn contains_key(&self, key: K) -> bool {
162        let hash = Self::hash_key(&key);
163
164        self.find_slot(key, hash).get().is_some()
165    }
166}
167
168/// A variant of the `Map` that includes a bloom filter using the
169/// `bloom` function for keys that can be represented as byte slices.
170///
171/// This is ideal for small maps for which querying for absent keys is
172/// a common behavior. In this case it will very likely outperform a
173/// `HashMap`, even one with a fast hashing algorithm.
174#[derive(Clone, Copy)]
175pub struct BloomMap<'arena, K, V> {
176    filter: CopyCell<u64>,
177    inner: Map<'arena, K, V>,
178}
179
180impl<'arena, K, V> BloomMap<'arena, K, V> {
181    /// Create a new, empty `BloomMap`.
182    pub const fn new() -> Self {
183        BloomMap {
184            filter: CopyCell::new(0),
185            inner: Map::new(),
186        }
187    }
188}
189
190impl<'arena, K, V: Copy> BloomMap<'arena, K, V> {
191    /// Get an iterator over key value pairs.
192    #[inline]
193    pub fn iter(&self) -> MapIter<'arena, K, V> {
194        self.inner.iter()
195    }
196
197    /// Returns true if the map contains no elements.
198    #[inline]
199    pub fn is_empty(&self) -> bool {
200        self.inner.is_empty()
201    }
202
203    /// Clears the map.
204    #[inline]
205    pub fn clear(&self) {
206        self.filter.set(0);
207        self.inner.clear();
208    }
209}
210
211impl<'arena, K, V> BloomMap<'arena, K, V>
212where
213    K: Eq + Hash + Copy + AsRef<[u8]>,
214    V: Copy,
215{
216    /// Inserts a key-value pair into the map. If the key was previously set,
217    /// old value is returned.
218    #[inline]
219    pub fn insert(&self, arena: &'arena Arena, key: K, value: V) -> Option<V> {
220        self.filter.set(self.filter.get() | bloom(key));
221        self.inner.insert(arena, key, value)
222    }
223
224    /// Returns the value corresponding to the key.
225    #[inline]
226    pub fn get(&self, key: K) -> Option<V> {
227        let b = bloom(key.as_ref());
228
229        if self.filter.get() & b == b {
230            self.inner.get(key)
231        } else {
232            None
233        }
234    }
235
236    /// Returns true if the map contains a value for the specified key.
237    #[inline]
238    pub fn contains_key(&self, key: K) -> bool {
239        let b = bloom(key);
240
241        self.filter.get() & b == b && self.inner.contains_key(key)
242    }
243}
244
245/// An iterator over the entries in the map.
246/// All entries are returned in insertion order.
247pub struct MapIter<'arena, K, V> {
248    next: Option<&'arena MapNode<'arena, K, V>>
249}
250
251impl<'arena, K, V: Copy> Iterator for MapIter<'arena, K, V> {
252    type Item = (&'arena K, V);
253
254    #[inline]
255    fn next(&mut self) -> Option<Self::Item> {
256        let next = self.next;
257
258        next.map(|map_node| {
259            let item = (&map_node.key, map_node.value.get());
260            self.next = map_node.next.get();
261            item
262        })
263    }
264}
265
266impl<'arena, K, V: Copy> IntoIterator for Map<'arena, K, V> {
267    type Item = (&'arena K, V);
268    type IntoIter = MapIter<'arena, K, V>;
269
270    #[inline]
271    fn into_iter(self) -> Self::IntoIter {
272        self.iter()
273    }
274}
275
276impl<'arena, K, V: Copy> IntoIterator for BloomMap<'arena, K, V> {
277    type Item = (&'arena K, V);
278    type IntoIter = MapIter<'arena, K, V>;
279
280    #[inline]
281    fn into_iter(self) -> Self::IntoIter {
282        self.iter()
283    }
284}
285
286impl<'arena, K, V> From<Map<'arena, K, V>> for BloomMap<'arena, K, V>
287where
288    K: Eq + Hash + Copy + AsRef<[u8]>,
289    V: Copy,
290{
291    fn from(map: Map<'arena, K, V>) -> BloomMap<'arena, K, V> {
292        let mut filter = 0;
293
294        for (key, _) in map.iter() {
295            filter |= bloom(key.as_ref());
296        }
297
298        BloomMap {
299            filter: CopyCell::new(filter),
300            inner: map,
301        }
302    }
303}
304
305impl<'arena, K, V> From<BloomMap<'arena, K, V>> for Map<'arena, K, V> {
306    #[inline]
307    fn from(bloom_map: BloomMap<'arena, K, V>) -> Map<'arena, K, V> {
308        bloom_map.inner
309    }
310}
311
312#[cfg(test)]
313mod test {
314    use super::*;
315
316    #[test]
317    fn map() {
318        let arena = Arena::new();
319        let map = Map::new();
320
321        map.insert(&arena, "foo", 10u64);
322        map.insert(&arena, "bar", 20);
323        map.insert(&arena, "doge", 30);
324
325        assert_eq!(map.contains_key("foo"), true);
326        assert_eq!(map.contains_key("bar"), true);
327        assert_eq!(map.contains_key("doge"), true);
328        assert_eq!(map.contains_key("moon"), false);
329
330        assert_eq!(map.get("foo"), Some(10));
331        assert_eq!(map.get("bar"), Some(20));
332        assert_eq!(map.get("doge"), Some(30));
333        assert_eq!(map.get("moon"), None);
334    }
335
336    #[test]
337    fn bloom_map() {
338        let arena = Arena::new();
339        let map = BloomMap::new();
340
341        map.insert(&arena, "foo", 10u64);
342        map.insert(&arena, "bar", 20);
343        map.insert(&arena, "doge", 30);
344
345        assert_eq!(map.contains_key("foo"), true);
346        assert_eq!(map.contains_key("bar"), true);
347        assert_eq!(map.contains_key("doge"), true);
348        assert_eq!(map.contains_key("moon"), false);
349
350        assert_eq!(map.get("foo"), Some(10));
351        assert_eq!(map.get("bar"), Some(20));
352        assert_eq!(map.get("doge"), Some(30));
353        assert_eq!(map.get("moon"), None);
354    }
355
356    #[test]
357    fn iter() {
358        let arena = Arena::new();
359        let map = Map::new();
360
361        map.insert(&arena, "foo", 10u64);
362        map.insert(&arena, "bar", 20);
363        map.insert(&arena, "doge", 30);
364
365        let mut iter = map.iter();
366
367        assert_eq!(iter.next(), Some((&"foo", 10)));
368        assert_eq!(iter.next(), Some((&"bar", 20)));
369        assert_eq!(iter.next(), Some((&"doge", 30)));
370        assert_eq!(iter.next(), None);
371    }
372
373    #[test]
374    fn insert_replace() {
375        let arena = Arena::new();
376        let map = Map::new();
377
378        map.insert(&arena, "foo", 10u64);
379        map.insert(&arena, "bar", 20);
380        map.insert(&arena, "doge", 30);
381
382        let mut iter = map.iter();
383
384        assert_eq!(iter.next(), Some((&"foo", 10)));
385        assert_eq!(iter.next(), Some((&"bar", 20)));
386        assert_eq!(iter.next(), Some((&"doge", 30)));
387        assert_eq!(iter.next(), None);
388
389        map.insert(&arena, "bar", 42);
390
391        let mut iter = map.iter();
392
393        assert_eq!(iter.next(), Some((&"foo", 10)));
394        assert_eq!(iter.next(), Some((&"bar", 42)));
395        assert_eq!(iter.next(), Some((&"doge", 30)));
396        assert_eq!(iter.next(), None);
397    }
398
399    #[test]
400    fn from_eq() {
401        let arena = Arena::new();
402        let map = Map::new();
403
404        map.insert(&arena, "foo", 10);
405        map.insert(&arena, "bar", 20);
406        map.insert(&arena, "doge", 30);
407
408        let bloom_map = BloomMap::new();
409
410        bloom_map.insert(&arena, "foo", 10);
411        bloom_map.insert(&arena, "bar", 20);
412        bloom_map.insert(&arena, "doge", 30);
413
414        assert_eq!(map, Map::from(bloom_map));
415        assert_eq!(BloomMap::from(map), bloom_map);
416    }
417}
toolshed/map.rs

toolshed/
map.rs
