From d14077c15823d4adefb97577d86ed48509cca9b9 Mon Sep 17 00:00:00 2001
From: Khyber Sen <kkysen@gmail.com>
Date: Sun, 19 Nov 2023 10:51:13 -0500
Subject: [PATCH] `enum_map_ty!`: Add a macro to avoid having to hardcode the
 `N` in `EnumMap` by using `EnumCount`.

This has to be a macro until `#![feature(const_generic_exprs)]` lands.
---
 include/dav1d/headers.rs | 9 +++++----
 src/decode.rs            | 3 ++-
 src/enum_map.rs          | 9 +++++++++
 src/filmgrain.rs         | 5 +++--
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/include/dav1d/headers.rs b/include/dav1d/headers.rs
index 88bce6127..2049eb298 100644
--- a/include/dav1d/headers.rs
+++ b/include/dav1d/headers.rs
@@ -2,6 +2,7 @@ use crate::src::enum_map::EnumKey;
 use std::ffi::c_int;
 use std::ffi::c_uint;
 use std::ops::BitAnd;
+use strum::EnumCount;
 
 /// This is so we can store both `*mut D` and `*mut R`
 /// for maintaining `dav1d` ABI compatibility,
@@ -201,7 +202,7 @@ pub const DAV1D_PIXEL_LAYOUT_I422: Dav1dPixelLayout = 2;
 pub const DAV1D_PIXEL_LAYOUT_I420: Dav1dPixelLayout = 1;
 pub const DAV1D_PIXEL_LAYOUT_I400: Dav1dPixelLayout = 0;
 
-#[derive(Clone, Copy, PartialEq, Eq)]
+#[derive(Clone, Copy, PartialEq, Eq, EnumCount)]
 pub(crate) enum Rav1dPixelLayout {
     I400,
     I420,
@@ -210,7 +211,7 @@ pub(crate) enum Rav1dPixelLayout {
 }
 
 impl EnumKey<4> for Rav1dPixelLayout {
-    const VALUES: [Self; 4] = [Self::I400, Self::I420, Self::I422, Self::I444];
+    const VALUES: [Self; Self::COUNT] = [Self::I400, Self::I420, Self::I422, Self::I444];
 
     fn as_usize(self) -> usize {
         self as usize
@@ -252,7 +253,7 @@ impl From<Rav1dPixelLayout> for Dav1dPixelLayout {
     }
 }
 
-#[derive(Clone, Copy, PartialEq, Eq)]
+#[derive(Clone, Copy, PartialEq, Eq, EnumCount)]
 pub(crate) enum Rav1dPixelLayoutSubSampled {
     I420,
     I422,
@@ -260,7 +261,7 @@ pub(crate) enum Rav1dPixelLayoutSubSampled {
 }
 
 impl EnumKey<3> for Rav1dPixelLayoutSubSampled {
-    const VALUES: [Self; 3] = [Self::I420, Self::I422, Self::I444];
+    const VALUES: [Self; Self::COUNT] = [Self::I420, Self::I422, Self::I444];
 
     fn as_usize(self) -> usize {
         self as usize
diff --git a/src/decode.rs b/src/decode.rs
index 908107470..dd7da975a 100644
--- a/src/decode.rs
+++ b/src/decode.rs
@@ -43,6 +43,7 @@ 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::enum_map_ty;
 use crate::src::enum_map::DefaultValue;
 use crate::src::enum_map::EnumMap;
 use crate::src::env::av1_get_bwd_ref_1_ctx;
@@ -3817,7 +3818,7 @@ impl DefaultValue for [u8; 2] {
 }
 
 /// `{ Y+U+V, Y+U } * 4`
-static ss_size_mul: EnumMap<Rav1dPixelLayout, [u8; 2], 4> = enum_map!(Rav1dPixelLayout => [u8; 2]; match key {
+static ss_size_mul: enum_map_ty!(Rav1dPixelLayout, [u8; 2]) = enum_map!(Rav1dPixelLayout => [u8; 2]; match key {
     I400 => [4, 4],
     I420 => [6, 5],
     I422 => [8, 6],
diff --git a/src/enum_map.rs b/src/enum_map.rs
index 64919b92e..1d8f01b6c 100644
--- a/src/enum_map.rs
+++ b/src/enum_map.rs
@@ -29,6 +29,15 @@ where
     _phantom: PhantomData<K>,
 }
 
+// Has to be a macro until we have `#![feature(generic_const_exprs)]`.
+macro_rules! enum_map_ty {
+    ($K:ty, $V:ty) => {
+        EnumMap<$K, $V, { <$K as ::strum::EnumCount>::COUNT }>
+    }
+}
+
+pub(crate) use enum_map_ty;
+
 impl<K, V, const N: usize> EnumMap<K, V, N>
 where
     K: EnumKey<N>,
diff --git a/src/filmgrain.rs b/src/filmgrain.rs
index 1d3e88bf1..c868b532a 100644
--- a/src/filmgrain.rs
+++ b/src/filmgrain.rs
@@ -12,6 +12,7 @@ use crate::include::dav1d::headers::Rav1dFilmGrainData;
 use crate::include::dav1d::headers::Rav1dPixelLayoutSubSampled;
 use crate::src::assume::assume;
 use crate::src::enum_map::enum_map;
+use crate::src::enum_map::enum_map_ty;
 use crate::src::enum_map::DefaultValue;
 use crate::src::enum_map::EnumMap;
 use crate::src::internal::GrainLut;
@@ -194,9 +195,9 @@ impl FnFGUV32x32xN {
 #[repr(C)]
 pub(crate) struct Rav1dFilmGrainDSPContext {
     pub generate_grain_y: FnGenerateGrainY,
-    pub generate_grain_uv: EnumMap<Rav1dPixelLayoutSubSampled, FnGenerateGrainUV, 3>,
+    pub generate_grain_uv: enum_map_ty!(Rav1dPixelLayoutSubSampled, FnGenerateGrainUV),
     pub fgy_32x32xn: FnFGY32x32xN,
-    pub fguv_32x32xn: EnumMap<Rav1dPixelLayoutSubSampled, FnFGUV32x32xN, 3>,
+    pub fguv_32x32xn: enum_map_ty!(Rav1dPixelLayoutSubSampled, FnFGUV32x32xN),
 }
 
 #[cfg(feature = "asm")]