1#![doc = include_str!(concat!("../", core::env!("CARGO_PKG_README")))]
2#![cfg_attr(not(feature = "std"), no_std)]
3
4#[cfg(feature = "alloc")]
5extern crate alloc;
6
7#[cfg(test)]
8pub extern crate self as enum_table;
9
10use core::marker::PhantomData;
11
12#[cfg(feature = "derive")]
13pub use enum_table_derive::Enumerable;
14
15mod builder;
16mod intrinsics;
17
18#[doc(hidden)]
19pub mod __private {
20 pub use crate::builder::EnumTableBuilder;
21 pub use crate::intrinsics::{sort_variants, variant_index_of};
22}
23
24mod impls;
25
26mod macros;
27
28pub unsafe trait Enumerable: Copy + 'static {
68 const VARIANTS: &'static [Self];
69 const COUNT: usize = Self::VARIANTS.len();
70
71 fn variant_index(&self) -> usize {
76 intrinsics::binary_search_index::<Self>(self)
77 }
78}
79
80pub struct EnumTable<K: Enumerable, V, const N: usize> {
108 table: [V; N],
109 _phantom: PhantomData<K>,
110}
111
112impl<K: Enumerable, V, const N: usize> EnumTable<K, V, N> {
113 pub(crate) const fn assert() {
114 const {
115 assert!(
116 N == K::COUNT,
117 "EnumTable: N must equal K::COUNT. The const generic N does not match the number of enum variants."
118 );
119 assert!(
120 intrinsics::is_sorted(K::VARIANTS),
121 "EnumTable: K::VARIANTS is not sorted in ascending order by unsigned bit-pattern. This is required by the `Enumerable` trait's safety contract; use `#[derive(Enumerable)]` instead of a hand-written `unsafe impl` to avoid this."
122 );
123 }
124 }
125
126 pub(crate) const fn new(table: [V; N]) -> Self {
127 const { Self::assert() };
128
129 Self {
130 table,
131 _phantom: PhantomData,
132 }
133 }
134
135 pub fn from_fn(mut f: impl FnMut(K) -> V) -> Self {
137 Self::new(core::array::from_fn(|i| f(K::VARIANTS[i])))
138 }
139
140 pub fn try_from_fn<E>(mut f: impl FnMut(K) -> Result<V, E>) -> Result<Self, E> {
166 let table = intrinsics::try_collect_array(|i| f(K::VARIANTS[i]))?;
167 Ok(Self::new(table))
168 }
169
170 pub fn checked_from_fn(mut f: impl FnMut(K) -> Option<V>) -> Option<Self> {
196 Self::try_from_fn(|k| f(k).ok_or(())).ok()
197 }
198
199 pub fn get(&self, variant: K) -> &V {
201 &self.table[variant.variant_index()]
202 }
203
204 pub fn get_mut(&mut self, variant: K) -> &mut V {
206 &mut self.table[variant.variant_index()]
207 }
208
209 pub fn set(&mut self, variant: K, value: V) -> V {
211 core::mem::replace(&mut self.table[variant.variant_index()], value)
212 }
213
214 pub const fn get_const(&self, variant: K) -> &V {
216 let idx = intrinsics::binary_search_index::<K>(&variant);
217 &self.table[idx]
218 }
219
220 pub const fn get_mut_const(&mut self, variant: K) -> &mut V {
222 let idx = intrinsics::binary_search_index::<K>(&variant);
223 &mut self.table[idx]
224 }
225
226 pub const fn set_const(&mut self, variant: K, value: V) -> V {
228 let idx = intrinsics::binary_search_index::<K>(&variant);
229 core::mem::replace(&mut self.table[idx], value)
230 }
231
232 pub fn zip_with<U, W>(
264 self,
265 other: EnumTable<K, U, N>,
266 mut f: impl FnMut(K, V, U) -> W,
267 ) -> EnumTable<K, W, N> {
268 let mut other_iter = other.table.into_iter();
269 self.map(|k, v| {
270 let u = unsafe { other_iter.next().unwrap_unchecked() };
272 f(k, v, u)
273 })
274 }
275
276 pub fn map<U>(self, mut f: impl FnMut(K, V) -> U) -> EnumTable<K, U, N> {
303 let mut i = 0;
304 EnumTable::new(self.table.map(|value| {
305 let key = K::VARIANTS[i];
306 i += 1;
307 f(key, value)
308 }))
309 }
310
311 pub fn for_each(&self, mut f: impl FnMut(K, &V)) {
336 self.table.iter().enumerate().for_each(|(i, value)| {
337 f(K::VARIANTS[i], value);
338 });
339 }
340
341 pub fn for_each_mut(&mut self, mut f: impl FnMut(K, &mut V)) {
368 self.table.iter_mut().enumerate().for_each(|(i, value)| {
369 f(K::VARIANTS[i], value);
370 });
371 }
372}
373
374impl<K: Enumerable, V: Copy, const N: usize> EnumTable<K, V, N> {
375 pub const fn from_elem(value: V) -> Self {
396 Self::new([value; N])
397 }
398}
399
400impl<K: Enumerable, V: Default, const N: usize> EnumTable<K, V, N> {
401 pub fn clear(&mut self) {
403 self.table.fill_with(V::default);
404 }
405
406 pub fn take(&mut self, variant: K) -> V {
409 core::mem::take(&mut self.table[variant.variant_index()])
410 }
411}
412
413#[cfg(test)]
414mod tests {
415 use super::*;
416
417 #[derive(Debug, Clone, Copy, PartialEq, Eq, Enumerable)]
418 enum Color {
419 Red = 33,
420 Green = 11,
421 Blue = 222,
422 }
423
424 const TABLES: EnumTable<Color, &'static str, { Color::COUNT }> =
425 crate::et!(Color, &'static str, |color| match color {
426 Color::Red => "Red",
427 Color::Green => "Green",
428 Color::Blue => "Blue",
429 });
430
431 #[test]
432 fn from_fn() {
433 let table =
434 EnumTable::<Color, &'static str, { Color::COUNT }>::from_fn(|color| match color {
435 Color::Red => "Red",
436 Color::Green => "Green",
437 Color::Blue => "Blue",
438 });
439
440 assert_eq!(table.get(Color::Red), &"Red");
441 assert_eq!(table.get(Color::Green), &"Green");
442 assert_eq!(table.get(Color::Blue), &"Blue");
443 }
444
445 #[test]
446 fn try_from_fn() {
447 let table =
448 EnumTable::<Color, &'static str, { Color::COUNT }>::try_from_fn(|color| match color {
449 Color::Red => Ok::<&'static str, core::convert::Infallible>("Red"),
450 Color::Green => Ok("Green"),
451 Color::Blue => Ok("Blue"),
452 });
453
454 assert!(table.is_ok());
455 let table = table.unwrap();
456
457 assert_eq!(table.get(Color::Red), &"Red");
458 assert_eq!(table.get(Color::Green), &"Green");
459 assert_eq!(table.get(Color::Blue), &"Blue");
460
461 let error_table =
462 EnumTable::<Color, &'static str, { Color::COUNT }>::try_from_fn(|color| match color {
463 Color::Red => Ok("Red"),
464 Color::Green => Err("Error on Green"),
465 Color::Blue => Ok("Blue"),
466 });
467
468 assert_eq!(error_table, Err("Error on Green"));
469 }
470
471 #[test]
472 fn checked_from_fn() {
473 let table =
474 EnumTable::<Color, &'static str, { Color::COUNT }>::checked_from_fn(
475 |color| match color {
476 Color::Red => Some("Red"),
477 Color::Green => Some("Green"),
478 Color::Blue => Some("Blue"),
479 },
480 );
481
482 assert!(table.is_some());
483 let table = table.unwrap();
484
485 assert_eq!(table.get(Color::Red), &"Red");
486 assert_eq!(table.get(Color::Green), &"Green");
487 assert_eq!(table.get(Color::Blue), &"Blue");
488
489 let none_table = EnumTable::<Color, &'static str, { Color::COUNT }>::checked_from_fn(
490 |color| match color {
491 Color::Red => Some("Red"),
492 Color::Green => None,
493 Color::Blue => Some("Blue"),
494 },
495 );
496
497 assert!(none_table.is_none());
498 }
499
500 #[test]
501 fn get() {
502 assert_eq!(TABLES.get(Color::Red), &"Red");
503 assert_eq!(TABLES.get(Color::Green), &"Green");
504 assert_eq!(TABLES.get(Color::Blue), &"Blue");
505 }
506
507 #[test]
508 fn get_mut() {
509 let mut table = TABLES;
510 assert_eq!(table.get_mut(Color::Red), &mut "Red");
511 assert_eq!(table.get_mut(Color::Green), &mut "Green");
512 assert_eq!(table.get_mut(Color::Blue), &mut "Blue");
513
514 *table.get_mut(Color::Red) = "Changed Red";
515 *table.get_mut(Color::Green) = "Changed Green";
516 *table.get_mut(Color::Blue) = "Changed Blue";
517
518 assert_eq!(table.get(Color::Red), &"Changed Red");
519 assert_eq!(table.get(Color::Green), &"Changed Green");
520 assert_eq!(table.get(Color::Blue), &"Changed Blue");
521 }
522
523 #[test]
524 fn set() {
525 let mut table = TABLES;
526 assert_eq!(table.set(Color::Red, "New Red"), "Red");
527 assert_eq!(table.set(Color::Green, "New Green"), "Green");
528 assert_eq!(table.set(Color::Blue, "New Blue"), "Blue");
529
530 assert_eq!(table.get(Color::Red), &"New Red");
531 assert_eq!(table.get(Color::Green), &"New Green");
532 assert_eq!(table.get(Color::Blue), &"New Blue");
533 }
534
535 #[test]
536 fn keys() {
537 let keys: Vec<_> = TABLES.keys().collect();
538 assert_eq!(keys, vec![&Color::Green, &Color::Red, &Color::Blue]);
539 }
540
541 #[test]
542 fn values() {
543 let values: Vec<_> = TABLES.values().collect();
544 assert_eq!(values, vec![&"Green", &"Red", &"Blue"]);
545 }
546
547 #[test]
548 fn iter() {
549 let iter: Vec<_> = TABLES.iter().collect();
550 assert_eq!(
551 iter,
552 vec![
553 (&Color::Green, &"Green"),
554 (&Color::Red, &"Red"),
555 (&Color::Blue, &"Blue")
556 ]
557 );
558 }
559
560 #[test]
561 fn iter_mut() {
562 let mut table = TABLES;
563 for (key, value) in table.iter_mut() {
564 *value = match key {
565 Color::Red => "Changed Red",
566 Color::Green => "Changed Green",
567 Color::Blue => "Changed Blue",
568 };
569 }
570 let iter: Vec<_> = table.iter().collect();
571 assert_eq!(
572 iter,
573 vec![
574 (&Color::Green, &"Changed Green"),
575 (&Color::Red, &"Changed Red"),
576 (&Color::Blue, &"Changed Blue")
577 ]
578 );
579 }
580
581 #[test]
582 fn map() {
583 let table = EnumTable::<Color, i32, { Color::COUNT }>::from_fn(|color| match color {
584 Color::Red => 1,
585 Color::Green => 2,
586 Color::Blue => 3,
587 });
588
589 let mapped = table.map(|key, value| match key {
590 Color::Red => value + 10,
591 Color::Green => value + 20,
592 Color::Blue => value + 30,
593 });
594
595 assert_eq!(mapped.get(Color::Red), &11);
596 assert_eq!(mapped.get(Color::Green), &22);
597 assert_eq!(mapped.get(Color::Blue), &33);
598 }
599
600 #[test]
601 fn for_each_mut() {
602 let mut table = EnumTable::<Color, i32, { Color::COUNT }>::from_fn(|color| match color {
603 Color::Red => 10,
604 Color::Green => 20,
605 Color::Blue => 30,
606 });
607
608 table.for_each_mut(|key, value| {
609 *value += match key {
610 Color::Red => 1,
611 Color::Green => 2,
612 Color::Blue => 3,
613 }
614 });
615
616 assert_eq!(table.get(Color::Red), &11);
617 assert_eq!(table.get(Color::Green), &22);
618 assert_eq!(table.get(Color::Blue), &33);
619 }
620
621 macro_rules! run_variants_test {
622 ($($variant:ident),+) => {{
623 #[derive(Debug, Clone, Copy, PartialEq, Eq, Enumerable)]
624 #[repr(u8)]
625 enum Test {
626 $($variant,)*
627 }
628
629 let map = EnumTable::<Test, &'static str, { Test::COUNT }>::from_fn(|t| match t {
630 $(Test::$variant => stringify!($variant),)*
631 });
632 $(
633 assert_eq!(map.get(Test::$variant), &stringify!($variant));
634 )*
635 }};
636 }
637
638 #[test]
639 fn binary_search_correct_variants() {
640 run_variants_test!(A);
641 run_variants_test!(A, B);
642 run_variants_test!(A, B, C);
643 run_variants_test!(A, B, C, D);
644 run_variants_test!(A, B, C, D, E);
645 }
646
647 #[test]
648 fn variant_index() {
649 assert_eq!(Color::Green.variant_index(), 0);
652 assert_eq!(Color::Red.variant_index(), 1);
653 assert_eq!(Color::Blue.variant_index(), 2);
654 }
655
656 #[derive(Debug, Clone, Copy, PartialEq, Eq, Enumerable)]
657 #[repr(i8)]
658 enum Signed {
659 Neg = -1,
660 Zero = 0,
661 Pos = 1,
662 }
663
664 #[test]
665 fn signed_repr_end_to_end() {
666 assert_eq!(Signed::VARIANTS, &[Signed::Zero, Signed::Pos, Signed::Neg]);
668 assert_eq!(Signed::Zero.variant_index(), 0);
669 assert_eq!(Signed::Pos.variant_index(), 1);
670 assert_eq!(Signed::Neg.variant_index(), 2);
671
672 let table = EnumTable::<Signed, &'static str, { Signed::COUNT }>::from_fn(|s| match s {
673 Signed::Neg => "neg",
674 Signed::Zero => "zero",
675 Signed::Pos => "pos",
676 });
677
678 assert_eq!(table.get(Signed::Neg), &"neg");
679 assert_eq!(table.get(Signed::Zero), &"zero");
680 assert_eq!(table.get(Signed::Pos), &"pos");
681 }
682
683 #[test]
684 fn get_const() {
685 const RED: &str = TABLES.get_const(Color::Red);
686 const GREEN: &str = TABLES.get_const(Color::Green);
687 const BLUE: &str = TABLES.get_const(Color::Blue);
688
689 assert_eq!(RED, "Red");
690 assert_eq!(GREEN, "Green");
691 assert_eq!(BLUE, "Blue");
692 }
693
694 #[test]
695 fn set_const() {
696 const fn make_table() -> EnumTable<Color, &'static str, { Color::COUNT }> {
697 let mut table = TABLES;
698 table.set_const(Color::Red, "New Red");
699 table
700 }
701 const TABLE: EnumTable<Color, &'static str, { Color::COUNT }> = make_table();
702 assert_eq!(TABLE.get_const(Color::Red), &"New Red");
703 assert_eq!(TABLE.get_const(Color::Green), &"Green");
704 }
705
706 #[test]
707 fn get_mut_const() {
708 const fn make_table() -> EnumTable<Color, &'static str, { Color::COUNT }> {
709 let mut table = TABLES;
710 *table.get_mut_const(Color::Green) = "Changed Green";
711 table
712 }
713 const TABLE: EnumTable<Color, &'static str, { Color::COUNT }> = make_table();
714 assert_eq!(TABLE.get_const(Color::Green), &"Changed Green");
715 }
716
717 #[test]
718 fn take_option() {
719 let mut table =
720 EnumTable::<Color, Option<i32>, { Color::COUNT }>::from_fn(|color| match color {
721 Color::Red => Some(1),
722 Color::Green => Some(2),
723 Color::Blue => None,
724 });
725
726 assert_eq!(table.take(Color::Red), Some(1));
727 assert_eq!(table.get(Color::Red), &None);
728
729 assert_eq!(table.take(Color::Blue), None);
730 assert_eq!(table.get(Color::Blue), &None);
731 }
732
733 #[test]
734 fn take_default() {
735 let mut table = EnumTable::<Color, i32, { Color::COUNT }>::from_fn(|color| match color {
736 Color::Red => 1,
737 Color::Green => 2,
738 Color::Blue => 3,
739 });
740
741 assert_eq!(table.take(Color::Red), 1);
742 assert_eq!(table.get(Color::Red), &0);
743 assert_eq!(table.get(Color::Green), &2);
744 }
745
746 #[test]
747 fn clear_option() {
748 let mut table =
749 EnumTable::<Color, Option<i32>, { Color::COUNT }>::from_fn(|color| match color {
750 Color::Red => Some(1),
751 Color::Green => Some(2),
752 Color::Blue => Some(3),
753 });
754
755 table.clear();
756
757 assert_eq!(table.get(Color::Red), &None);
758 assert_eq!(table.get(Color::Green), &None);
759 assert_eq!(table.get(Color::Blue), &None);
760 }
761
762 #[test]
763 fn zip_with() {
764 let a = EnumTable::<Color, i32, { Color::COUNT }>::from_fn(|c| match c {
765 Color::Red => -10,
766 Color::Green => -20,
767 Color::Blue => -30,
768 });
769 let b = EnumTable::<Color, u32, { Color::COUNT }>::from_fn(|c| match c {
770 Color::Red => 1,
771 Color::Green => 2,
772 Color::Blue => 3,
773 });
774
775 let sum = a.zip_with(b, |key, x, y| match key {
776 Color::Blue => x + y as i32 - 100, _ => x + y as i32,
778 });
779 assert_eq!(sum.get(Color::Red), &-9);
780 assert_eq!(sum.get(Color::Green), &-18);
781 assert_eq!(sum.get(Color::Blue), &-127);
782 }
783
784 #[test]
785 fn clear_default() {
786 let mut table = EnumTable::<Color, i32, { Color::COUNT }>::from_fn(|color| match color {
787 Color::Red => 1,
788 Color::Green => 2,
789 Color::Blue => 3,
790 });
791
792 table.clear();
793
794 assert_eq!(table.get(Color::Red), &0);
795 assert_eq!(table.get(Color::Green), &0);
796 assert_eq!(table.get(Color::Blue), &0);
797 }
798}