enum_map.rs

#![feature(const_generics, const_fn)]

/// Local Variables:
/// rmsbolt-command: "rustc -C opt-level=3"
/// rmsbolt-disassemble: nil
/// End:
use std::{fmt, marker::PhantomData, ops};

pub trait AsIndex {
    const SIZE: usize;

    fn as_idx(&self) -> usize;
}

pub struct EnumMap<I, O, const S: usize>
where
    O: Sized,
{
    inner: [O; S],
    _marker: PhantomData<I>,
}

impl<I, O, const S: usize> fmt::Debug for EnumMap<I, O, S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("EnumMap")
            .field("inner_size", &self.inner.len())
            .field("size", &S)
            .finish()
    }
}

impl<I, O, const S: usize> EnumMap<I, O, S>
where
    I: AsIndex + Sized,
    O: Sized,
{
    #[inline]
    const fn len(&self) -> usize {
        S
    }
}

impl<I, O, const S: usize> ops::Index<I> for EnumMap<I, O, S>
where
    I: AsIndex,
    O: Sized,
{
    type Output = O;

    #[inline]
    fn index(&self, idx: I) -> &Self::Output {
        &self.inner[I::as_idx(&idx)]
    }
}

impl<I, O, const S: usize> ops::IndexMut<I> for EnumMap<I, O, S>
where
    I: AsIndex,
    O: Sized,
{
    #[inline]
    fn index_mut(&mut self, idx: I) -> &mut Self::Output {
        &mut self.inner[I::as_idx(&idx)]
    }
}

macro_rules! enum_map {
    { $idx_ty:ty => $val_ty:ty ; $($key:expr => $val:expr,)+ } => {
        {

            const SIZE : usize = <$idx_ty as AsIndex>::SIZE;

            let mut data = [<$val_ty>::default(); SIZE];
            
            // we only pay for this price (but direct initialization will need to do this anyway)
            $(data[AsIndex::as_idx(&$key)] = $val;)+

            EnumMap::<$idx_ty, $val_ty, SIZE> {
                inner: data, // copy the array (this might be fast since data usually small enough & memcpy usually use simd optimized copy
                _marker: PhantomData::<$idx_ty>,
            }
        }
    }
}

#[cfg(test)]
pub mod tests {
    use super::*;

    #[test]
    fn test_enum_map() {
        enum A {
            X,
            Y,
        }

        impl AsIndex for A {
            const SIZE: usize = 2;

            #[inline]
            fn as_idx(&self) -> usize {
                match self {
                    A::X => 0,
                    A::Y => 1,
                }
            }
        }

        enum B {
            C,
            D,
            A(A),
        }

        impl AsIndex for B {
            const SIZE: usize = 2 + A::SIZE;

            #[inline]
            fn as_idx(&self) -> usize {
                match self {
                    B::C => 0,
                    B::D => 1,
                    B::A(a) => 1 + AsIndex::as_idx(a),
                }
            }
        }

        let data =
            enum_map! { B => usize; B::C => 0, B::D => 1, B::A(A::Y) => 2, B::A(A::X) => 3, };

        assert_eq!(data[B::C], 0);
        assert_eq!(data[B::A(A::X)], 3);

        let mut data =
            enum_map! { B => usize; B::C => 0, B::D => 1, B::A(A::Y) => 2, B::A(A::X) => 3, };

        data[B::C] = 1;

        assert_eq!(data[B::C], 1);
    }
}

fn main() {
    enum A {
        X,
        Y,
    }

    impl AsIndex for A {
        const SIZE: usize = 2;

        #[inline]
        fn as_idx(&self) -> usize {
            match self {
                A::X => 0,
                A::Y => 1,
            }
        }
    }

    enum B {
        C,
        D,
        A(A),
    }

    impl AsIndex for B {
        const SIZE: usize = 2 + A::SIZE;

        #[inline]
        fn as_idx(&self) -> usize {
            match self {
                B::C => 0,
                B::D => 1,
                B::A(a) => 1 + AsIndex::as_idx(a),
            }
        }
    }

    let data = enum_map! { B => usize; B::C => 0, B::D => 1, B::A(A::Y) => 2, B::A(A::X) => 3, };

    assert_eq!(data[B::C], 0);
    assert_eq!(data[B::A(A::X)], 3);

    println!("data: {:?}", data);

    let mut data =
        enum_map! { B => usize; B::C => 0, B::D => 1, B::A(A::Y) => 2, B::A(A::X) => 3, };
}

enum_map_maybe_uninit.rs

#![feature(const_generics, const_fn, maybe_uninit_uninit_array, maybe_uninit_ref)]

use std::mem::MaybeUninit;

/// Local Variables:
/// rmsbolt-command: "rustc -C opt-level=3"
/// rmsbolt-disassemble: nil
/// End:
use std::{fmt, marker::PhantomData, ops};

pub trait IndexOf {
    const SIZE: usize;

    fn index_of(&self) -> usize;
}

pub struct EnumMap<I, O, const S: usize>
where
    O: Sized,
{
    inner: [MaybeUninit<O>; S],
    _marker: PhantomData<I>,
}

impl<I, O, const S: usize> fmt::Debug for EnumMap<I, O, S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("EnumMap")
            .field("inner_size", &self.inner.len())
            .field("size", &S)
            .finish()
    }
}

impl<I, O, const S: usize> EnumMap<I, O, S>
where
    I: IndexOf + Sized,
    O: Sized,
{
    #[inline]
    const fn len(&self) -> usize {
        S
    }
}

impl<I, O, const S: usize> ops::Index<I> for EnumMap<I, O, S>
where
    I: IndexOf,
    O: Sized,
{
    type Output = O;

    #[inline]
    fn index(&self, idx: I) -> &Self::Output {
        let idx = I::index_of(&idx);
        unsafe { self.inner[idx].get_ref() }
    }
}

impl<I, O, const S: usize> ops::IndexMut<I> for EnumMap<I, O, S>
where
    I: IndexOf,
    O: Sized,
{
    #[inline]
    fn index_mut(&mut self, idx: I) -> &mut Self::Output {
        let idx = I::index_of(&idx);
        unsafe { self.inner[idx].get_mut() }
    }
}

macro_rules! enum_map {
    { $idx_ty:ty => $val_ty:ty ; $($key:expr => $val:expr,)+ } => {
        {

            const SIZE : usize = <$idx_ty as IndexOf>::SIZE;

            let mut data = [MaybeUninit::<$val_ty>::uninit(); SIZE];
            let mut counter = 0;
            // we only pay for this price (but direct initialization will need to do this anyway)
            $(
                let value : $val_ty = $val;
                data[IndexOf::index_of(&$key)] = MaybeUninit::new::<>(value);
                counter += 1;
            )+

            // runtime error :'(
            assert!(counter != SIZE, "you're missing some cases for enum or having a repeated enum as an indexes");

            EnumMap::<$idx_ty, $val_ty, SIZE> {
                inner: data, // copy the array (this might be fast since data usually small enough & memcpy usually use simd optimized copy
                _marker: PhantomData::<$idx_ty>,
            }
        }
    }
}

#[cfg(test)]
pub mod tests {
    use super::*;

    #[test]
    fn test_enum_map() {
        enum A {
            X,
            Y,
        }

        impl IndexOf for A {
            const SIZE: usize = 2;

            #[inline]
            fn index_of(&self) -> usize {
                match self {
                    A::X => 0,
                    A::Y => 1,
                }
            }
        }

        enum B {
            C,
            D,
            A(A),
        }

        impl IndexOf for B {
            const SIZE: usize = 2 + A::SIZE;

            #[inline]
            fn index_of(&self) -> usize {
                match self {
                    B::C => 0,
                    B::D => 1,
                    B::A(a) => 1 + IndexOf::index_of(a),
                }
            }
        }

        let data =
            enum_map! { B => usize; B::C => 0, B::D => 1, B::A(A::Y) => 2, B::A(A::X) => 3, };

        assert_eq!(data[B::C], 0);
        assert_eq!(data[B::A(A::X)], 3);

        let mut data =
            enum_map! { B => usize; B::C => 0, B::D => 1, B::A(A::Y) => 2, B::A(A::X) => 3, };

        data[B::C] = 1;

        assert_eq!(data[B::C], 1);
    }
}

fn main() {
    enum A {
        X,
        Y,
    }

    impl IndexOf for A {
        const SIZE: usize = 2;

        #[inline]
        fn index_of(&self) -> usize {
            match self {
                A::X => 0,
                A::Y => 1,
            }
        }
    }

    enum B {
        C,
        D,
        A(A),
    }

    impl IndexOf for B {
        const SIZE: usize = 2 + A::SIZE;

        #[inline]
        fn index_of(&self) -> usize {
            match self {
                B::C => 0,
                B::D => 1,
                B::A(a) => 1 + IndexOf::index_of(a),
            }
        }
    }

    let data = enum_map! { B => usize; B::C => 0, B::D => 1, B::A(A::Y) => 2, B::A(A::X) => 3, };

    assert_eq!(data[B::C], 0);
    assert_eq!(data[B::A(A::X)], 3);

    println!("data: {:?}", data);

    let mut data =
        enum_map! { B => usize; B::C => 0, B::D => 1, B::A(A::Y) => 2, B::A(A::X) => 3, };
}