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