array_map/
map.rs

1use crate::Indexable;
2use core::marker::PhantomData;
3use core::mem::MaybeUninit;
4
5/// A Map backed by an array. The keys must be known statically.
6///
7/// This map is O(1) for all operations.
8/// All members must be initialized at creation time so there is no option
9/// to insert or remove items from the map. Because of this it can be indexed into using the key type.
10#[must_use]
11pub struct ArrayMap<K, V, const N: usize> {
12  pub(crate) array: [V; N],
13  pub(crate) phantom: PhantomData<*const K>,
14}
15
16#[allow(unsafe_code)]
17unsafe impl<K, V: Send, const N: usize> Send for ArrayMap<K, V, N> {}
18#[allow(unsafe_code)]
19unsafe impl<K, V: Sync, const N: usize> Sync for ArrayMap<K, V, N> {}
20impl<K, V: Unpin, const N: usize> Unpin for ArrayMap<K, V, N> {}
21#[cfg(feature = "std")]
22impl<K, V: std::panic::UnwindSafe, const N: usize> std::panic::UnwindSafe for ArrayMap<K, V, N> {}
23
24impl<K, V, const N: usize> ArrayMap<K, V, N> {
25  /// Function that can be used in const contexts to get a new `ArrayMap`.
26  ///
27  /// Callers should ensure that each `V` in the array matches the return of `index()` from `K` to get back expected results.
28  pub const fn new(array: [V; N]) -> Self {
29    Self {
30      array,
31      phantom: PhantomData,
32    }
33  }
34}
35
36impl<K: Indexable, V: Clone, const N: usize> ArrayMap<K, V, N> {
37  /// Returns a new [`ArrayMap`] where all the values are initialized to the same value
38  #[inline(always)]
39  pub fn from_value(v: V) -> Self {
40    Self::from_closure(|_| v.clone())
41  }
42}
43
44impl<K, V: Clone, const N: usize> Clone for ArrayMap<K, V, N> {
45  #[inline(always)]
46  fn clone(&self) -> Self {
47    Self {
48      array: self.array.clone(),
49      phantom: PhantomData,
50    }
51  }
52}
53
54impl<K, V: Copy, const N: usize> Copy for ArrayMap<K, V, N> {}
55
56impl<K, V: PartialEq, const N: usize> PartialEq for ArrayMap<K, V, N> {
57  fn eq(&self, other: &Self) -> bool {
58    self.array.eq(&other.array)
59  }
60}
61
62impl<K, V: Eq, const N: usize> Eq for ArrayMap<K, V, N> {}
63
64impl<K: core::fmt::Debug + Indexable, V: core::fmt::Debug, const N: usize> core::fmt::Debug for ArrayMap<K, V, N> {
65  fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
66    assert_eq!(N, K::SIZE);
67    f.debug_map().entries(self.iter()).finish()
68  }
69}
70
71impl<K: Indexable, V, const N: usize> ArrayMap<K, V, N> {
72  /// Returns a new [`ArrayMap`] where all the values are initialized to the value returned from each call of the closure.
73  /// # Panics
74  /// Panics if [`K::iter()`](Indexable::iter()) returns anything other than `N` items or if [`K::index()`](Indexable::iter()) returns a value >= `N`
75  pub fn from_closure<F: FnMut(&K) -> V>(mut f: F) -> Self {
76    let mut array = PanicSafeInit::new();
77    for k in K::iter() {
78      let v = f(&k);
79      array.write(k, v);
80    }
81    array.into_map().unwrap()
82  }
83
84  /// Returns a new [`ArrayMap`] where all the values are initialized to the same value
85  /// # Errors
86  /// Returns an error the first time that the provided closure returns an error.
87  /// # Panics
88  /// Panics if [`K::iter()`](Indexable::iter()) returns anything other than `N` items or if [`K::index()`](Indexable::iter()) returns a value >= `N`
89  pub fn try_from_closure<E, F: FnMut(&K) -> Result<V, E>>(mut f: F) -> Result<Self, E> {
90    let mut array = PanicSafeInit::new();
91    for k in K::iter() {
92      match f(&k) {
93        Ok(v) => {
94          array.write(k, v);
95        }
96        Err(e) => {
97          return Err(e);
98        }
99      }
100    }
101    Ok(array.into_map().unwrap())
102  }
103
104  /// Collects an iterator of `(K, Result<V, E>)` into a `Result<Self, E>`
105  ///
106  /// The trait [`FromIterator`](core::iter::FromIterator) cannot be implemented for [`Result`](core::result::Result)
107  /// like it is on a Vec because of the orphan rules.
108  ///
109  /// # Errors
110  /// Returns an error with the first emitted error from the provided iterator
111  /// # Panics
112  /// Panics if any of the safety requirements in the [`Indexable`](crate::Indexable) trait are wrong
113  pub fn from_results<E, I: IntoIterator<Item = (K, Result<V, E>)>>(iter: I) -> Result<Self, E> {
114    let mut array = PanicSafeInit::new();
115    for (k, r) in iter {
116      match r {
117        Ok(v) => {
118          array.write(k, v);
119        }
120        Err(e) => {
121          return Err(e);
122        }
123      }
124    }
125    Ok(array.into_map().unwrap_or_else(|| {
126      panic!(
127        "Not all indexes have been set. Indexable::index() for {} probably isn't unique",
128        core::any::type_name::<K>()
129      )
130    }))
131  }
132
133  /// Collects an iterator of `(K, Option<V>)` into an `Option<Self>`
134  ///
135  /// Returns `None` if any of the values is `None` or if the iterator isn't long enough.
136  ///
137  /// # Panics
138  /// Panics if any of the safety requirements in the [`Indexable`](crate::Indexable) trait are wrong
139  pub fn from_options<E, I: IntoIterator<Item = (K, Option<V>)>>(iter: I) -> Option<Self> {
140    let mut array = PanicSafeInit::new();
141    for (k, r) in iter {
142      #[allow(clippy::single_match_else)]
143      match r {
144        Some(v) => array.write(k, v),
145        None => {
146          if !array.filled[k.index()] {
147            return None;
148          }
149        }
150      }
151    }
152    array.into_map()
153  }
154
155  /// Collects an iterator of `(K, V)` into an `Option<Self>`
156  ///
157  /// Returns `None` if there aren't enough elements, or not all values of `K` have been seen.
158  ///
159  /// # Panics
160  /// Panics if any of the safety requirements in the [`Indexable`](crate::Indexable) trait are wrong
161  #[allow(clippy::should_implement_trait)]
162  pub fn from_iter<E, I: IntoIterator<Item = (K, V)>>(iter: I) -> Option<Self> {
163    let mut array = PanicSafeInit::new();
164    for (k, v) in iter {
165      array.write(k, v);
166    }
167    array.into_map()
168  }
169
170  /// Returns an iterator over all the items in the map. Note that the keys are owned.
171  /// # Panics
172  /// Panics if any of the safety requirements in the [`Indexable`](crate::Indexable) trait are wrong
173  #[inline(always)]
174  pub fn iter(&self) -> impl Iterator<Item = (K, &V)> + '_ {
175    assert_eq!(N, K::SIZE);
176    crate::assert_indexable_safe::<K>();
177    K::iter().zip(self.array.iter())
178  }
179
180  /// Returns a mutable iterator over all the items in the map
181  /// # Panics
182  /// Panics if any of the safety requirements in the [`Indexable`](crate::Indexable) trait are wrong
183  #[inline(always)]
184  pub fn iter_mut(&mut self) -> impl Iterator<Item = (K, &mut V)> + '_ {
185    assert_eq!(N, K::SIZE);
186    crate::assert_indexable_safe::<K>();
187    K::iter().zip(self.array.iter_mut())
188  }
189
190  /// Returns an iterator over all the keys in the map. Note that the keys are owned.
191  #[inline(always)]
192  #[allow(clippy::unused_self)]
193  pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
194    K::iter()
195  }
196
197  /// Turns this into an array of mutable references to the original.
198  /// Useful for calling `map` when all you need are mutable references
199  ///
200  /// # Panics
201  /// Can only panic if `K::iter()` is unstable or doesn't meet requirements
202  pub fn each_mut(&mut self) -> ArrayMap<K, &mut V, N> {
203    let mut array = PanicSafeInit::new();
204    for (k, v) in self.iter_mut() {
205      array.write(k, v);
206    }
207    array.into_map().unwrap()
208  }
209
210  /// Turns this into an array of references to the original.
211  /// Useful for calling `map` when all you need are references
212  ///
213  /// # Panics
214  /// Can only panic if `K::iter()` is unstable or doesn't meet requirements
215  pub fn each_ref(&self) -> ArrayMap<K, &V, N> {
216    let mut array = PanicSafeInit::new();
217    for (k, v) in self.iter() {
218      array.write(k, v);
219    }
220    array.into_map().unwrap()
221  }
222
223  /// Maps all values in this map using the provided function
224  ///
225  /// # Panics
226  /// Can only panic if `K::iter()` is unstable or doesn't meet requirements
227  pub fn map<U, F: FnMut(&K, V) -> U>(self, mut f: F) -> ArrayMap<K, U, N> {
228    let mut array = PanicSafeInit::new();
229
230    for (k, v) in self {
231      let u = f(&k, v);
232      array.write(k, u);
233    }
234
235    array.into_map().unwrap()
236  }
237}
238
239impl<K, V, const N: usize> ArrayMap<K, V, N> {
240  /// Returns an iterator over all the values in the map
241  #[inline]
242  pub fn values(&self) -> impl Iterator<Item = &V> + '_ {
243    self.array.iter()
244  }
245  /// Returns a mutable iterator over all the values in the map
246  #[inline]
247  pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> + '_ {
248    self.array.iter_mut()
249  }
250}
251
252impl<K: Indexable, V, const N: usize> core::iter::IntoIterator for ArrayMap<K, V, N> {
253  type Item = (K, V);
254  type IntoIter = core::iter::Zip<K::Iter, core::array::IntoIter<V, N>>;
255
256  fn into_iter(self) -> Self::IntoIter {
257    assert_eq!(N, K::SIZE);
258    crate::assert_indexable_safe::<K>();
259    K::iter().zip(self.array.into_iter())
260  }
261}
262
263impl<K: Indexable, V, const N: usize> core::iter::FromIterator<(K, V)> for ArrayMap<K, V, N> {
264  fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
265    let mut array = PanicSafeInit::new();
266    for (k, v) in iter {
267      array.write(k, v);
268    }
269    array.into_map().unwrap()
270  }
271}
272
273impl<K: Indexable, V, const N: usize> core::iter::FromIterator<(K, V)> for ArrayMap<K, Option<V>, N> {
274  fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
275    let mut this = Self::from_closure(|_| None);
276    for (t, u) in iter {
277      this.array[t.index()] = Some(u);
278    }
279    this
280  }
281}
282
283impl<K: Indexable, V: Default, const N: usize> Default for ArrayMap<K, V, N> {
284  #[inline(always)]
285  fn default() -> Self {
286    Self::from_closure(|_| V::default())
287  }
288}
289
290impl<K, V, const N: usize> AsRef<[V; N]> for ArrayMap<K, V, N> {
291  #[inline(always)]
292  fn as_ref(&self) -> &[V; N] {
293    &self.array
294  }
295}
296
297impl<K, V, const N: usize> AsMut<[V; N]> for ArrayMap<K, V, N> {
298  #[inline(always)]
299  fn as_mut(&mut self) -> &mut [V; N] {
300    &mut self.array
301  }
302}
303
304impl<K, V, const N: usize> AsRef<[V]> for ArrayMap<K, V, N> {
305  #[inline(always)]
306  fn as_ref(&self) -> &[V] {
307    &self.array
308  }
309}
310
311impl<K, V, const N: usize> AsMut<[V]> for ArrayMap<K, V, N> {
312  #[inline(always)]
313  fn as_mut(&mut self) -> &mut [V] {
314    &mut self.array
315  }
316}
317
318impl<K: Indexable, V, const N: usize> core::ops::Index<K> for ArrayMap<K, V, N> {
319  type Output = V;
320
321  #[inline(always)]
322  fn index(&self, index: K) -> &Self::Output {
323    assert_eq!(N, K::SIZE);
324    &self.array[index.index()]
325  }
326}
327
328impl<K: Indexable, V, const N: usize> core::ops::IndexMut<K> for ArrayMap<K, V, N> {
329  #[inline(always)]
330  fn index_mut(&mut self, index: K) -> &mut Self::Output {
331    assert_eq!(N, K::SIZE);
332    &mut self.array[index.index()]
333  }
334}
335
336impl<K: Indexable, V, const N: usize> core::ops::Index<usize> for ArrayMap<K, V, N> {
337  type Output = V;
338
339  #[inline(always)]
340  fn index(&self, index: usize) -> &Self::Output {
341    assert_eq!(N, K::SIZE);
342    &self.array[index]
343  }
344}
345
346impl<K: Indexable, V, const N: usize> core::ops::IndexMut<usize> for ArrayMap<K, V, N> {
347  #[inline(always)]
348  fn index_mut(&mut self, index: usize) -> &mut Self::Output {
349    assert_eq!(N, K::SIZE);
350    &mut self.array[index]
351  }
352}
353
354struct PanicSafeInit<K: Indexable, V, const N: usize> {
355  filled: [bool; N],
356  array: MaybeUninit<[V; N]>,
357  _phantom: PhantomData<*const K>,
358}
359
360impl<K: Indexable, V, const N: usize> PanicSafeInit<K, V, N> {
361  pub fn new() -> Self {
362    assert_eq!(K::SIZE, N);
363    crate::assert_indexable_safe::<K>();
364    Self {
365      filled: [false; N],
366      array: MaybeUninit::uninit(),
367      _phantom: PhantomData,
368    }
369  }
370  pub fn write(&mut self, k: K, v: V) {
371    let index = k.index();
372    #[allow(unsafe_code)]
373    // Safety: We are only writing to values and never read from them.
374    unsafe {
375      self.array.as_mut_ptr().cast::<V>().add(index).write(v);
376      self.filled[index] = true;
377    }
378  }
379  pub fn into_map(self) -> Option<ArrayMap<K, V, N>> {
380    self.filled.iter().all(|x| *x).then(|| {
381      let no_drop = core::mem::ManuallyDrop::new(self);
382      #[allow(unsafe_code)]
383      // The array is known to be initialized, though this might be UB?
384      let array = unsafe { core::ptr::read(&no_drop.array).assume_init() };
385      ArrayMap::new(array)
386    })
387  }
388}
389
390impl<K: Indexable, V, const N: usize> Drop for PanicSafeInit<K, V, N> {
391  fn drop(&mut self) {
392    if core::mem::needs_drop::<V>() {
393      for k in K::iter() {
394        let index = k.index();
395        if self.filled[index] {
396          #[allow(unsafe_code)]
397          // Safety: We only drop items that are previously initialized
398          unsafe {
399            core::ptr::drop_in_place(self.array.as_mut_ptr().cast::<V>().add(index));
400          }
401        }
402      }
403    }
404  }
405}
406
407#[cfg(test)]
408mod test {
409  use super::*;
410  use crate::test::*;
411  use crate::IndexU8;
412
413  fn send_sync_traits<
414    K: Indexable,
415    V,
416    T: Send
417      + Sync
418      + Unpin
419      + Default
420      + Clone
421      + Copy
422      + PartialEq
423      + Eq
424      + core::fmt::Debug
425      + AsRef<[V; N]>
426      + AsMut<[V; N]>
427      + core::ops::Index<K>
428      + core::ops::IndexMut<K>,
429    const N: usize,
430  >() {
431  }
432
433  fn test_traits<
434    I: Indexable + Copy + PartialEq + core::fmt::Debug,
435    V: Send + Sync + Unpin + Default + Clone + Copy + PartialEq + Eq + core::fmt::Debug,
436    F: FnMut(I) -> V,
437    const N: usize,
438  >(
439    mut f: F,
440  ) {
441    send_sync_traits::<I, V, ArrayMap<I, V, N>, N>();
442    let mut whole_map = ArrayMap::<I, V, N>::default();
443    for i in I::iter() {
444      assert_eq!(whole_map[i], V::default());
445      let new_value = f(i);
446      whole_map[i] = new_value;
447      assert_eq!(whole_map[i], new_value);
448      assert!(whole_map.iter().any(|(k, v)| k == i && v == &new_value));
449    }
450    assert_eq!(whole_map.keys().count(), I::SIZE);
451    assert_eq!(whole_map.iter().count(), I::SIZE);
452    assert_eq!(whole_map.iter_mut().count(), I::SIZE);
453  }
454
455  #[test]
456  fn test_impls() {
457    test_traits::<bool, Option<(u8, u8)>, _, { bool::SIZE }>(|b| Some((b as u8 + 1, b as u8)));
458    test_traits::<u8, u32, _, { u8::SIZE }>(|u| u as u32 + 1);
459    test_traits::<Lowercase, &'static str, _, { Lowercase::SIZE }>(|_| "abc");
460  }
461
462  #[test]
463  fn test_map() {
464    let mut m1 = ArrayMap::<IndexU8<5>, u8, 5>::new([0_u8, 1, 2, 3, 4]);
465    let r = m1.each_mut();
466    for (k, v) in r.into_iter() {
467      assert_eq!(k.get(), *v);
468      *v += 1;
469    }
470    let r = m1.each_ref();
471    for (k, v) in r.into_iter() {
472      assert_eq!(k.get() + 1, *v);
473    }
474  }
475}
array_map/map.rs

array_map/
map.rs
