exhaustive_map/
map.rs

1use std::{
2    borrow::Borrow,
3    collections::{BTreeMap, HashMap},
4    fmt::Debug,
5    hash::{BuildHasher, Hash},
6    marker::PhantomData,
7    mem::MaybeUninit,
8    ops::{Index, IndexMut},
9};
10
11use crate::{
12    finite::{Finite, FiniteExt},
13    IterAll,
14};
15
16/// A map which is guaranteed to always contain a value for each possible key of type `K`.
17/// ```
18/// use exhaustive_map::ExhaustiveMap;
19///
20/// let mut map = ExhaustiveMap::<u8, u16>::from_fn(|i| i as u16 + 100);
21/// assert_eq!(map.len(), 256);
22///
23/// assert_eq!(map[3], 103);
24///
25/// map[7] = 9999;
26/// assert_eq!(map[7], 9999);
27///
28/// map.swap(7, 3);
29/// assert_eq!(map[3], 9999);
30/// assert_eq!(map[7], 103);
31/// ```
32#[repr(transparent)]
33pub struct ExhaustiveMap<K: Finite, V> {
34    // Replace with [V; { K::INHABITANTS }] when Rust supports it
35    array: Box<[V]>,
36    _phantom: PhantomData<K>,
37}
38
39impl<K: Finite, V> ExhaustiveMap<K, V> {
40    /// Creates a map by providing a mapping function from `K` to `V`.
41    ///
42    /// Similar to [`array::from_fn`](std::array::from_fn).
43    #[must_use]
44    pub fn from_fn(f: impl FnMut(K) -> V) -> Self {
45        Self {
46            array: K::iter_all().map(f).collect(),
47            _phantom: PhantomData,
48        }
49    }
50
51    /// Tries to create a map by providing a mapping function from `K` to `Result<V, E>`.
52    ///
53    /// # Errors
54    ///
55    /// Returns the first error if any of the mappings fails.
56    pub fn try_from_fn<E>(f: impl FnMut(K) -> Result<V, E>) -> Result<Self, E> {
57        Ok(Self {
58            array: K::iter_all().map(f).collect::<Result<_, E>>()?,
59            _phantom: PhantomData,
60        })
61    }
62
63    /// Creates a map by providing a mapping function from `usize` to `V`.
64    /// The map is filled according to the [`Finite`] implementation of `K`.
65    ///
66    /// ```
67    /// use exhaustive_map::{ExhaustiveMap, Finite};
68    ///
69    /// #[derive(Finite, Debug)]
70    /// enum Color {
71    ///     Red,
72    ///     Green,
73    ///     Blue,
74    /// }
75    ///
76    /// let map = ExhaustiveMap::from_usize_fn(|i| i);
77    /// assert_eq!(map[Color::Red], 0);
78    /// assert_eq!(map[Color::Green], 1);
79    /// assert_eq!(map[Color::Blue], 2);
80    /// ```
81    #[must_use]
82    pub fn from_usize_fn(f: impl FnMut(usize) -> V) -> Self {
83        Self {
84            array: (0..K::INHABITANTS).map(f).collect(),
85            _phantom: PhantomData,
86        }
87    }
88
89    /// Returns the number of elements in the map.
90    ///
91    /// Always equal to `K::INHABITANTS`.
92    #[must_use]
93    pub const fn len(&self) -> usize {
94        K::INHABITANTS
95    }
96
97    /// Returns `true` if the map contains no elements.
98    ///
99    /// The map can only be empty if `K::INHABITANTS` is zero,
100    /// meaning the type `K` is uninhabitable.
101    #[must_use]
102    pub const fn is_empty(&self) -> bool {
103        K::INHABITANTS == 0
104    }
105
106    /// Replace the value stored for `k` with `v`, returning the previous stored value.
107    pub fn replace<Q: Borrow<K>>(&mut self, k: Q, v: V) -> V {
108        std::mem::replace(&mut self[k], v)
109    }
110
111    /// Swaps the values at stored at `k1` and `k2`.
112    pub fn swap<Q1: Borrow<K>, Q2: Borrow<K>>(&mut self, k1: Q1, k2: Q2) {
113        self.array
114            .swap(k1.borrow().to_usize(), k2.borrow().to_usize());
115    }
116
117    /// Replace the value stored for `k` with the default value of `V`, returning the previous stored value.
118    pub fn take<Q: Borrow<K>>(&mut self, k: Q) -> V
119    where
120        V: Default,
121    {
122        std::mem::take(&mut self[k])
123    }
124
125    /// Change the values of the stored values via a mapping function.
126    ///
127    /// ```
128    /// use exhaustive_map::ExhaustiveMap;
129    ///
130    /// let bool_to_int = ExhaustiveMap::from_fn(|k| if k { 1 } else { 0 });
131    /// let bool_to_int_string = bool_to_int.map_values(|v| v.to_string());
132    ///
133    /// assert_eq!(bool_to_int_string[false], "0");
134    /// assert_eq!(bool_to_int_string[true], "1");
135    /// ```
136    #[must_use]
137    pub fn map_values<U>(self, f: impl FnMut(V) -> U) -> ExhaustiveMap<K, U> {
138        ExhaustiveMap {
139            array: self.into_values().map(f).collect(),
140            _phantom: PhantomData,
141        }
142    }
143
144    /// An iterator visiting all keys in the order provided by [`Finite`].
145    ///
146    /// This creates new keys by calling [`K::from_usize`](Finite::from_usize) for each key.
147    pub fn keys() -> IterAll<K> {
148        K::iter_all()
149    }
150
151    /// An iterator visiting all values stored in the map, ordered by the keys order provided by [`Finite`].
152    pub fn values(&self) -> Values<'_, V> {
153        Values(self.array.iter())
154    }
155
156    /// A mutable iterator visiting all values stored in the map, ordered by the keys order provided by [`Finite`].
157    pub fn values_mut(&mut self) -> ValuesMut<'_, V> {
158        ValuesMut(self.array.iter_mut())
159    }
160
161    /// Creates a consuming iterator visiting all the values, ordered by the keys order provided by [`Finite`].
162    /// The map cannot be used after calling this.
163    pub fn into_values(self) -> IntoValues<V> {
164        IntoValues(self.array.into_vec().into_iter())
165    }
166
167    /// An iterator visiting all entries stored in the map, ordered by the keys order provided by [`Finite`].
168    ///
169    /// This creates new keys by calling [`K::from_usize`](Finite::from_usize) for each key.
170    pub fn iter(&self) -> Iter<'_, K, V> {
171        Iter(Self::keys().zip(self.values()))
172    }
173
174    /// A mutable iterator visiting all entries stored in the map, ordered by the keys order provided by [`Finite`].
175    ///
176    /// This creates new keys by calling [`K::from_usize`](Finite::from_usize) for each key.
177    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
178        IterMut(Self::keys().zip(self.values_mut()))
179    }
180
181    /// Creates a map with [`MaybeUninit`] values.
182    ///
183    /// After every value have been initialized [`assume_init`](ExhaustiveMap::assume_init) can be
184    /// called to obtain a map with values of type `V`.
185    #[must_use]
186    pub fn new_uninit() -> ExhaustiveMap<K, MaybeUninit<V>> {
187        ExhaustiveMap::from_usize_fn(|_| MaybeUninit::uninit())
188    }
189}
190
191impl<K: Finite, V> ExhaustiveMap<K, Option<V>> {
192    /// Tries to convert an `ExhaustiveMap<K, Option<V>>` to an `ExhaustiveMap<K, V>`.
193    ///
194    /// # Errors
195    ///
196    /// If any of the values are `None`, this returns `Err` containing the input map.
197    pub fn try_unwrap_values(self) -> Result<ExhaustiveMap<K, V>, ExhaustiveMap<K, Option<V>>> {
198        if !self.array.iter().all(Option::is_some) {
199            return Err(self);
200        }
201        #[allow(clippy::missing_panics_doc)]
202        let values: Box<[V]> = self
203            .array
204            .into_vec()
205            .into_iter()
206            .map(|v| v.unwrap())
207            .collect();
208        // SAFETY: `values` has the correct length as we used `map`.
209        Ok(unsafe { values.try_into().unwrap_unchecked() })
210    }
211}
212
213impl<K: Finite, V> ExhaustiveMap<K, MaybeUninit<V>> {
214    /// # Safety
215    ///
216    /// All elements must have been initialized.
217    #[must_use]
218    pub unsafe fn assume_init(self) -> ExhaustiveMap<K, V> {
219        ExhaustiveMap {
220            array: std::mem::transmute::<Box<[MaybeUninit<V>]>, Box<[V]>>(self.array),
221            _phantom: PhantomData,
222        }
223    }
224}
225
226impl<K: Finite, V> TryFrom<Box<[V]>> for ExhaustiveMap<K, V> {
227    type Error = Box<[V]>;
228
229    fn try_from(value: Box<[V]>) -> Result<Self, Self::Error> {
230        if value.len() != K::INHABITANTS {
231            return Err(value);
232        }
233        Ok(Self {
234            array: value,
235            _phantom: PhantomData,
236        })
237    }
238}
239
240impl<K: Finite, V> From<ExhaustiveMap<K, V>> for Box<[V]> {
241    fn from(value: ExhaustiveMap<K, V>) -> Self {
242        value.array
243    }
244}
245
246impl<K: Finite, V> TryFrom<Vec<V>> for ExhaustiveMap<K, V> {
247    type Error = Vec<V>;
248
249    fn try_from(value: Vec<V>) -> Result<Self, Self::Error> {
250        if value.len() != K::INHABITANTS {
251            return Err(value);
252        }
253        Ok(Self {
254            array: value.into(),
255            _phantom: PhantomData,
256        })
257    }
258}
259
260impl<const N: usize, K: Finite, V> TryFrom<[V; N]> for ExhaustiveMap<K, V> {
261    type Error = [V; N];
262
263    fn try_from(value: [V; N]) -> Result<Self, Self::Error> {
264        if N != K::INHABITANTS {
265            return Err(value);
266        }
267        Ok(Self {
268            array: value.into(),
269            _phantom: PhantomData,
270        })
271    }
272}
273
274impl<K: Finite + Eq + Hash, V> TryFrom<HashMap<K, V>> for ExhaustiveMap<K, V> {
275    type Error = K;
276
277    fn try_from(mut value: HashMap<K, V>) -> Result<Self, Self::Error> {
278        Self::try_from_fn(|k| value.remove(&k).ok_or(k))
279    }
280}
281
282impl<K: Finite + Eq + Hash, V, S: BuildHasher + Default> From<ExhaustiveMap<K, V>>
283    for HashMap<K, V, S>
284{
285    fn from(value: ExhaustiveMap<K, V>) -> Self {
286        Self::from_iter(value)
287    }
288}
289
290impl<K: Finite + Ord, V> From<ExhaustiveMap<K, V>> for BTreeMap<K, V> {
291    fn from(value: ExhaustiveMap<K, V>) -> Self {
292        Self::from_iter(value)
293    }
294}
295
296/// An iterator over the values of an [`ExhaustiveMap`].
297///
298/// This `struct` is created by the [`ExhaustiveMap::values`] method.
299#[must_use = "iterators are lazy and do nothing unless consumed"]
300pub struct Values<'a, V>(std::slice::Iter<'a, V>);
301
302impl<'a, V> Iterator for Values<'a, V> {
303    type Item = &'a V;
304
305    fn next(&mut self) -> Option<Self::Item> {
306        self.0.next()
307    }
308}
309
310/// A mutable iterator over the values of an [`ExhaustiveMap`].
311///
312/// This `struct` is created by the [`ExhaustiveMap::values_mut`] method.
313#[must_use = "iterators are lazy and do nothing unless consumed"]
314pub struct ValuesMut<'a, V>(std::slice::IterMut<'a, V>);
315
316impl<'a, V> Iterator for ValuesMut<'a, V> {
317    type Item = &'a mut V;
318
319    fn next(&mut self) -> Option<Self::Item> {
320        self.0.next()
321    }
322}
323
324/// An owning iterator over the values of an [`ExhaustiveMap`].
325///
326/// This `struct` is created by the [`ExhaustiveMap::into_values`] method.
327#[must_use = "iterators are lazy and do nothing unless consumed"]
328pub struct IntoValues<V>(std::vec::IntoIter<V>);
329
330impl<V> Iterator for IntoValues<V> {
331    type Item = V;
332
333    fn next(&mut self) -> Option<Self::Item> {
334        self.0.next()
335    }
336}
337
338impl<K: Finite, V: Default> Default for ExhaustiveMap<K, V> {
339    fn default() -> Self {
340        Self::from_fn(|_| V::default())
341    }
342}
343
344/// An iterator over the entries of an [`ExhaustiveMap`].
345///
346/// This `struct` is created by the [`ExhaustiveMap::iter`] method.
347#[must_use = "iterators are lazy and do nothing unless consumed"]
348pub struct Iter<'a, K: Finite, V>(std::iter::Zip<IterAll<K>, Values<'a, V>>);
349
350impl<'a, K: Finite, V> Iterator for Iter<'a, K, V> {
351    type Item = (K, &'a V);
352
353    fn next(&mut self) -> Option<Self::Item> {
354        self.0.next()
355    }
356}
357
358/// A mutable iterator over the entries of an [`ExhaustiveMap`].
359///
360/// This `struct` is created by the [`ExhaustiveMap::iter_mut`] method.
361#[must_use = "iterators are lazy and do nothing unless consumed"]
362pub struct IterMut<'a, K: Finite, V>(std::iter::Zip<IterAll<K>, ValuesMut<'a, V>>);
363
364impl<'a, K: Finite, V> Iterator for IterMut<'a, K, V> {
365    type Item = (K, &'a mut V);
366
367    fn next(&mut self) -> Option<Self::Item> {
368        self.0.next()
369    }
370}
371
372/// An owning iterator over the entries of an [`ExhaustiveMap`].
373///
374/// This `struct` is created by the [`into_iter`](IntoIterator::into_iter) method on [`ExhaustiveMap`]
375/// (provided by the [`IntoIterator`] trait).
376#[must_use = "iterators are lazy and do nothing unless consumed"]
377pub struct IntoIter<K: Finite, V>(std::iter::Zip<IterAll<K>, IntoValues<V>>);
378
379impl<K: Finite, V> Iterator for IntoIter<K, V> {
380    type Item = (K, V);
381
382    fn next(&mut self) -> Option<Self::Item> {
383        self.0.next()
384    }
385}
386
387impl<K: Finite, V> IntoIterator for ExhaustiveMap<K, V> {
388    type Item = (K, V);
389
390    type IntoIter = IntoIter<K, V>;
391
392    fn into_iter(self) -> Self::IntoIter {
393        IntoIter(Self::keys().zip(self.into_values()))
394    }
395}
396
397impl<'a, K: Finite, V> IntoIterator for &'a ExhaustiveMap<K, V> {
398    type Item = (K, &'a V);
399
400    type IntoIter = Iter<'a, K, V>;
401
402    fn into_iter(self) -> Self::IntoIter {
403        self.iter()
404    }
405}
406
407impl<'a, K: Finite, V> IntoIterator for &'a mut ExhaustiveMap<K, V> {
408    type Item = (K, &'a mut V);
409
410    type IntoIter = IterMut<'a, K, V>;
411
412    fn into_iter(self) -> Self::IntoIter {
413        self.iter_mut()
414    }
415}
416
417impl<K: Finite + Debug, V: Debug> Debug for ExhaustiveMap<K, V> {
418    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
419        f.debug_map().entries(self).finish()
420    }
421}
422
423impl<K: Finite, V, Q: Borrow<K>> Index<Q> for ExhaustiveMap<K, V> {
424    type Output = V;
425
426    fn index(&self, index: Q) -> &Self::Output {
427        &self.array[K::to_usize(index.borrow())]
428    }
429}
430
431impl<K: Finite, V, Q: Borrow<K>> IndexMut<Q> for ExhaustiveMap<K, V> {
432    fn index_mut(&mut self, index: Q) -> &mut Self::Output {
433        &mut self.array[K::to_usize(index.borrow())]
434    }
435}
436
437// The following traits could have been implemented using a derive macro,
438// however that would put an unnecessary trait bound on the key.
439
440impl<K: Finite, V: Clone> Clone for ExhaustiveMap<K, V> {
441    fn clone(&self) -> Self {
442        Self {
443            array: self.array.clone(),
444            _phantom: PhantomData,
445        }
446    }
447}
448
449impl<K: Finite, V: PartialEq> PartialEq for ExhaustiveMap<K, V> {
450    fn eq(&self, other: &Self) -> bool {
451        self.array.eq(&other.array)
452    }
453}
454
455impl<K: Finite, V: Eq> Eq for ExhaustiveMap<K, V> {}
456
457impl<K: Finite, V: PartialOrd> PartialOrd for ExhaustiveMap<K, V> {
458    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
459        self.array.partial_cmp(&other.array)
460    }
461}
462
463impl<K: Finite, V: Ord> Ord for ExhaustiveMap<K, V> {
464    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
465        self.array.cmp(&other.array)
466    }
467}
468
469impl<K: Finite, V: Hash> Hash for ExhaustiveMap<K, V> {
470    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
471        self.array.hash(state);
472    }
473}
474
475// SAFETY: `ExhaustiveMap<K, V>` is just a transparent wrapper around `Box<[V]>`.
476unsafe impl<K: Finite, V> Send for ExhaustiveMap<K, V> where Box<[V]>: Send {}
477// SAFETY: `ExhaustiveMap<K, V>` is just a transparent wrapper around `Box<[V]>`.
478unsafe impl<K: Finite, V> Sync for ExhaustiveMap<K, V> where Box<[V]>: Sync {}
479
480#[cfg(test)]
481mod test {
482    use super::*;
483
484    #[derive(Finite)]
485    struct Key(PhantomData<*mut u8>);
486
487    #[allow(unused)]
488    const fn assert_implements_traits<
489        T: Send + Sync + Default + Clone + PartialEq + Eq + PartialOrd + Ord + Hash,
490    >() {
491    }
492
493    const _: () = assert_implements_traits::<ExhaustiveMap<Key, bool>>();
494
495    #[test]
496    fn test_uninit() {
497        let mut m = ExhaustiveMap::<bool, u8>::new_uninit();
498        m[true].write(123);
499        m[false].write(45);
500        // SAFETY: All elements has been initialized.
501        let m = unsafe { m.assume_init() };
502        println!("{m:?}");
503    }
504
505    #[test]
506    fn test_conversion() {
507        let m: ExhaustiveMap<bool, u8> = [2, 3].try_into().unwrap();
508        assert_eq!(m[false], 2);
509        assert_eq!(m[true], 3);
510    }
511
512    #[test]
513    fn test_try_unrwap_values() {
514        let m: ExhaustiveMap<bool, Option<u8>> = ExhaustiveMap::from_fn(|_| None);
515        let mut m = m.try_unwrap_values().unwrap_err();
516        m[false] = Some(2);
517        let mut m = m.try_unwrap_values().unwrap_err();
518        m[true] = Some(3);
519        let m = m.try_unwrap_values().unwrap();
520        let expected: ExhaustiveMap<bool, u8> = [2, 3].try_into().unwrap();
521        assert_eq!(m, expected);
522    }
523}
exhaustive_map/map.rs

exhaustive_map/
map.rs
