enum_table/
lib.rs

1#![doc = include_str!(concat!("../", std::env!("CARGO_PKG_README")))]
2
3#[cfg(feature = "derive")]
4pub use enum_table_derive::Enumable;
5
6pub mod builder;
7mod impls;
8mod macros;
9
10use dev_macros::*;
11
12/// A trait for enumerations that can be used with `EnumTable`.
13///
14/// This trait requires that the enumeration provides a static array of its variants
15/// and a constant representing the count of these variants.
16pub trait Enumable: Sized + 'static {
17    const VARIANTS: &'static [Self];
18    const COUNT: usize = Self::VARIANTS.len();
19}
20
21const fn to_usize<T: Copy>(t: T) -> usize {
22    #[inline(always)]
23    const fn cast<U>(t: &impl Sized) -> &U {
24        // SAFETY: This is safe because we ensure that the type T is a valid representation
25        unsafe { std::mem::transmute(t) }
26    }
27
28    let t = &t;
29
30    match const { core::mem::size_of::<T>() } {
31        1 => *cast::<u8>(t) as usize,
32        2 => *cast::<u16>(t) as usize,
33        4 => *cast::<u32>(t) as usize,
34        #[cfg(target_pointer_width = "64")]
35        8 => *cast::<u64>(t) as usize,
36        #[cfg(target_pointer_width = "32")]
37        8 => panic!("Unsupported size: 64-bit value found on a 32-bit architecture"),
38        _ => panic!("Values larger than u64 are not supported"),
39    }
40}
41
42const fn from_usize<T>(u: &usize) -> &T {
43    unsafe {
44        // SAFETY: This is safe because we ensure that the usize is derived from a valid T
45        std::mem::transmute::<&usize, &T>(u)
46    }
47}
48
49/// A table that associates each variant of an enumeration with a value.
50///
51/// `EnumTable` is a generic struct that uses an enumeration as keys and stores
52/// associated values. It provides constant-time access to the values based on
53/// the enumeration variant. This is particularly useful when you want to map
54/// enum variants to specific values without the overhead of a `HashMap`.
55///
56/// # Type Parameters
57///
58/// * `K`: The enumeration type that implements the `Enumable` trait. This trait
59///   ensures that the enum provides a static array of its variants and a count
60///   of these variants.
61/// * `V`: The type of values to be associated with each enum variant.
62/// * `N`: The number of variants in the enum, which should match the length of
63///   the static array of variants provided by the `Enumable` trait.
64///
65/// # Note
66/// The `new` method allows for the creation of an `EnumTable` in `const` contexts,
67/// but it does not perform compile-time checks. For enhanced compile-time safety
68/// and convenience, it is advisable to use the [`crate::et`] macro or
69/// [`crate::builder::EnumTableBuilder`], which provide these checks.
70///
71/// # Examples
72///
73/// ```rust
74/// use enum_table::{EnumTable, Enumable};
75///
76/// #[derive(Enumable)]
77/// enum Color {
78///     Red,
79///     Green,
80///     Blue,
81/// }
82///
83/// // Create an EnumTable using the new_with_fn method
84/// let table = EnumTable::<Color, &'static str, { Color::COUNT }>::new_with_fn(|color| match color {
85///     Color::Red => "Red",
86///     Color::Green => "Green",
87///     Color::Blue => "Blue",
88/// });
89///
90/// // Access values associated with enum variants
91/// assert_eq!(table.get(&Color::Red), &"Red");
92/// assert_eq!(table.get(&Color::Green), &"Green");
93/// assert_eq!(table.get(&Color::Blue), &"Blue");
94/// ```
95pub struct EnumTable<K: Enumable, V, const N: usize> {
96    table: [(usize, V); N],
97    _phantom: core::marker::PhantomData<K>,
98}
99
100impl<K: Enumable, V, const N: usize> EnumTable<K, V, N> {
101    /// Creates a new `EnumTable` with the given table of discriminants and values.
102    /// Typically, you would use the [`crate::et`] macro or the [`crate::builder::EnumTableBuilder`] instead.
103    ///
104    ///
105    /// # Arguments
106    ///
107    /// * `table` - An array of tuples where each tuple contains a discriminant of
108    ///   an enumeration variant and its associated value.
109    pub const fn new(table: [(usize, V); N]) -> Self {
110        Self {
111            table,
112            _phantom: core::marker::PhantomData,
113        }
114    }
115
116    /// Create a new EnumTable with a function that takes a variant and returns a value.
117    /// If you want to define it in const, use [`crate::et`] macro
118    /// Creates a new `EnumTable` using a function to generate values for each variant.
119    ///
120    /// # Arguments
121    ///
122    /// * `f` - A function that takes a reference to an enumeration variant and returns
123    ///   a value to be associated with that variant.
124    pub fn new_with_fn(mut f: impl FnMut(&K) -> V) -> Self {
125        let table = core::array::from_fn(|i| {
126            let k = &K::VARIANTS[i];
127            (to_usize(core::mem::discriminant(k)), f(k))
128        });
129
130        Self {
131            table,
132            _phantom: core::marker::PhantomData,
133        }
134    }
135
136    /// Returns a reference to the value associated with the given enumeration variant.
137    ///
138    /// # Arguments
139    ///
140    /// * `variant` - A reference to an enumeration variant.
141    pub const fn get(&self, variant: &K) -> &V {
142        use_variant_value!(self, variant, i, {
143            return &self.table[i].1;
144        });
145    }
146
147    /// Returns a mutable reference to the value associated with the given enumeration variant.
148    ///
149    /// # Arguments
150    ///
151    /// * `variant` - A reference to an enumeration variant.
152    pub const fn get_mut(&mut self, variant: &K) -> &mut V {
153        use_variant_value!(self, variant, i, {
154            return &mut self.table[i].1;
155        });
156    }
157
158    /// Sets the value associated with the given enumeration variant.
159    ///
160    /// # Arguments
161    ///
162    /// * `variant` - A reference to an enumeration variant.
163    /// * `value` - The new value to associate with the variant.
164    /// # Returns
165    /// The old value associated with the variant.
166    pub const fn set(&mut self, variant: &K, value: V) -> V {
167        use_variant_value!(self, variant, i, {
168            return core::mem::replace(&mut self.table[i].1, value);
169        });
170    }
171
172    /// Returns the number of generic N
173    pub const fn len(&self) -> usize {
174        N
175    }
176
177    /// Returns `false` since the table is never empty.
178    pub const fn is_empty(&self) -> bool {
179        false
180    }
181
182    /// Returns an iterator over references to the keys in the table.
183    pub fn keys(&self) -> impl Iterator<Item = &K> {
184        self.table
185            .iter()
186            .map(|(discriminant, _)| from_usize(discriminant))
187    }
188
189    /// Returns an iterator over references to the values in the table.
190    pub fn values(&self) -> impl Iterator<Item = &V> {
191        self.table.iter().map(|(_, value)| value)
192    }
193
194    /// Returns an iterator over mutable references to the values in the table.
195    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
196        self.table.iter_mut().map(|(_, value)| value)
197    }
198
199    /// Returns an iterator over mutable references to the values in the table.
200    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
201        self.table
202            .iter()
203            .map(|(discriminant, value)| (from_usize(discriminant), value))
204    }
205
206    /// Returns an iterator over mutable references to the values in the table.
207    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
208        self.table
209            .iter_mut()
210            .map(|(discriminant, value)| (from_usize(discriminant), value))
211    }
212}
213
214mod dev_macros {
215    macro_rules! use_variant_value {
216        ($self:ident, $variant:ident, $i:ident,{$($tt:tt)+}) => {
217            let discriminant = to_usize(core::mem::discriminant($variant));
218
219            let mut $i = 0;
220            while $i < $self.table.len() {
221                if $self.table[$i].0 == discriminant {
222                    $($tt)+
223                }
224                $i += 1;
225            }
226            unreachable!();
227        };
228    }
229
230    pub(super) use use_variant_value;
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
238    enum Color {
239        Red,
240        Green,
241        Blue,
242    }
243
244    impl Enumable for Color {
245        const VARIANTS: &'static [Self] = &[Color::Red, Color::Green, Color::Blue];
246    }
247
248    const TABLES: EnumTable<Color, &'static str, { Color::COUNT }> =
249        crate::et!(Color, &'static str, |color| match color {
250            Color::Red => "Red",
251            Color::Green => "Green",
252            Color::Blue => "Blue",
253        });
254
255    #[test]
256    fn new_with_fn() {
257        let table =
258            EnumTable::<Color, &'static str, { Color::COUNT }>::new_with_fn(|color| match color {
259                Color::Red => "Red",
260                Color::Green => "Green",
261                Color::Blue => "Blue",
262            });
263
264        assert_eq!(table.get(&Color::Red), &"Red");
265        assert_eq!(table.get(&Color::Green), &"Green");
266        assert_eq!(table.get(&Color::Blue), &"Blue");
267    }
268
269    #[test]
270    fn get() {
271        assert_eq!(TABLES.get(&Color::Red), &"Red");
272        assert_eq!(TABLES.get(&Color::Green), &"Green");
273        assert_eq!(TABLES.get(&Color::Blue), &"Blue");
274    }
275
276    #[test]
277    fn get_mut() {
278        let mut table = TABLES;
279        assert_eq!(table.get_mut(&Color::Red), &mut "Red");
280        assert_eq!(table.get_mut(&Color::Green), &mut "Green");
281        assert_eq!(table.get_mut(&Color::Blue), &mut "Blue");
282
283        *table.get_mut(&Color::Red) = "Changed Red";
284        *table.get_mut(&Color::Green) = "Changed Green";
285        *table.get_mut(&Color::Blue) = "Changed Blue";
286
287        assert_eq!(table.get(&Color::Red), &"Changed Red");
288        assert_eq!(table.get(&Color::Green), &"Changed Green");
289        assert_eq!(table.get(&Color::Blue), &"Changed Blue");
290    }
291
292    #[test]
293    fn set() {
294        let mut table = TABLES;
295        assert_eq!(table.set(&Color::Red, "New Red"), "Red");
296        assert_eq!(table.set(&Color::Green, "New Green"), "Green");
297        assert_eq!(table.set(&Color::Blue, "New Blue"), "Blue");
298
299        assert_eq!(table.get(&Color::Red), &"New Red");
300        assert_eq!(table.get(&Color::Green), &"New Green");
301        assert_eq!(table.get(&Color::Blue), &"New Blue");
302    }
303
304    #[test]
305    fn keys() {
306        let keys: Vec<_> = TABLES.keys().collect();
307        assert_eq!(keys, vec![&Color::Red, &Color::Green, &Color::Blue]);
308    }
309
310    #[test]
311    fn values() {
312        let values: Vec<_> = TABLES.values().collect();
313        assert_eq!(values, vec![&"Red", &"Green", &"Blue"]);
314    }
315
316    #[test]
317    fn iter() {
318        let iter: Vec<_> = TABLES.iter().collect();
319        assert_eq!(
320            iter,
321            vec![
322                (&Color::Red, &"Red"),
323                (&Color::Green, &"Green"),
324                (&Color::Blue, &"Blue")
325            ]
326        );
327    }
328
329    #[test]
330    fn iter_mut() {
331        let mut table = TABLES;
332        for (key, value) in table.iter_mut() {
333            *value = match key {
334                Color::Red => "Changed Red",
335                Color::Green => "Changed Green",
336                Color::Blue => "Changed Blue",
337            };
338        }
339        let iter: Vec<_> = table.iter().collect();
340        assert_eq!(
341            iter,
342            vec![
343                (&Color::Red, &"Changed Red"),
344                (&Color::Green, &"Changed Green"),
345                (&Color::Blue, &"Changed Blue")
346            ]
347        );
348    }
349}