struct EnumMap, enum_map!: Add a const-compatible enum map an…

…d use it for `static ss_size_mul` (#552)

Using an `enum` map for these is a lot more clear and harder to mess up
accidentally than manually casting the `enum` keys to `usize`, making
sure all variants are initialized, etc. Existing `EnumMap`
implementations either weren't `const`-compatible, which we need to
initialize `static`s like this, and/or weren't quite zero-overhead. This
relatively simple version I made based on the existing `as usize`
approach and existing `EnumMap`s works at `const`-time and is
zero-overhead when I checked it in godbolt.

The downside is a bit of extra verbosity in the type paremeters and
having to `impl DefaultValue` for the value type. The former isn't an
issue; just annoying. For the latter, not all types have a default value
that makes sense, but all the places I've seen this pattern used in
`rav1d`, they do, so that should be fine. The reason we use `trait
DefaultValue` with a `const DEFAULT: Self` is so it's `const`, unlike
`trait Default`, as `trait` `fn`s can't (yet) be `const`.

include/dav1d/headers.rs

@@ -1,3 +1,4 @@
+use crate::src::enum_map::EnumKey;
 use std::ffi::c_int;
 use std::ffi::c_uint;
 use std::ops::BitAnd;
@@ -208,6 +209,14 @@ pub(crate) enum Rav1dPixelLayout {
     I444,
 }
 
+impl EnumKey<4> for Rav1dPixelLayout {
+    const VALUES: [Self; 4] = [Self::I400, Self::I420, Self::I422, Self::I444];
+
+    fn as_usize(self) -> usize {
+        self as usize
+    }
+}
+
 impl BitAnd for Rav1dPixelLayout {
     type Output = bool;
 

lib.rs

@@ -46,6 +46,7 @@ pub mod src {
     mod data;
     mod decode;
     mod dequant_tables;
+    pub(crate) mod enum_map;
     mod env;
     pub(crate) mod error;
     mod fg_apply;

src/decode.rs

@@ -42,6 +42,9 @@ use crate::src::ctx::CaseSet;
 use crate::src::data::rav1d_data_props_copy;
 use crate::src::data::rav1d_data_unref_internal;
 use crate::src::dequant_tables::dav1d_dq_tbl;
+use crate::src::enum_map::enum_map;
+use crate::src::enum_map::DefaultValue;
+use crate::src::enum_map::EnumMap;
 use crate::src::env::av1_get_bwd_ref_1_ctx;
 use crate::src::env::av1_get_bwd_ref_ctx;
 use crate::src::env::av1_get_fwd_ref_1_ctx;
@@ -3809,15 +3812,17 @@ fn reset_context(ctx: &mut BlockContext, keyframe: bool, pass: c_int) {
     ctx.pal_sz.0.fill(0);
 }
 
+impl DefaultValue for [u8; 2] {
+    const DEFAULT: Self = [0; 2];
+}
+
 /// `{ Y+U+V, Y+U } * 4`
-static ss_size_mul: [[u8; 2]; 4] = {
-    let mut a = [[0; 2]; 4];
-    a[Rav1dPixelLayout::I400 as usize] = [4, 4];
-    a[Rav1dPixelLayout::I420 as usize] = [6, 5];
-    a[Rav1dPixelLayout::I422 as usize] = [8, 6];
-    a[Rav1dPixelLayout::I444 as usize] = [12, 8];
-    a
-};
+static ss_size_mul: EnumMap<Rav1dPixelLayout, [u8; 2], 4> = enum_map!(Rav1dPixelLayout => [u8; 2]; match key {
+    I400 => [4, 4],
+    I420 => [6, 5],
+    I422 => [8, 6],
+    I444 => [12, 8],
+});
 
 unsafe fn setup_tile(
     ts: &mut Rav1dTileState,
@@ -3834,7 +3839,7 @@ unsafe fn setup_tile(
     let row_sb_end = (*f.frame_hdr).tiling.row_start_sb[tile_row + 1] as c_int;
     let sb_shift = f.sb_shift;
 
-    let size_mul = &ss_size_mul[f.cur.p.layout as usize];
+    let size_mul = &ss_size_mul[f.cur.p.layout];
     for p in 0..2 {
         ts.frame_thread[p].pal_idx = if !(f.frame_thread.pal_idx).is_null() {
             f.frame_thread
@@ -4290,7 +4295,7 @@ pub(crate) unsafe fn rav1d_decode_frame_init(f: &mut Rav1dFrameContext) -> Rav1d
     }
 
     let num_sb128 = f.sb128w * f.sb128h;
-    let size_mul = &ss_size_mul[f.cur.p.layout as usize];
+    let size_mul = &ss_size_mul[f.cur.p.layout];
     let hbd = ((*f.seq_hdr).hbd != 0) as c_int;
     if c.n_fc > 1 {
         let mut tile_idx = 0;

src/enum_map.rs

@@ -0,0 +1,107 @@
+use std::marker::PhantomData;
+use std::ops::Index;
+use std::ops::IndexMut;
+
+pub trait EnumKey<const N: usize>: Sized + Copy {
+    const VALUES: [Self; N];
+
+    fn as_usize(self) -> usize;
+}
+
+/// This is a `const` version of [`Default::default`].
+/// `trait` `fn`s can't be `const` (yet),
+/// but we can make this an associated `const`.
+pub trait DefaultValue {
+    const DEFAULT: Self;
+}
+
+/// A map from an `enum` key `K` to `V`s.
+/// `N` is the number of possible `enum` values.
+pub struct EnumMap<K, V, const N: usize>
+where
+    K: EnumKey<N>,
+{
+    array: [V; N],
+    _phantom: PhantomData<K>,
+}
+
+impl<K, V, const N: usize> EnumMap<K, V, N>
+where
+    K: EnumKey<N>,
+{
+    /// Create an [`EnumMap`] from an existing array
+    /// where the array's indices correspond to `K`'s values `as usize`.
+    pub const fn new(array: [V; N]) -> Self {
+        Self {
+            array,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<K, V, const N: usize> EnumMap<K, V, N>
+where
+    K: EnumKey<N>,
+    V: DefaultValue,
+{
+    /// Create an [`EnumMap`] with default values when `V: ` [`DefaultValue`].
+    #[allow(dead_code)] // TODO(kkysen) remove when used
+    pub const fn default() -> Self {
+        Self {
+            array: [V::DEFAULT; N],
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<K, V, const N: usize> Index<K> for EnumMap<K, V, N>
+where
+    K: EnumKey<N>,
+{
+    type Output = V;
+
+    fn index(&self, index: K) -> &Self::Output {
+        &self.array[index.as_usize()]
+    }
+}
+
+impl<K, V, const N: usize> IndexMut<K> for EnumMap<K, V, N>
+where
+    K: EnumKey<N>,
+{
+    fn index_mut(&mut self, index: K) -> &mut Self::Output {
+        &mut self.array[index.as_usize()]
+    }
+}
+
+/// Create an [`EnumMap`] where `V: ` [`DefaultValue`]
+/// using a `match` from `K` to `V`.
+///
+/// The [`DefaultValue::DEFAULT`] is not actually ever used,
+/// but needs to exist to be able to
+/// create the array safely and at `const` time.
+///
+/// [`MaybeUninit`] can do this without [`DefaultValue`]
+/// by using `unsafe` initialization,
+/// but it also doesn't yet work in `const` contexts.
+///
+/// [`MaybeUninit`]: std::mem::MaybeUninit
+macro_rules! enum_map {
+    ($K:ty => $V:ty; match key { $($t:tt)* }) => {{
+        use $crate::src::enum_map::EnumKey;
+        use $crate::src::enum_map::EnumMap;
+
+        let mut a = [<$V>::DEFAULT; <$K>::VALUES.len()];
+        let mut i = 0;
+        while i < <$K>::VALUES.len() {
+            let key = <$K>::VALUES[i];
+            use $K::*;
+            let value = match key { $($t)* };
+            a[key as usize] = value;
+            i += 1;
+        }
+        EnumMap::<$K, $V, { <$K>::VALUES.len() }>::new(a)
+    }};
+}
+
+pub(crate) use enum_map;