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    fn size_hint(&self) -> (usize, Option<usize>) {
310        (self.0.len(), Some(self.0.len()))
311    }
312}
313
314impl<T> ExactSizeIterator for Values<'_, T> {
315    fn len(&self) -> usize {
316        self.0.len()
317    }
318}
319
320impl<T> DoubleEndedIterator for Values<'_, T> {
321    fn next_back(&mut self) -> Option<Self::Item> {
322        self.0.next_back()
323    }
324}
325
326/// A mutable iterator over the values of an [`ExhaustiveMap`].
327///
328/// This `struct` is created by the [`ExhaustiveMap::values_mut`] method.
329#[must_use = "iterators are lazy and do nothing unless consumed"]
330pub struct ValuesMut<'a, V>(std::slice::IterMut<'a, V>);
331
332impl<'a, V> Iterator for ValuesMut<'a, V> {
333    type Item = &'a mut V;
334
335    fn next(&mut self) -> Option<Self::Item> {
336        self.0.next()
337    }
338
339    fn size_hint(&self) -> (usize, Option<usize>) {
340        (self.0.len(), Some(self.0.len()))
341    }
342}
343
344impl<T> ExactSizeIterator for ValuesMut<'_, T> {
345    fn len(&self) -> usize {
346        self.0.len()
347    }
348}
349
350impl<T> DoubleEndedIterator for ValuesMut<'_, T> {
351    fn next_back(&mut self) -> Option<Self::Item> {
352        self.0.next_back()
353    }
354}
355
356/// An owning iterator over the values of an [`ExhaustiveMap`].
357///
358/// This `struct` is created by the [`ExhaustiveMap::into_values`] method.
359#[must_use = "iterators are lazy and do nothing unless consumed"]
360pub struct IntoValues<V>(std::vec::IntoIter<V>);
361
362impl<V> Iterator for IntoValues<V> {
363    type Item = V;
364
365    fn next(&mut self) -> Option<Self::Item> {
366        self.0.next()
367    }
368
369    fn size_hint(&self) -> (usize, Option<usize>) {
370        (self.0.len(), Some(self.0.len()))
371    }
372}
373
374impl<T> ExactSizeIterator for IntoValues<T> {
375    fn len(&self) -> usize {
376        self.0.len()
377    }
378}
379
380impl<T> DoubleEndedIterator for IntoValues<T> {
381    fn next_back(&mut self) -> Option<Self::Item> {
382        self.0.next_back()
383    }
384}
385
386impl<K: Finite, V: Default> Default for ExhaustiveMap<K, V> {
387    fn default() -> Self {
388        Self::from_fn(|_| V::default())
389    }
390}
391
392/// An iterator over the entries of an [`ExhaustiveMap`].
393///
394/// This `struct` is created by the [`ExhaustiveMap::iter`] method.
395#[must_use = "iterators are lazy and do nothing unless consumed"]
396pub struct Iter<'a, K, V>(std::iter::Zip<IterAll<K>, Values<'a, V>>);
397
398impl<'a, K, V> Iterator for Iter<'a, K, V> {
399    type Item = (K, &'a V);
400
401    fn next(&mut self) -> Option<Self::Item> {
402        self.0.next()
403    }
404
405    fn size_hint(&self) -> (usize, Option<usize>) {
406        (self.0.len(), Some(self.0.len()))
407    }
408}
409
410impl<K, V> ExactSizeIterator for Iter<'_, K, V> {
411    fn len(&self) -> usize {
412        self.0.len()
413    }
414}
415
416impl<K, V> DoubleEndedIterator for Iter<'_, K, V> {
417    fn next_back(&mut self) -> Option<Self::Item> {
418        self.0.next_back()
419    }
420}
421
422/// A mutable iterator over the entries of an [`ExhaustiveMap`].
423///
424/// This `struct` is created by the [`ExhaustiveMap::iter_mut`] method.
425#[must_use = "iterators are lazy and do nothing unless consumed"]
426pub struct IterMut<'a, K, V>(std::iter::Zip<IterAll<K>, ValuesMut<'a, V>>);
427
428impl<'a, K, V> Iterator for IterMut<'a, K, V> {
429    type Item = (K, &'a mut V);
430
431    fn next(&mut self) -> Option<Self::Item> {
432        self.0.next()
433    }
434
435    fn size_hint(&self) -> (usize, Option<usize>) {
436        (self.0.len(), Some(self.0.len()))
437    }
438}
439
440impl<K, V> ExactSizeIterator for IterMut<'_, K, V> {
441    fn len(&self) -> usize {
442        self.0.len()
443    }
444}
445
446impl<K, V> DoubleEndedIterator for IterMut<'_, K, V> {
447    fn next_back(&mut self) -> Option<Self::Item> {
448        self.0.next_back()
449    }
450}
451
452/// An owning iterator over the entries of an [`ExhaustiveMap`].
453///
454/// This `struct` is created by the [`into_iter`](IntoIterator::into_iter) method on [`ExhaustiveMap`]
455/// (provided by the [`IntoIterator`] trait).
456#[must_use = "iterators are lazy and do nothing unless consumed"]
457pub struct IntoIter<K, V>(std::iter::Zip<IterAll<K>, IntoValues<V>>);
458
459impl<K, V> Iterator for IntoIter<K, V> {
460    type Item = (K, V);
461
462    fn next(&mut self) -> Option<Self::Item> {
463        self.0.next()
464    }
465
466    fn size_hint(&self) -> (usize, Option<usize>) {
467        (self.0.len(), Some(self.0.len()))
468    }
469}
470
471impl<K, V> ExactSizeIterator for IntoIter<K, V> {
472    fn len(&self) -> usize {
473        self.0.len()
474    }
475}
476
477impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
478    fn next_back(&mut self) -> Option<Self::Item> {
479        self.0.next_back()
480    }
481}
482
483impl<K: Finite, V> IntoIterator for ExhaustiveMap<K, V> {
484    type Item = (K, V);
485
486    type IntoIter = IntoIter<K, V>;
487
488    fn into_iter(self) -> Self::IntoIter {
489        IntoIter(Self::keys().zip(self.into_values()))
490    }
491}
492
493impl<'a, K: Finite, V> IntoIterator for &'a ExhaustiveMap<K, V> {
494    type Item = (K, &'a V);
495
496    type IntoIter = Iter<'a, K, V>;
497
498    fn into_iter(self) -> Self::IntoIter {
499        self.iter()
500    }
501}
502
503impl<'a, K: Finite, V> IntoIterator for &'a mut ExhaustiveMap<K, V> {
504    type Item = (K, &'a mut V);
505
506    type IntoIter = IterMut<'a, K, V>;
507
508    fn into_iter(self) -> Self::IntoIter {
509        self.iter_mut()
510    }
511}
512
513impl<K: Finite + Debug, V: Debug> Debug for ExhaustiveMap<K, V> {
514    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
515        f.debug_map().entries(self).finish()
516    }
517}
518
519impl<K: Finite, V, Q: Borrow<K>> Index<Q> for ExhaustiveMap<K, V> {
520    type Output = V;
521
522    fn index(&self, index: Q) -> &Self::Output {
523        &self.array[K::to_usize(index.borrow())]
524    }
525}
526
527impl<K: Finite, V, Q: Borrow<K>> IndexMut<Q> for ExhaustiveMap<K, V> {
528    fn index_mut(&mut self, index: Q) -> &mut Self::Output {
529        &mut self.array[K::to_usize(index.borrow())]
530    }
531}
532
533// The following traits could have been implemented using a derive macro,
534// however that would put an unnecessary trait bound on the key.
535
536impl<K: Finite, V: Clone> Clone for ExhaustiveMap<K, V> {
537    fn clone(&self) -> Self {
538        Self {
539            array: self.array.clone(),
540            _phantom: PhantomData,
541        }
542    }
543}
544
545impl<K: Finite, V: PartialEq> PartialEq for ExhaustiveMap<K, V> {
546    fn eq(&self, other: &Self) -> bool {
547        self.array.eq(&other.array)
548    }
549}
550
551impl<K: Finite, V: Eq> Eq for ExhaustiveMap<K, V> {}
552
553impl<K: Finite, V: PartialOrd> PartialOrd for ExhaustiveMap<K, V> {
554    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
555        self.array.partial_cmp(&other.array)
556    }
557}
558
559impl<K: Finite, V: Ord> Ord for ExhaustiveMap<K, V> {
560    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
561        self.array.cmp(&other.array)
562    }
563}
564
565impl<K: Finite, V: Hash> Hash for ExhaustiveMap<K, V> {
566    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
567        self.array.hash(state);
568    }
569}
570
571// SAFETY: `ExhaustiveMap<K, V>` is just a transparent wrapper around `Box<[V]>`.
572unsafe impl<K: Finite, V> Send for ExhaustiveMap<K, V> where Box<[V]>: Send {}
573// SAFETY: `ExhaustiveMap<K, V>` is just a transparent wrapper around `Box<[V]>`.
574unsafe impl<K: Finite, V> Sync for ExhaustiveMap<K, V> where Box<[V]>: Sync {}
575
576#[cfg(test)]
577mod test {
578    use super::*;
579
580    #[derive(Finite)]
581    struct Key(PhantomData<*mut u8>);
582
583    #[allow(unused)]
584    const fn assert_implements_traits<
585        T: Send + Sync + Default + Clone + PartialEq + Eq + PartialOrd + Ord + Hash,
586    >() {
587    }
588
589    const _: () = assert_implements_traits::<ExhaustiveMap<Key, bool>>();
590
591    #[test]
592    fn test_uninit() {
593        let mut m = ExhaustiveMap::<bool, u8>::new_uninit();
594        m[true].write(123);
595        m[false].write(45);
596        // SAFETY: All elements has been initialized.
597        let m = unsafe { m.assume_init() };
598        println!("{m:?}");
599    }
600
601    #[test]
602    fn test_conversion() {
603        let m: ExhaustiveMap<bool, u8> = [2, 3].try_into().unwrap();
604        assert_eq!(m[false], 2);
605        assert_eq!(m[true], 3);
606    }
607
608    #[test]
609    fn test_try_unrwap_values() {
610        let m: ExhaustiveMap<bool, Option<u8>> = ExhaustiveMap::from_fn(|_| None);
611        let mut m = m.try_unwrap_values().unwrap_err();
612        m[false] = Some(2);
613        let mut m = m.try_unwrap_values().unwrap_err();
614        m[true] = Some(3);
615        let m = m.try_unwrap_values().unwrap();
616        let expected: ExhaustiveMap<bool, u8> = [2, 3].try_into().unwrap();
617        assert_eq!(m, expected);
618    }
619}
exhaustive_map/map.rs

exhaustive_map/
map.rs
