cantor/
map.rs

1use crate::array::Array;
2use crate::*;
3use core::ops::{Index, IndexMut};
4
5/// A complete mapping from keys of type `K` to values of type `V`, implemented using an array
6/// indexed by [`Finite::index_of`] of the key.
7///
8/// # Example
9/// ```
10/// use cantor::{Finite, ArrayMap};
11///
12/// // Define key type
13/// #[derive(Finite, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Debug)]
14/// enum MyType {
15///     A,
16///     B(bool),
17///     C(bool, bool)
18/// };
19///
20/// // Initialize map
21/// let mut map = ArrayMap::new(|x: MyType| match x {
22///     MyType::A => false,
23///     MyType::B(a) => a,
24///     MyType::C(a, _) => a,
25/// });
26///
27/// // Use map
28/// map[MyType::C(true, true)] = false;
29/// assert_eq!(map[MyType::A], false);
30/// assert_eq!(map[MyType::B(true)], true);
31/// assert_eq!(map[MyType::C(true, true)], false);
32/// ```
33pub struct ArrayMap<K: ArrayFinite<V>, V>(K::Array);
34
35/// The trait required to use [`ArrayMap`]. Theoretically, this should apply to all
36/// [`Finite`] types, but due to limitations in const generics, a blanket implementation is not
37/// currently possible.
38///
39/// This is automatically implemented on concrete types that derive [`Finite`]. It can also be
40/// implemented on a particular concrete type using [`impl_concrete_finite`].
41#[doc(hidden)]
42#[allow(clippy::missing_safety_doc)] // Should never be manually implemented.
43pub unsafe trait ArrayFinite<V>: Finite {
44    #[allow(missing_docs)]
45    type Array: Array<V>;
46}
47
48impl<K: ArrayFinite<V>, V> ArrayMap<K, V> {
49    /// Constructs a new [`ArrayMap`] with initial values populated using the given function.
50    pub fn new(mut f: impl FnMut(K) -> V) -> Self {
51        ArrayMap(K::Array::new(|k| {
52            f(unsafe { K::nth(k).unwrap_unchecked() })
53        }))
54    }
55
56    /// Constructs a new [`ArrayMap`] from an array of values, each corresponding to the key
57    /// determined by [`Finite::nth`].
58    ///
59    /// # Example
60    /// ```
61    /// use cantor::*;
62    /// let map = ArrayMap::from([1, 3]);
63    /// assert_eq!(map[false], 1);
64    /// assert_eq!(map[true], 3);
65    /// ```
66    pub fn from(array: K::Array) -> Self {
67        Self(array)
68    }
69
70    /// Applies a mapping function the values of this map.
71    pub fn map_with_key<N>(&self, mut f: impl FnMut(K, &V) -> N) -> ArrayMap<K, N>
72    where
73        K: ArrayFinite<N>,
74    {
75        ArrayMap(<K as ArrayFinite<N>>::Array::new(|k| unsafe {
76            f(
77                K::nth(k).unwrap_unchecked(),
78                self.0.as_slice().get_unchecked(k),
79            )
80        }))
81    }
82
83    /// Applies a mapping function the values of this map.
84    pub fn map<N>(&self, mut f: impl FnMut(&V) -> N) -> ArrayMap<K, N>
85    where
86        K: ArrayFinite<N>,
87    {
88        self.map_with_key(|_, v| f(v))
89    }
90}
91
92impl<K: ArrayFinite<V>, V: Default> Default for ArrayMap<K, V> {
93    fn default() -> Self {
94        ArrayMap(K::Array::new(|_| Default::default()))
95    }
96}
97
98impl<K: ArrayFinite<V>, V> Index<K> for ArrayMap<K, V> {
99    type Output = V;
100    fn index(&self, index: K) -> &Self::Output {
101        let index = K::index_of(index);
102        unsafe { self.0.as_slice().get_unchecked(index) }
103    }
104}
105
106impl<K: ArrayFinite<V>, V> IndexMut<K> for ArrayMap<K, V> {
107    fn index_mut(&mut self, index: K) -> &mut Self::Output {
108        let index = K::index_of(index);
109        unsafe { self.0.as_slice_mut().get_unchecked_mut(index) }
110    }
111}
112
113impl<K: CompressFinite + ArrayFinite<V>, V> Index<Compress<K>> for ArrayMap<K, V> {
114    type Output = V;
115    fn index(&self, index: Compress<K>) -> &Self::Output {
116        let index = Compress::index_of(index);
117        unsafe { self.0.as_slice().get_unchecked(index) }
118    }
119}
120
121impl<K: CompressFinite + ArrayFinite<V>, V> IndexMut<Compress<K>> for ArrayMap<K, V> {
122    fn index_mut(&mut self, index: Compress<K>) -> &mut Self::Output {
123        let index = Compress::index_of(index);
124        unsafe { self.0.as_slice_mut().get_unchecked_mut(index) }
125    }
126}
127
128impl<K: ArrayFinite<V>, V> Clone for ArrayMap<K, V>
129where
130    K::Array: Clone,
131{
132    fn clone(&self) -> Self {
133        Self(self.0.clone())
134    }
135}
136
137impl<K: ArrayFinite<V>, V> Copy for ArrayMap<K, V> where K::Array: Copy {}
138
139impl<K: ArrayFinite<V>, V> PartialEq for ArrayMap<K, V>
140where
141    K::Array: PartialEq,
142{
143    fn eq(&self, other: &Self) -> bool {
144        self.0 == other.0
145    }
146}
147
148impl<K: ArrayFinite<V>, V> Eq for ArrayMap<K, V> where K::Array: Eq {}
149
150impl<K: ArrayFinite<V>, V> PartialOrd for ArrayMap<K, V>
151where
152    K::Array: PartialOrd,
153{
154    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
155        self.0.partial_cmp(&other.0)
156    }
157}
158
159impl<K: ArrayFinite<V>, V> Ord for ArrayMap<K, V>
160where
161    K::Array: Ord,
162{
163    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
164        self.0.cmp(&other.0)
165    }
166}
167
168#[test]
169fn test_map_with_key() {
170    let map = ArrayMap::new(|x| if x { 1 } else { 0 });
171    let map = map.map_with_key(|k, v| if k { *v * 2 } else { *v + 5 });
172    assert_eq!(map[false], 5);
173    assert_eq!(map[true], 2);
174}