#![doc = include_str!(concat!("../", core::env!("CARGO_PKG_README")))]
#![cfg_attr(not(feature = "std"), no_std)]
#[cfg(test)]
pub extern crate self as enum_table;
use core::marker::PhantomData;
#[cfg(feature = "derive")]
pub use enum_table_derive::Enumable;
pub mod builder;
mod intrinsics;
pub mod __private {
pub use crate::intrinsics::{sort_variants, variant_index_of};
}
mod impls;
pub use impls::*;
mod macros;
pub trait Enumable: Copy + 'static {
const VARIANTS: &'static [Self];
const COUNT: usize = Self::VARIANTS.len();
fn variant_index(&self) -> usize {
intrinsics::binary_search_index::<Self>(self)
}
}
pub struct EnumTable<K: Enumable, V, const N: usize> {
table: [V; N],
_phantom: PhantomData<K>,
}
impl<K: Enumable, V, const N: usize> EnumTable<K, V, N> {
pub(crate) const fn new(table: [V; N]) -> Self {
const {
assert!(
N == K::COUNT,
"EnumTable: N must equal K::COUNT. The const generic N does not match the number of enum variants."
);
}
#[cfg(debug_assertions)]
const {
if !intrinsics::is_sorted(K::VARIANTS) {
panic!(
"Enumable: variants are not sorted by discriminant. Use `enum_table::Enumable` derive macro to ensure correct ordering."
);
}
}
Self {
table,
_phantom: PhantomData,
}
}
pub fn new_with_fn(mut f: impl FnMut(&K) -> V) -> Self {
Self::new(core::array::from_fn(|i| f(&K::VARIANTS[i])))
}
pub fn try_new_with_fn<E>(mut f: impl FnMut(&K) -> Result<V, E>) -> Result<Self, (K, E)> {
let table = intrinsics::try_collect_array(|i| {
let variant = &K::VARIANTS[i];
f(variant).map_err(|e| (*variant, e))
})?;
Ok(Self::new(table))
}
pub fn checked_new_with_fn(mut f: impl FnMut(&K) -> Option<V>) -> Result<Self, K> {
let table = intrinsics::try_collect_array(|i| {
let variant = &K::VARIANTS[i];
f(variant).ok_or(*variant)
})?;
Ok(Self::new(table))
}
pub fn get(&self, variant: &K) -> &V {
&self.table[variant.variant_index()]
}
pub fn get_mut(&mut self, variant: &K) -> &mut V {
&mut self.table[variant.variant_index()]
}
pub fn set(&mut self, variant: &K, value: V) -> V {
core::mem::replace(&mut self.table[variant.variant_index()], value)
}
pub const fn get_const(&self, variant: &K) -> &V {
let idx = intrinsics::binary_search_index::<K>(variant);
&self.table[idx]
}
pub const fn get_mut_const(&mut self, variant: &K) -> &mut V {
let idx = intrinsics::binary_search_index::<K>(variant);
&mut self.table[idx]
}
pub const fn set_const(&mut self, variant: &K, value: V) -> V {
let idx = intrinsics::binary_search_index::<K>(variant);
core::mem::replace(&mut self.table[idx], value)
}
pub const fn len(&self) -> usize {
N
}
pub const fn is_empty(&self) -> bool {
N == 0
}
pub const fn as_slice(&self) -> &[V] {
&self.table
}
pub const fn as_mut_slice(&mut self) -> &mut [V] {
&mut self.table
}
pub fn into_array(self) -> [V; N] {
self.table
}
pub fn zip<U, W>(
self,
other: EnumTable<K, U, N>,
mut f: impl FnMut(V, U) -> W,
) -> EnumTable<K, W, N> {
let mut other_iter = other.table.into_iter();
EnumTable::new(self.table.map(|v| {
let u = unsafe { other_iter.next().unwrap_unchecked() };
f(v, u)
}))
}
pub fn map<U>(self, f: impl FnMut(V) -> U) -> EnumTable<K, U, N> {
EnumTable::new(self.table.map(f))
}
pub fn map_with_key<U>(self, mut f: impl FnMut(&K, V) -> U) -> EnumTable<K, U, N> {
let mut i = 0;
EnumTable::new(self.table.map(|value| {
let key = &K::VARIANTS[i];
i += 1;
f(key, value)
}))
}
pub fn map_mut(&mut self, f: impl FnMut(&mut V)) {
self.table.iter_mut().for_each(f);
}
pub fn map_mut_with_key(&mut self, mut f: impl FnMut(&K, &mut V)) {
self.table.iter_mut().enumerate().for_each(|(i, value)| {
f(&K::VARIANTS[i], value);
});
}
}
impl<K: Enumable, V, const N: usize> EnumTable<K, Option<V>, N> {
pub const fn new_fill_with_none() -> Self {
Self::new([const { None }; N])
}
pub fn clear_to_none(&mut self) {
for value in &mut self.table {
*value = None;
}
}
pub fn remove(&mut self, variant: &K) -> Option<V> {
self.table[variant.variant_index()].take()
}
pub const fn remove_const(&mut self, variant: &K) -> Option<V> {
let idx = intrinsics::binary_search_index::<K>(variant);
self.table[idx].take()
}
}
impl<K: Enumable, V: Copy, const N: usize> EnumTable<K, V, N> {
pub const fn new_fill_with_copy(value: V) -> Self {
Self::new([value; N])
}
}
impl<K: Enumable, V: Default, const N: usize> EnumTable<K, V, N> {
pub fn new_fill_with_default() -> Self {
Self::new(core::array::from_fn(|_| V::default()))
}
pub fn clear_to_default(&mut self) {
self.table.fill_with(V::default);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Enumable)]
enum Color {
Red = 33,
Green = 11,
Blue = 222,
}
const TABLES: EnumTable<Color, &'static str, { Color::COUNT }> =
crate::et!(Color, &'static str, |color| match color {
Color::Red => "Red",
Color::Green => "Green",
Color::Blue => "Blue",
});
#[test]
fn new_with_fn() {
let table =
EnumTable::<Color, &'static str, { Color::COUNT }>::new_with_fn(|color| match color {
Color::Red => "Red",
Color::Green => "Green",
Color::Blue => "Blue",
});
assert_eq!(table.get(&Color::Red), &"Red");
assert_eq!(table.get(&Color::Green), &"Green");
assert_eq!(table.get(&Color::Blue), &"Blue");
}
#[test]
fn try_new_with_fn() {
let table =
EnumTable::<Color, &'static str, { Color::COUNT }>::try_new_with_fn(
|color| match color {
Color::Red => Ok::<&'static str, core::convert::Infallible>("Red"),
Color::Green => Ok("Green"),
Color::Blue => Ok("Blue"),
},
);
assert!(table.is_ok());
let table = table.unwrap();
assert_eq!(table.get(&Color::Red), &"Red");
assert_eq!(table.get(&Color::Green), &"Green");
assert_eq!(table.get(&Color::Blue), &"Blue");
let error_table = EnumTable::<Color, &'static str, { Color::COUNT }>::try_new_with_fn(
|color| match color {
Color::Red => Ok("Red"),
Color::Green => Err("Error on Green"),
Color::Blue => Ok("Blue"),
},
);
assert!(error_table.is_err());
let (variant, error) = error_table.unwrap_err();
assert_eq!(variant, Color::Green);
assert_eq!(error, "Error on Green");
}
#[test]
fn checked_new_with_fn() {
let table =
EnumTable::<Color, &'static str, { Color::COUNT }>::checked_new_with_fn(|color| {
match color {
Color::Red => Some("Red"),
Color::Green => Some("Green"),
Color::Blue => Some("Blue"),
}
});
assert!(table.is_ok());
let table = table.unwrap();
assert_eq!(table.get(&Color::Red), &"Red");
assert_eq!(table.get(&Color::Green), &"Green");
assert_eq!(table.get(&Color::Blue), &"Blue");
let error_table =
EnumTable::<Color, &'static str, { Color::COUNT }>::checked_new_with_fn(|color| {
match color {
Color::Red => Some("Red"),
Color::Green => None,
Color::Blue => Some("Blue"),
}
});
assert!(error_table.is_err());
let variant = error_table.unwrap_err();
assert_eq!(variant, Color::Green);
}
#[test]
fn get() {
assert_eq!(TABLES.get(&Color::Red), &"Red");
assert_eq!(TABLES.get(&Color::Green), &"Green");
assert_eq!(TABLES.get(&Color::Blue), &"Blue");
}
#[test]
fn get_mut() {
let mut table = TABLES;
assert_eq!(table.get_mut(&Color::Red), &mut "Red");
assert_eq!(table.get_mut(&Color::Green), &mut "Green");
assert_eq!(table.get_mut(&Color::Blue), &mut "Blue");
*table.get_mut(&Color::Red) = "Changed Red";
*table.get_mut(&Color::Green) = "Changed Green";
*table.get_mut(&Color::Blue) = "Changed Blue";
assert_eq!(table.get(&Color::Red), &"Changed Red");
assert_eq!(table.get(&Color::Green), &"Changed Green");
assert_eq!(table.get(&Color::Blue), &"Changed Blue");
}
#[test]
fn set() {
let mut table = TABLES;
assert_eq!(table.set(&Color::Red, "New Red"), "Red");
assert_eq!(table.set(&Color::Green, "New Green"), "Green");
assert_eq!(table.set(&Color::Blue, "New Blue"), "Blue");
assert_eq!(table.get(&Color::Red), &"New Red");
assert_eq!(table.get(&Color::Green), &"New Green");
assert_eq!(table.get(&Color::Blue), &"New Blue");
}
#[test]
fn keys() {
let keys: Vec<_> = TABLES.keys().collect();
assert_eq!(keys, vec![&Color::Green, &Color::Red, &Color::Blue]);
}
#[test]
fn values() {
let values: Vec<_> = TABLES.values().collect();
assert_eq!(values, vec![&"Green", &"Red", &"Blue"]);
}
#[test]
fn iter() {
let iter: Vec<_> = TABLES.iter().collect();
assert_eq!(
iter,
vec![
(&Color::Green, &"Green"),
(&Color::Red, &"Red"),
(&Color::Blue, &"Blue")
]
);
}
#[test]
fn iter_mut() {
let mut table = TABLES;
for (key, value) in table.iter_mut() {
*value = match key {
Color::Red => "Changed Red",
Color::Green => "Changed Green",
Color::Blue => "Changed Blue",
};
}
let iter: Vec<_> = table.iter().collect();
assert_eq!(
iter,
vec![
(&Color::Green, &"Changed Green"),
(&Color::Red, &"Changed Red"),
(&Color::Blue, &"Changed Blue")
]
);
}
#[test]
fn map() {
let table = EnumTable::<Color, i32, { Color::COUNT }>::new_with_fn(|color| match color {
Color::Red => 1,
Color::Green => 2,
Color::Blue => 3,
});
let doubled = table.map(|value| value * 2);
assert_eq!(doubled.get(&Color::Red), &2);
assert_eq!(doubled.get(&Color::Green), &4);
assert_eq!(doubled.get(&Color::Blue), &6);
}
#[test]
fn map_with_key() {
let table = EnumTable::<Color, i32, { Color::COUNT }>::new_with_fn(|color| match color {
Color::Red => 1,
Color::Green => 2,
Color::Blue => 3,
});
let mapped = table.map_with_key(|key, value| match key {
Color::Red => value + 10, Color::Green => value + 20, Color::Blue => value + 30, });
assert_eq!(mapped.get(&Color::Red), &11);
assert_eq!(mapped.get(&Color::Green), &22);
assert_eq!(mapped.get(&Color::Blue), &33);
}
#[test]
fn map_mut() {
let mut table =
EnumTable::<Color, i32, { Color::COUNT }>::new_with_fn(|color| match color {
Color::Red => 10,
Color::Green => 20,
Color::Blue => 30,
});
table.map_mut(|value| *value += 5);
assert_eq!(table.get(&Color::Red), &15);
assert_eq!(table.get(&Color::Green), &25);
assert_eq!(table.get(&Color::Blue), &35);
}
#[test]
fn map_mut_with_key() {
let mut table =
EnumTable::<Color, i32, { Color::COUNT }>::new_with_fn(|color| match color {
Color::Red => 10,
Color::Green => 20,
Color::Blue => 30,
});
table.map_mut_with_key(|key, value| {
*value += match key {
Color::Red => 1, Color::Green => 2, Color::Blue => 3, }
});
assert_eq!(table.get(&Color::Red), &11);
assert_eq!(table.get(&Color::Green), &22);
assert_eq!(table.get(&Color::Blue), &33);
}
macro_rules! run_variants_test {
($($variant:ident),+) => {{
#[derive(Debug, Clone, Copy, PartialEq, Eq, Enumable)]
#[repr(u8)]
enum Test {
$($variant,)*
}
let map = EnumTable::<Test, &'static str, { Test::COUNT }>::new_with_fn(|t| match t {
$(Test::$variant => stringify!($variant),)*
});
$(
assert_eq!(map.get(&Test::$variant), &stringify!($variant));
)*
}};
}
#[test]
fn binary_search_correct_variants() {
run_variants_test!(A);
run_variants_test!(A, B);
run_variants_test!(A, B, C);
run_variants_test!(A, B, C, D);
run_variants_test!(A, B, C, D, E);
}
#[test]
fn variant_index() {
assert_eq!(Color::Green.variant_index(), 0);
assert_eq!(Color::Red.variant_index(), 1);
assert_eq!(Color::Blue.variant_index(), 2);
}
#[test]
fn get_const() {
const RED: &str = TABLES.get_const(&Color::Red);
const GREEN: &str = TABLES.get_const(&Color::Green);
const BLUE: &str = TABLES.get_const(&Color::Blue);
assert_eq!(RED, "Red");
assert_eq!(GREEN, "Green");
assert_eq!(BLUE, "Blue");
}
#[test]
fn set_const() {
const fn make_table() -> EnumTable<Color, &'static str, { Color::COUNT }> {
let mut table = TABLES;
table.set_const(&Color::Red, "New Red");
table
}
const TABLE: EnumTable<Color, &'static str, { Color::COUNT }> = make_table();
assert_eq!(TABLE.get_const(&Color::Red), &"New Red");
assert_eq!(TABLE.get_const(&Color::Green), &"Green");
}
#[test]
fn get_mut_const() {
const fn make_table() -> EnumTable<Color, &'static str, { Color::COUNT }> {
let mut table = TABLES;
*table.get_mut_const(&Color::Green) = "Changed Green";
table
}
const TABLE: EnumTable<Color, &'static str, { Color::COUNT }> = make_table();
assert_eq!(TABLE.get_const(&Color::Green), &"Changed Green");
}
#[test]
fn remove_option() {
let mut table =
EnumTable::<Color, Option<i32>, { Color::COUNT }>::new_with_fn(|color| match color {
Color::Red => Some(1),
Color::Green => Some(2),
Color::Blue => None,
});
assert_eq!(table.remove(&Color::Red), Some(1));
assert_eq!(table.get(&Color::Red), &None);
assert_eq!(table.remove(&Color::Blue), None);
assert_eq!(table.get(&Color::Blue), &None);
}
#[test]
fn remove_const_option() {
const fn make_table() -> EnumTable<Color, Option<i32>, { Color::COUNT }> {
let mut table = EnumTable::new_fill_with_none();
table.set_const(&Color::Red, Some(42));
table.set_const(&Color::Green, Some(99));
table
}
let mut table = make_table();
assert_eq!(table.remove_const(&Color::Red), Some(42));
assert_eq!(table.get(&Color::Red), &None);
}
#[test]
fn as_slice() {
let slice = TABLES.as_slice();
assert_eq!(slice.len(), 3);
assert_eq!(slice[0], "Green");
assert_eq!(slice[1], "Red");
assert_eq!(slice[2], "Blue");
}
#[test]
fn as_mut_slice() {
let mut table = TABLES;
let slice = table.as_mut_slice();
slice[0] = "Changed Green";
assert_eq!(table.get(&Color::Green), &"Changed Green");
}
#[test]
fn into_array() {
let arr = TABLES.into_array();
assert_eq!(arr, ["Green", "Red", "Blue"]);
}
#[test]
fn zip() {
let a = EnumTable::<Color, i32, { Color::COUNT }>::new_with_fn(|c| match c {
Color::Red => -10,
Color::Green => -20,
Color::Blue => -30,
});
let b = EnumTable::<Color, u32, { Color::COUNT }>::new_with_fn(|c| match c {
Color::Red => 1,
Color::Green => 2,
Color::Blue => 3,
});
let sum = a.zip(b, |x, y| (x + y as i32) as i8);
assert_eq!(sum.get(&Color::Red), &-9);
assert_eq!(sum.get(&Color::Green), &-18);
assert_eq!(sum.get(&Color::Blue), &-27);
}
}