style/rule_tree/
map.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
4
5#![forbid(unsafe_code)]
6
7use fxhash::FxHashMap;
8use malloc_size_of::{MallocShallowSizeOf, MallocSizeOfOps};
9use std::collections::hash_map;
10use std::hash::Hash;
11use std::mem;
12
13pub(super) struct Map<K, V> {
14    inner: MapInner<K, V>,
15}
16
17enum MapInner<K, V> {
18    Empty,
19    One(V),
20    Map(Box<FxHashMap<K, V>>),
21}
22
23pub(super) struct MapIter<'a, K, V> {
24    inner: MapIterInner<'a, K, V>,
25}
26
27enum MapIterInner<'a, K, V> {
28    One(std::option::IntoIter<&'a V>),
29    Map(std::collections::hash_map::Values<'a, K, V>),
30}
31
32pub(super) enum Entry<'a, K, V> {
33    Occupied(&'a mut V),
34    Vacant(VacantEntry<'a, K, V>),
35}
36
37pub(super) struct VacantEntry<'a, K, V> {
38    inner: VacantEntryInner<'a, K, V>,
39}
40
41enum VacantEntryInner<'a, K, V> {
42    One(&'a mut MapInner<K, V>),
43    Map(hash_map::VacantEntry<'a, K, V>),
44}
45
46impl<K, V> Default for Map<K, V> {
47    fn default() -> Self {
48        Map {
49            inner: MapInner::Empty,
50        }
51    }
52}
53
54impl<'a, K, V> IntoIterator for &'a Map<K, V> {
55    type Item = &'a V;
56    type IntoIter = MapIter<'a, K, V>;
57
58    fn into_iter(self) -> Self::IntoIter {
59        MapIter {
60            inner: match &self.inner {
61                MapInner::Empty => MapIterInner::One(None.into_iter()),
62                MapInner::One(one) => MapIterInner::One(Some(one).into_iter()),
63                MapInner::Map(map) => MapIterInner::Map(map.values()),
64            },
65        }
66    }
67}
68
69impl<'a, K, V> Iterator for MapIter<'a, K, V> {
70    type Item = &'a V;
71
72    fn next(&mut self) -> Option<Self::Item> {
73        match &mut self.inner {
74            MapIterInner::One(one_iter) => one_iter.next(),
75            MapIterInner::Map(map_iter) => map_iter.next(),
76        }
77    }
78}
79
80impl<K, V> Map<K, V>
81where
82    K: Eq + Hash,
83{
84    pub(super) fn is_empty(&self) -> bool {
85        match &self.inner {
86            MapInner::Empty => true,
87            MapInner::One(_) => false,
88            MapInner::Map(map) => map.is_empty(),
89        }
90    }
91
92    #[cfg(debug_assertions)]
93    pub(super) fn len(&self) -> usize {
94        match &self.inner {
95            MapInner::Empty => 0,
96            MapInner::One(_) => 1,
97            MapInner::Map(map) => map.len(),
98        }
99    }
100
101    pub(super) fn get(&self, key: &K, key_from_value: impl FnOnce(&V) -> K) -> Option<&V> {
102        match &self.inner {
103            MapInner::One(one) if *key == key_from_value(one) => Some(one),
104            MapInner::Map(map) => map.get(key),
105            MapInner::Empty | MapInner::One(_) => None,
106        }
107    }
108
109    pub(super) fn entry(
110        &mut self,
111        key: K,
112        key_from_value: impl FnOnce(&V) -> K,
113    ) -> Entry<'_, K, V> {
114        match self.inner {
115            ref mut inner @ MapInner::Empty => Entry::Vacant(VacantEntry {
116                inner: VacantEntryInner::One(inner),
117            }),
118            MapInner::One(_) => {
119                let one = match mem::replace(&mut self.inner, MapInner::Empty) {
120                    MapInner::One(one) => one,
121                    _ => unreachable!(),
122                };
123                // If this panics, the child `one` will be lost.
124                let one_key = key_from_value(&one);
125                // Same for the equality test.
126                if key == one_key {
127                    self.inner = MapInner::One(one);
128                    let one = match &mut self.inner {
129                        MapInner::One(one) => one,
130                        _ => unreachable!(),
131                    };
132                    return Entry::Occupied(one);
133                }
134                self.inner = MapInner::Map(Box::new(FxHashMap::with_capacity_and_hasher(
135                    2,
136                    Default::default(),
137                )));
138                let map = match &mut self.inner {
139                    MapInner::Map(map) => map,
140                    _ => unreachable!(),
141                };
142                map.insert(one_key, one);
143                match map.entry(key) {
144                    hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry {
145                        inner: VacantEntryInner::Map(entry),
146                    }),
147                    _ => unreachable!(),
148                }
149            },
150            MapInner::Map(ref mut map) => match map.entry(key) {
151                hash_map::Entry::Occupied(entry) => Entry::Occupied(entry.into_mut()),
152                hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry {
153                    inner: VacantEntryInner::Map(entry),
154                }),
155            },
156        }
157    }
158
159    pub(super) fn remove(&mut self, key: &K, key_from_value: impl FnOnce(&V) -> K) -> Option<V> {
160        match &mut self.inner {
161            MapInner::One(one) if *key == key_from_value(one) => {
162                match mem::replace(&mut self.inner, MapInner::Empty) {
163                    MapInner::One(one) => Some(one),
164                    _ => unreachable!(),
165                }
166            },
167            MapInner::Map(map) => map.remove(key),
168            MapInner::Empty | MapInner::One(_) => None,
169        }
170    }
171}
172
173impl<'a, K, V> VacantEntry<'a, K, V> {
174    pub(super) fn insert(self, value: V) -> &'a mut V {
175        match self.inner {
176            VacantEntryInner::One(map) => {
177                *map = MapInner::One(value);
178                match map {
179                    MapInner::One(one) => one,
180                    _ => unreachable!(),
181                }
182            },
183            VacantEntryInner::Map(entry) => entry.insert(value),
184        }
185    }
186}
187
188impl<K, V> MallocShallowSizeOf for Map<K, V>
189where
190    K: Eq + Hash,
191{
192    fn shallow_size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
193        match &self.inner {
194            MapInner::Map(m) => {
195                // We want to account for both the box and the hashmap.
196                m.shallow_size_of(ops) + (**m).shallow_size_of(ops)
197            },
198            MapInner::One(_) | MapInner::Empty => 0,
199        }
200    }
201}
style/rule_tree/map.rs

style/rule_tree/
map.rs
