enum_table/
lib.rs

1#[cfg(feature = "derive")]
2pub use enum_table_derive::Enumable;
3
4pub mod builder;
5mod impls;
6mod macros;
7
8use core::mem::Discriminant;
9use dev_macros::*;
10
11/// A trait for enumerations that can be used with `EnumTable`.
12///
13/// This trait requires that the enumeration provides a static array of its variants
14/// and a constant representing the count of these variants.
15pub trait Enumable: Sized + 'static {
16    const VARIANTS: &'static [Self];
17    const COUNT: usize = Self::VARIANTS.len();
18}
19
20const fn to_usize<T>(t: T) -> usize {
21    #[inline(always)]
22    const fn cast<T, U>(t: T) -> U {
23        use core::mem::ManuallyDrop;
24        unsafe { core::mem::transmute_copy::<ManuallyDrop<T>, U>(&ManuallyDrop::new(t)) }
25    }
26
27    match const { core::mem::size_of::<T>() } {
28        1 => cast::<T, u8>(t) as usize,
29        2 => cast::<T, u16>(t) as usize,
30        4 => cast::<T, u32>(t) as usize,
31        8 => cast::<T, u64>(t) as usize,
32        _ => panic!("Unsupported size"),
33    }
34}
35
36/// A table that associates each variant of an enumeration with a value.
37///
38/// `EnumTable` is a generic struct that uses an enumeration as keys and stores
39/// associated values. It provides constant-time access to the values based on
40/// the enumeration variant.
41#[derive(Debug, Clone, Copy)]
42pub struct EnumTable<K: Enumable, V, const N: usize> {
43    table: [(Discriminant<K>, V); N],
44}
45
46impl<K: Enumable, V, const N: usize> EnumTable<K, V, N> {
47    /// Creates a new `EnumTable` with the given table of discriminants and values.
48    ///
49    /// # Arguments
50    ///
51    /// * `table` - An array of tuples where each tuple contains a discriminant of
52    ///   an enumeration variant and its associated value.
53    pub const fn new(table: [(Discriminant<K>, V); N]) -> Self {
54        Self { table }
55    }
56
57    /// Create a new EnumTable with a function that takes a variant and returns a value.
58    /// If you want to define it in const, use [`crate::et`] macro
59    /// Creates a new `EnumTable` using a function to generate values for each variant.
60    ///
61    /// # Arguments
62    ///
63    /// * `f` - A function that takes a reference to an enumeration variant and returns
64    ///   a value to be associated with that variant.
65    pub fn new_with_fn(mut f: impl FnMut(&K) -> V) -> Self {
66        let table = core::array::from_fn(|i| {
67            let k = &K::VARIANTS[i];
68            (core::mem::discriminant(k), f(k))
69        });
70
71        Self { table }
72    }
73
74    /// Returns a reference to the value associated with the given enumeration variant.
75    ///
76    /// # Arguments
77    ///
78    /// * `variant` - A reference to an enumeration variant.
79    pub const fn get(&self, variant: &K) -> &V {
80        use_variant_value!(self, variant, i, {
81            return &self.table[i].1;
82        });
83    }
84
85    /// Returns a mutable reference to the value associated with the given enumeration variant.
86    ///
87    /// # Arguments
88    ///
89    /// * `variant` - A reference to an enumeration variant.
90    pub const fn get_mut(&mut self, variant: &K) -> &mut V {
91        use_variant_value!(self, variant, i, {
92            return &mut self.table[i].1;
93        });
94    }
95
96    /// const function is not callable drop.
97    /// So, we use forget to avoid calling drop.
98    /// Careful, not to call drop on the old value.
99    /// Sets the value associated with the given enumeration variant in a constant context.
100    ///
101    /// # Arguments
102    ///
103    /// * `variant` - A reference to an enumeration variant.
104    /// * `value` - The new value to associate with the variant.
105    ///
106    /// # Safety
107    ///
108    /// This method uses `std::mem::forget` to avoid calling `drop` on the old value.
109    /// Be careful not to call `drop` on the old value manually.
110    pub const fn const_set(&mut self, variant: &K, value: V) {
111        use_variant_value!(self, variant, i, {
112            let old = core::mem::replace(&mut self.table[i].1, value);
113            std::mem::forget(old);
114            return;
115        });
116    }
117
118    /// Sets the value associated with the given enumeration variant.
119    ///
120    /// # Arguments
121    ///
122    /// * `variant` - A reference to an enumeration variant.
123    /// * `value` - The new value to associate with the variant.
124    pub fn set(&mut self, variant: &K, value: V) {
125        use_variant_value!(self, variant, i, {
126            self.table[i].1 = value;
127            return;
128        });
129    }
130
131    /// Returns the number of generic N
132    pub const fn len(&self) -> usize {
133        N
134    }
135
136    /// Returns `false` since the table is never empty.
137    pub const fn is_empty(&self) -> bool {
138        false
139    }
140}
141
142mod dev_macros {
143    macro_rules! use_variant_value {
144        ($self:ident, $variant:ident, $i:ident,{$($tt:tt)+}) => {
145            let discriminant = core::mem::discriminant($variant);
146
147            let mut $i = 0;
148            while $i < $self.table.len() {
149                if to_usize($self.table[$i].0) == to_usize(discriminant) {
150                    $($tt)+
151                }
152                $i += 1;
153            }
154            unreachable!();
155        };
156    }
157
158    pub(super) use use_variant_value;
159}