1use crate::symbol::Symbol;
2use crate::variable::Variable;
3use ahash::{HashMap, HashMapExt};
4use smallvec::SmallVec;
5use std::fmt::{Debug, Formatter};
6
7const INLINE: usize = 8;
8
9const SPILL_AT: usize = 32;
10
11type Entries = SmallVec<[(Symbol, Variable); INLINE]>;
12
13#[derive(Clone)]
14enum Repr {
15 Small(Entries),
16 Large(HashMap<Symbol, Variable>),
17}
18
19#[derive(Clone)]
20pub struct VariableMap(Repr);
21
22impl VariableMap {
23 pub fn new() -> Self {
24 Self(Repr::Small(SmallVec::new()))
25 }
26
27 pub fn with_capacity(capacity: usize) -> Self {
28 match capacity > SPILL_AT {
29 true => Self(Repr::Large(HashMap::with_capacity(capacity))),
30 false => Self(Repr::Small(SmallVec::with_capacity(capacity))),
31 }
32 }
33
34 pub fn len(&self) -> usize {
35 match &self.0 {
36 Repr::Small(entries) => entries.len(),
37 Repr::Large(map) => map.len(),
38 }
39 }
40
41 pub fn is_empty(&self) -> bool {
42 self.len() == 0
43 }
44
45 pub fn clear(&mut self) {
46 match &mut self.0 {
47 Repr::Small(entries) => entries.clear(),
48 Repr::Large(map) => map.clear(),
49 }
50 }
51
52 #[inline]
53 pub fn get(&self, key: &Symbol) -> Option<&Variable> {
54 match &self.0 {
55 Repr::Small(entries) => entries
56 .iter()
57 .find(|(k, _)| k.as_str() == key.as_str())
58 .map(|(_, v)| v),
59 Repr::Large(map) => map.get(key),
60 }
61 }
62
63 #[inline]
64 pub fn get_str(&self, key: &str) -> Option<&Variable> {
65 match &self.0 {
66 Repr::Small(entries) => entries
67 .iter()
68 .find(|(k, _)| k.as_str() == key)
69 .map(|(_, v)| v),
70 Repr::Large(map) => map.get(key),
71 }
72 }
73
74 pub fn get_mut(&mut self, key: &Symbol) -> Option<&mut Variable> {
75 match &mut self.0 {
76 Repr::Small(entries) => entries
77 .iter_mut()
78 .find(|(k, _)| k.as_str() == key.as_str())
79 .map(|(_, v)| v),
80 Repr::Large(map) => map.get_mut(key),
81 }
82 }
83
84 pub fn get_key_value(&self, key: &Symbol) -> Option<(&Symbol, &Variable)> {
85 match &self.0 {
86 Repr::Small(entries) => entries
87 .iter()
88 .find(|(k, _)| k.as_str() == key.as_str())
89 .map(|(k, v)| (k, v)),
90 Repr::Large(map) => map.get_key_value(key),
91 }
92 }
93
94 pub fn contains_key(&self, key: &Symbol) -> bool {
95 self.get(key).is_some()
96 }
97
98 pub fn contains_key_str(&self, key: &str) -> bool {
99 self.get_str(key).is_some()
100 }
101
102 pub fn remove_str(&mut self, key: &str) -> Option<Variable> {
103 match &mut self.0 {
104 Repr::Small(entries) => entries
105 .iter()
106 .position(|(k, _)| k.as_str() == key)
107 .map(|index| entries.remove(index).1),
108 Repr::Large(map) => map.remove(key),
109 }
110 }
111
112 pub fn insert(&mut self, key: Symbol, value: Variable) -> Option<Variable> {
113 match &mut self.0 {
114 Repr::Small(entries) => {
115 if let Some(slot) = entries.iter_mut().find(|(k, _)| k.as_str() == key.as_str()) {
116 return Some(std::mem::replace(&mut slot.1, value));
117 }
118 if entries.len() >= SPILL_AT {
119 self.spill();
120 let Repr::Large(map) = &mut self.0 else {
121 unreachable!("just spilled")
122 };
123 return map.insert(key, value);
124 }
125 entries.push((key, value));
126 None
127 }
128 Repr::Large(map) => map.insert(key, value),
129 }
130 }
131
132 pub fn remove(&mut self, key: &Symbol) -> Option<Variable> {
133 match &mut self.0 {
134 Repr::Small(entries) => entries
135 .iter()
136 .position(|(k, _)| k.as_str() == key.as_str())
137 .map(|index| entries.remove(index).1),
138 Repr::Large(map) => map.remove(key),
139 }
140 }
141
142 fn spill(&mut self) {
143 let Repr::Small(entries) = &mut self.0 else {
144 return;
145 };
146 let mut map = HashMap::with_capacity(entries.len() * 2);
147 for (key, value) in entries.drain(..) {
148 map.insert(key, value);
149 }
150 self.0 = Repr::Large(map);
151 }
152
153 pub fn entry(&mut self, key: Symbol) -> Entry<'_> {
154 if matches!(&self.0, Repr::Small(entries)
155 if entries.len() >= SPILL_AT && self.get(&key).is_none())
156 {
157 self.spill();
158 }
159
160 match self.contains_key(&key) {
161 true => Entry::Occupied(OccupiedEntry { map: self, key }),
162 false => Entry::Vacant(VacantEntry { map: self, key }),
163 }
164 }
165
166 pub fn iter(&self) -> Iter<'_> {
167 match &self.0 {
168 Repr::Small(entries) => Iter::Small(entries.iter()),
169 Repr::Large(map) => Iter::Large(map.iter()),
170 }
171 }
172
173 pub fn iter_mut(&mut self) -> IterMut<'_> {
174 match &mut self.0 {
175 Repr::Small(entries) => IterMut::Small(entries.iter_mut()),
176 Repr::Large(map) => IterMut::Large(map.iter_mut()),
177 }
178 }
179
180 pub fn keys(&self) -> impl Iterator<Item = &Symbol> + '_ {
181 self.iter().map(|(key, _)| key)
182 }
183
184 pub fn values(&self) -> impl Iterator<Item = &Variable> {
185 self.iter().map(|(_, value)| value)
186 }
187
188 pub fn values_mut(&mut self) -> impl Iterator<Item = &mut Variable> {
189 self.iter_mut().map(|(_, value)| value)
190 }
191}
192
193pub enum Entry<'a> {
194 Occupied(OccupiedEntry<'a>),
195 Vacant(VacantEntry<'a>),
196}
197
198pub struct OccupiedEntry<'a> {
199 map: &'a mut VariableMap,
200 key: Symbol,
201}
202
203pub struct VacantEntry<'a> {
204 map: &'a mut VariableMap,
205 key: Symbol,
206}
207
208impl<'a> Entry<'a> {
209 pub fn or_insert(self, default: Variable) -> &'a mut Variable {
210 match self {
211 Entry::Occupied(entry) => entry.into_mut(),
212 Entry::Vacant(entry) => entry.insert(default),
213 }
214 }
215
216 pub fn or_insert_with<F: FnOnce() -> Variable>(self, default: F) -> &'a mut Variable {
217 match self {
218 Entry::Occupied(entry) => entry.into_mut(),
219 Entry::Vacant(entry) => entry.insert(default()),
220 }
221 }
222}
223
224impl<'a> OccupiedEntry<'a> {
225 pub fn get(&self) -> &Variable {
226 self.map.get(&self.key).expect("occupied")
227 }
228
229 pub fn get_mut(&mut self) -> &mut Variable {
230 self.map.get_mut(&self.key).expect("occupied")
231 }
232
233 pub fn into_mut(self) -> &'a mut Variable {
234 let key = self.key;
235 self.map.get_mut(&key).expect("occupied")
236 }
237
238 pub fn insert(&mut self, value: Variable) -> Variable {
239 std::mem::replace(self.get_mut(), value)
240 }
241}
242
243impl<'a> VacantEntry<'a> {
244 pub fn insert(self, value: Variable) -> &'a mut Variable {
245 let key = self.key;
246 self.map.insert(key.clone(), value);
247 self.map.get_mut(&key).expect("just inserted")
248 }
249}
250
251pub enum Iter<'a> {
252 Small(std::slice::Iter<'a, (Symbol, Variable)>),
253 Large(std::collections::hash_map::Iter<'a, Symbol, Variable>),
254}
255
256impl<'a> Iterator for Iter<'a> {
257 type Item = (&'a Symbol, &'a Variable);
258
259 fn next(&mut self) -> Option<Self::Item> {
260 match self {
261 Iter::Small(iter) => iter.next().map(|(key, value)| (key, value)),
262 Iter::Large(iter) => iter.next(),
263 }
264 }
265
266 fn size_hint(&self) -> (usize, Option<usize>) {
267 match self {
268 Iter::Small(iter) => iter.size_hint(),
269 Iter::Large(iter) => iter.size_hint(),
270 }
271 }
272}
273
274pub enum IterMut<'a> {
275 Small(std::slice::IterMut<'a, (Symbol, Variable)>),
276 Large(std::collections::hash_map::IterMut<'a, Symbol, Variable>),
277}
278
279impl<'a> Iterator for IterMut<'a> {
280 type Item = (&'a Symbol, &'a mut Variable);
281
282 fn next(&mut self) -> Option<Self::Item> {
283 match self {
284 IterMut::Small(iter) => iter.next().map(|(key, value)| (&*key, value)),
285 IterMut::Large(iter) => iter.next(),
286 }
287 }
288}
289
290pub enum IntoIter {
291 Small(smallvec::IntoIter<[(Symbol, Variable); INLINE]>),
292 Large(std::collections::hash_map::IntoIter<Symbol, Variable>),
293}
294
295impl Iterator for IntoIter {
296 type Item = (Symbol, Variable);
297
298 fn next(&mut self) -> Option<Self::Item> {
299 match self {
300 IntoIter::Small(iter) => iter.next(),
301 IntoIter::Large(iter) => iter.next(),
302 }
303 }
304}
305
306impl IntoIterator for VariableMap {
307 type Item = (Symbol, Variable);
308 type IntoIter = IntoIter;
309
310 fn into_iter(self) -> IntoIter {
311 match self.0 {
312 Repr::Small(entries) => IntoIter::Small(entries.into_iter()),
313 Repr::Large(map) => IntoIter::Large(map.into_iter()),
314 }
315 }
316}
317
318impl<'a> IntoIterator for &'a VariableMap {
319 type Item = (&'a Symbol, &'a Variable);
320 type IntoIter = Iter<'a>;
321
322 fn into_iter(self) -> Iter<'a> {
323 self.iter()
324 }
325}
326
327impl Default for VariableMap {
328 fn default() -> Self {
329 Self::new()
330 }
331}
332
333impl FromIterator<(Symbol, Variable)> for VariableMap {
334 fn from_iter<T: IntoIterator<Item = (Symbol, Variable)>>(iter: T) -> Self {
335 let iter = iter.into_iter();
336 let mut map = VariableMap::with_capacity(iter.size_hint().0);
337 for (key, value) in iter {
338 map.insert(key, value);
339 }
340 map
341 }
342}
343
344impl Extend<(Symbol, Variable)> for VariableMap {
345 fn extend<T: IntoIterator<Item = (Symbol, Variable)>>(&mut self, iter: T) {
346 for (key, value) in iter {
347 self.insert(key, value);
348 }
349 }
350}
351
352impl PartialEq for VariableMap {
353 fn eq(&self, other: &Self) -> bool {
354 self.len() == other.len()
355 && self
356 .iter()
357 .all(|(key, value)| other.get(key).is_some_and(|o| o == value))
358 }
359}
360
361impl VariableMap {
362 pub fn insert_str(&mut self, key: &str, value: Variable) -> Option<Variable> {
363 self.insert(Symbol::from(key), value)
364 }
365}
366
367impl Debug for VariableMap {
368 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
369 f.debug_map()
370 .entries(self.iter().map(|(k, v)| (k.as_str(), v)))
371 .finish()
372 }
373}
374
375impl serde::Serialize for VariableMap {
376 fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
377 use serde::ser::SerializeMap;
378
379 let mut map = serializer.serialize_map(Some(self.len()))?;
380 for (key, value) in self.iter() {
381 map.serialize_entry(key.as_str(), value)?;
382 }
383 map.end()
384 }
385}