diff --git a/include/dav1d/headers.rs b/include/dav1d/headers.rs index d4eeee3fb..23d5f4117 100644 --- a/include/dav1d/headers.rs +++ b/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; diff --git a/lib.rs b/lib.rs index ae2cb0c1d..c1d7b2495 100644 --- a/lib.rs +++ b/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; diff --git a/src/decode.rs b/src/decode.rs index 189198bc8..82ee4d9e6 100644 --- a/src/decode.rs +++ b/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 = 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; diff --git a/src/enum_map.rs b/src/enum_map.rs new file mode 100644 index 000000000..0cc6d96be --- /dev/null +++ b/src/enum_map.rs @@ -0,0 +1,107 @@ +use std::marker::PhantomData; +use std::ops::Index; +use std::ops::IndexMut; + +pub trait EnumKey: 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 +where + K: EnumKey, +{ + array: [V; N], + _phantom: PhantomData, +} + +impl EnumMap +where + K: EnumKey, +{ + /// 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 EnumMap +where + K: EnumKey, + 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 Index for EnumMap +where + K: EnumKey, +{ + type Output = V; + + fn index(&self, index: K) -> &Self::Output { + &self.array[index.as_usize()] + } +} + +impl IndexMut for EnumMap +where + K: EnumKey, +{ + 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;