enum_map/
lib.rs

Help
1// SPDX-FileCopyrightText: 2017 - 2023 Kamila Borowska <kamila@borowska.pw>
2// SPDX-FileCopyrightText: 2019 Riey <creeper844@gmail.com>
3// SPDX-FileCopyrightText: 2021 Alex Sayers <alex@asayers.com>
4// SPDX-FileCopyrightText: 2021 Bruno Corrêa Zimmermann <brunoczim@gmail.com>
5// SPDX-FileCopyrightText: 2022 Cass Fridkin <cass@cloudflare.com>
6// SPDX-FileCopyrightText: 2022 Mateusz Kowalczyk <fuuzetsu@fuuzetsu.co.uk>
7//
8// SPDX-License-Identifier: MIT OR Apache-2.0
9
10//! An enum mapping type.
11//!
12//! It is implemented using an array type, so using it is as fast as using Rust
13//! arrays.
14//!
15//! # Examples
16//!
17//! ```
18//! use enum_map::{enum_map, Enum, EnumMap};
19//!
20//! #[derive(Debug, Enum)]
21//! enum Example {
22//!     A(bool),
23//!     B,
24//!     C,
25//! }
26//!
27//! let mut map = enum_map! {
28//!     Example::A(false) => 0,
29//!     Example::A(true) => 1,
30//!     Example::B => 2,
31//!     Example::C => 3,
32//! };
33//! map[Example::C] = 4;
34//!
35//! assert_eq!(map[Example::A(true)], 1);
36//!
37//! for (key, &value) in &map {
38//!     println!("{:?} has {} as value.", key, value);
39//! }
40//! ```
41
42#![no_std]
43#![deny(missing_docs)]
44#![warn(clippy::pedantic)]
45
46#[cfg(feature = "arbitrary")]
47mod arbitrary;
48mod enum_map_impls;
49mod internal;
50mod iter;
51#[cfg(feature = "serde")]
52mod serde;
53
54#[doc(hidden)]
55pub use core::mem::{self, ManuallyDrop, MaybeUninit};
56#[doc(hidden)]
57pub use core::primitive::usize;
58use core::slice;
59#[doc(hidden)]
60// unreachable needs to be exported for compatibility with older versions of enum-map-derive
61pub use core::{panic, ptr, unreachable};
62pub use enum_map_derive::Enum;
63#[doc(hidden)]
64pub use internal::out_of_bounds;
65use internal::Array;
66pub use internal::{Enum, EnumArray};
67pub use iter::{IntoIter, IntoValues, Iter, IterMut, Values, ValuesMut};
68
69// SAFETY: initialized needs to represent number of initialized elements
70#[doc(hidden)]
71pub struct Guard<'a, K, V>
72where
73    K: EnumArray<V>,
74{
75    array_mut: &'a mut MaybeUninit<K::Array>,
76    initialized: usize,
77}
78
79impl<K, V> Drop for Guard<'_, K, V>
80where
81    K: EnumArray<V>,
82{
83    fn drop(&mut self) {
84        // This is safe as arr[..len] is initialized due to
85        // Guard's type invariant.
86        unsafe {
87            ptr::slice_from_raw_parts_mut(self.as_mut_ptr(), self.initialized).drop_in_place();
88        }
89    }
90}
91
92impl<'a, K, V> Guard<'a, K, V>
93where
94    K: EnumArray<V>,
95{
96    #[doc(hidden)]
97    pub fn as_mut_ptr(&mut self) -> *mut V {
98        self.array_mut.as_mut_ptr().cast::<V>()
99    }
100
101    #[doc(hidden)]
102    #[must_use]
103    pub fn new(array_mut: &'a mut MaybeUninit<K::Array>) -> Self {
104        Self {
105            array_mut,
106            initialized: 0,
107        }
108    }
109
110    #[doc(hidden)]
111    #[must_use]
112    #[allow(clippy::unused_self)]
113    pub fn storage_length(&self) -> usize {
114        // SAFETY: We need to use LENGTH from K::Array, as K::LENGTH is
115        // untrustworthy.
116        K::Array::LENGTH
117    }
118
119    #[doc(hidden)]
120    #[must_use]
121    pub fn get_key(&self) -> K {
122        K::from_usize(self.initialized)
123    }
124
125    #[doc(hidden)]
126    // Unsafe as it can write out of bounds.
127    pub unsafe fn push(&mut self, value: V) {
128        self.as_mut_ptr().add(self.initialized).write(value);
129        self.initialized += 1;
130    }
131}
132
133#[doc(hidden)]
134pub struct TypeEqualizer<'a, K, V>
135where
136    K: EnumArray<V>,
137{
138    pub enum_map: [EnumMap<K, V>; 0],
139    pub guard: Guard<'a, K, V>,
140}
141
142/// Enum map constructor.
143///
144/// This macro allows to create a new enum map in a type safe way. It takes
145/// a list of `,` separated pairs separated by `=>`. Left side is `|`
146/// separated list of enum keys, or `_` to match all unmatched enum keys,
147/// while right side is a value.
148///
149/// The iteration order when using this macro is not guaranteed to be
150/// consistent. Future releases of this crate may change it, and this is not
151/// considered to be a breaking change.
152///
153/// # Examples
154///
155/// ```
156/// use enum_map::{enum_map, Enum};
157///
158/// #[derive(Enum)]
159/// enum Example {
160///     A,
161///     B,
162///     C,
163///     D,
164/// }
165///
166/// let enum_map = enum_map! {
167///     Example::A | Example::B => 1,
168///     Example::C => 2,
169///     _ => 3,
170/// };
171/// assert_eq!(enum_map[Example::A], 1);
172/// assert_eq!(enum_map[Example::B], 1);
173/// assert_eq!(enum_map[Example::C], 2);
174/// assert_eq!(enum_map[Example::D], 3);
175/// ```
176#[macro_export]
177macro_rules! enum_map {
178    {$($t:tt)*} => {{
179        let mut uninit = $crate::MaybeUninit::uninit();
180        let mut eq = $crate::TypeEqualizer {
181            enum_map: [],
182            guard: $crate::Guard::new(&mut uninit),
183        };
184        if false {
185            // Safe because this code is unreachable
186            unsafe { (&mut eq.enum_map).as_mut_ptr().read() }
187        } else {
188            for _ in 0..(&eq.guard).storage_length() {
189                struct __PleaseDoNotUseBreakWithoutLabel;
190                let _please_do_not_use_continue_without_label;
191                let value;
192                #[allow(unreachable_code)]
193                loop {
194                    _please_do_not_use_continue_without_label = ();
195                    value = match (&eq.guard).get_key() { $($t)* };
196                    break __PleaseDoNotUseBreakWithoutLabel;
197                };
198
199                unsafe { (&mut eq.guard).push(value); }
200            }
201            $crate::mem::forget(eq);
202            // Safe because the array was fully initialized.
203            $crate::EnumMap::from_array(unsafe { uninit.assume_init() })
204        }
205    }};
206}
207
208/// An enum mapping.
209///
210/// This internally uses an array which stores a value for each possible
211/// enum value. To work, it requires implementation of internal (private,
212/// although public due to macro limitations) trait which allows extracting
213/// information about an enum, which can be automatically generated using
214/// `#[derive(Enum)]` macro.
215///
216/// Additionally, `bool` and `u8` automatically derives from `Enum`. While
217/// `u8` is not technically an enum, it's convenient to consider it like one.
218/// In particular, [reverse-complement in benchmark game] could be using `u8`
219/// as an enum.
220///
221/// # Examples
222///
223/// ```
224/// use enum_map::{enum_map, Enum, EnumMap};
225///
226/// #[derive(Enum)]
227/// enum Example {
228///     A,
229///     B,
230///     C,
231/// }
232///
233/// let mut map = EnumMap::default();
234/// // new initializes map with default values
235/// assert_eq!(map[Example::A], 0);
236/// map[Example::A] = 3;
237/// assert_eq!(map[Example::A], 3);
238/// ```
239///
240/// [reverse-complement in benchmark game]:
241///     http://benchmarksgame.alioth.debian.org/u64q/program.php?test=revcomp&lang=rust&id=2
242pub struct EnumMap<K: EnumArray<V>, V> {
243    array: K::Array,
244}
245
246impl<K: EnumArray<V>, V: Default> EnumMap<K, V> {
247    /// Clear enum map with default values.
248    ///
249    /// # Examples
250    ///
251    /// ```
252    /// use enum_map::{Enum, EnumMap};
253    ///
254    /// #[derive(Enum)]
255    /// enum Example {
256    ///     A,
257    ///     B,
258    /// }
259    ///
260    /// let mut enum_map = EnumMap::<_, String>::default();
261    /// enum_map[Example::B] = "foo".into();
262    /// enum_map.clear();
263    /// assert_eq!(enum_map[Example::A], "");
264    /// assert_eq!(enum_map[Example::B], "");
265    /// ```
266    #[inline]
267    pub fn clear(&mut self) {
268        for v in self.as_mut_slice() {
269            *v = V::default();
270        }
271    }
272}
273
274#[allow(clippy::len_without_is_empty)]
275impl<K: EnumArray<V>, V> EnumMap<K, V> {
276    /// Creates an enum map from array.
277    #[inline]
278    pub const fn from_array(array: K::Array) -> EnumMap<K, V> {
279        EnumMap { array }
280    }
281
282    /// Create an enum map, where each value is the returned value from `cb`
283    /// using provided enum key.
284    ///
285    /// ```
286    /// # use enum_map_derive::*;
287    /// use enum_map::{enum_map, Enum, EnumMap};
288    ///
289    /// #[derive(Enum, PartialEq, Debug)]
290    /// enum Example {
291    ///     A,
292    ///     B,
293    /// }
294    ///
295    /// let map = EnumMap::from_fn(|k| k == Example::A);
296    /// assert_eq!(map, enum_map! { Example::A => true, Example::B => false })
297    /// ```
298    pub fn from_fn<F>(mut cb: F) -> Self
299    where
300        F: FnMut(K) -> V,
301    {
302        enum_map! { k => cb(k) }
303    }
304
305    /// Returns an iterator over enum map.
306    ///
307    /// The iteration order is deterministic, and when using [macro@Enum] derive
308    /// it will be the order in which enum variants are declared.
309    ///
310    /// # Examples
311    ///
312    /// ```
313    /// use enum_map::{enum_map, Enum};
314    ///
315    /// #[derive(Enum, PartialEq)]
316    /// enum E {
317    ///     A,
318    ///     B,
319    ///     C,
320    /// }
321    ///
322    /// let map = enum_map! { E::A => 1, E::B => 2, E::C => 3};
323    /// assert!(map.iter().eq([(E::A, &1), (E::B, &2), (E::C, &3)]));
324    /// ```
325    #[inline]
326    pub fn iter(&self) -> Iter<K, V> {
327        self.into_iter()
328    }
329
330    /// Returns a mutable iterator over enum map.
331    #[inline]
332    pub fn iter_mut(&mut self) -> IterMut<K, V> {
333        self.into_iter()
334    }
335
336    /// Returns number of elements in enum map.
337    #[inline]
338    #[allow(clippy::unused_self)]
339    pub const fn len(&self) -> usize {
340        K::Array::LENGTH
341    }
342
343    /// Swaps two indexes.
344    ///
345    /// # Examples
346    ///
347    /// ```
348    /// use enum_map::enum_map;
349    ///
350    /// let mut map = enum_map! { false => 0, true => 1 };
351    /// map.swap(false, true);
352    /// assert_eq!(map[false], 1);
353    /// assert_eq!(map[true], 0);
354    /// ```
355    #[inline]
356    pub fn swap(&mut self, a: K, b: K) {
357        self.as_mut_slice().swap(a.into_usize(), b.into_usize());
358    }
359
360    /// Consumes an enum map and returns the underlying array.
361    ///
362    /// The order of elements is deterministic, and when using [macro@Enum]
363    /// derive it will be the order in which enum variants are declared.
364    ///
365    /// # Examples
366    ///
367    /// ```
368    /// use enum_map::{enum_map, Enum};
369    ///
370    /// #[derive(Enum, PartialEq)]
371    /// enum E {
372    ///     A,
373    ///     B,
374    ///     C,
375    /// }
376    ///
377    /// let map = enum_map! { E::A => 1, E::B => 2, E::C => 3};
378    /// assert_eq!(map.into_array(), [1, 2, 3]);
379    /// ```
380    pub fn into_array(self) -> K::Array {
381        self.array
382    }
383
384    /// Returns a reference to the underlying array.
385    ///
386    /// The order of elements is deterministic, and when using [macro@Enum]
387    /// derive it will be the order in which enum variants are declared.
388    ///
389    /// # Examples
390    ///
391    /// ```
392    /// use enum_map::{enum_map, Enum};
393    ///
394    /// #[derive(Enum, PartialEq)]
395    /// enum E {
396    ///     A,
397    ///     B,
398    ///     C,
399    /// }
400    ///
401    /// let map = enum_map! { E::A => 1, E::B => 2, E::C => 3};
402    /// assert_eq!(map.as_array(), &[1, 2, 3]);
403    /// ```
404    pub const fn as_array(&self) -> &K::Array {
405        &self.array
406    }
407
408    /// Returns a mutable reference to the underlying array.
409    ///
410    /// The order of elements is deterministic, and when using [macro@Enum]
411    /// derive it will be the order in which enum variants are declared.
412    ///
413    /// # Examples
414    ///
415    /// ```
416    /// use enum_map::{enum_map, Enum};
417    ///
418    /// #[derive(Enum, PartialEq)]
419    /// enum E {
420    ///     A,
421    ///     B,
422    ///     C,
423    /// }
424    ///
425    /// let mut map = enum_map! { E::A => 1, E::B => 2, E::C => 3};
426    /// map.as_mut_array()[1] = 42;
427    /// assert_eq!(map.as_array(), &[1, 42, 3]);
428    /// ```
429    pub fn as_mut_array(&mut self) -> &mut K::Array {
430        &mut self.array
431    }
432
433    /// Converts an enum map to a slice representing values.
434    ///
435    /// The order of elements is deterministic, and when using [macro@Enum]
436    /// derive it will be the order in which enum variants are declared.
437    ///
438    /// # Examples
439    ///
440    /// ```
441    /// use enum_map::{enum_map, Enum};
442    ///
443    /// #[derive(Enum, PartialEq)]
444    /// enum E {
445    ///     A,
446    ///     B,
447    ///     C,
448    /// }
449    ///
450    /// let map = enum_map! { E::A => 1, E::B => 2, E::C => 3};
451    /// assert_eq!(map.as_slice(), &[1, 2, 3]);
452    /// ```
453    #[inline]
454    pub fn as_slice(&self) -> &[V] {
455        unsafe { slice::from_raw_parts(ptr::addr_of!(self.array).cast(), K::Array::LENGTH) }
456    }
457
458    /// Converts a mutable enum map to a mutable slice representing values.
459    #[inline]
460    pub fn as_mut_slice(&mut self) -> &mut [V] {
461        unsafe { slice::from_raw_parts_mut(ptr::addr_of_mut!(self.array).cast(), K::Array::LENGTH) }
462    }
463
464    /// Returns an enum map with function `f` applied to each element in order.
465    ///
466    /// # Examples
467    ///
468    /// ```
469    /// use enum_map::enum_map;
470    ///
471    /// let a = enum_map! { false => 0, true => 1 };
472    /// let b = a.map(|_, x| f64::from(x) + 0.5);
473    /// assert_eq!(b, enum_map! { false => 0.5, true => 1.5 });
474    /// ```
475    pub fn map<F, T>(self, mut f: F) -> EnumMap<K, T>
476    where
477        F: FnMut(K, V) -> T,
478        K: EnumArray<T>,
479    {
480        struct DropOnPanic<K, V>
481        where
482            K: EnumArray<V>,
483        {
484            position: usize,
485            map: ManuallyDrop<EnumMap<K, V>>,
486        }
487        impl<K, V> Drop for DropOnPanic<K, V>
488        where
489            K: EnumArray<V>,
490        {
491            fn drop(&mut self) {
492                unsafe {
493                    ptr::drop_in_place(&mut self.map.as_mut_slice()[self.position..]);
494                }
495            }
496        }
497        let mut drop_protect = DropOnPanic {
498            position: 0,
499            map: ManuallyDrop::new(self),
500        };
501        enum_map! {
502            k => {
503                let value = unsafe { ptr::read(&drop_protect.map.as_slice()[drop_protect.position]) };
504                drop_protect.position += 1;
505                f(k, value)
506            }
507        }
508    }
509}
enum_map/lib.rs

enum_map/
lib.rs
