-
Notifications
You must be signed in to change notification settings - Fork 22
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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`.
- Loading branch information
Showing
4 changed files
with
132 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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; |