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

style/rule_tree/
map.rs
