enum Rav1dTxfmMode: make a real enum (#783)

memorysafety · Mar 4, 2024 · 4932987 · 4932987
2 parents 33f1413 + b831466
commit 4932987
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 24 deletions.
diff --git a/include/dav1d/headers.rs b/include/dav1d/headers.rs
@@ -82,16 +82,31 @@ pub enum Rav1dObuType {
 }
 
 pub type Dav1dTxfmMode = c_uint;
-pub const DAV1D_N_TX_MODES: usize = 3;
-pub const DAV1D_TX_SWITCHABLE: Dav1dTxfmMode = 2;
-pub const DAV1D_TX_LARGEST: Dav1dTxfmMode = 1;
-pub const DAV1D_TX_4X4_ONLY: Dav1dTxfmMode = 0;
+pub const DAV1D_N_TX_MODES: usize = Rav1dTxfmMode::COUNT;
+pub const DAV1D_TX_4X4_ONLY: Dav1dTxfmMode = Rav1dTxfmMode::Only4x4 as Dav1dTxfmMode;
+pub const DAV1D_TX_LARGEST: Dav1dTxfmMode = Rav1dTxfmMode::Largest as Dav1dTxfmMode;
+pub const DAV1D_TX_SWITCHABLE: Dav1dTxfmMode = Rav1dTxfmMode::Switchable as Dav1dTxfmMode;
 
-pub(crate) type Rav1dTxfmMode = c_uint;
-pub(crate) const _RAV1D_N_TX_MODES: usize = DAV1D_N_TX_MODES;
-pub(crate) const RAV1D_TX_SWITCHABLE: Rav1dTxfmMode = DAV1D_TX_SWITCHABLE;
-pub(crate) const RAV1D_TX_LARGEST: Rav1dTxfmMode = DAV1D_TX_LARGEST;
-pub(crate) const RAV1D_TX_4X4_ONLY: Rav1dTxfmMode = DAV1D_TX_4X4_ONLY;
+#[derive(Clone, Copy, PartialEq, Eq, FromRepr, EnumCount)]
+pub enum Rav1dTxfmMode {
+    Only4x4 = 0,
+    Largest = 1,
+    Switchable = 2,
+}
+
+impl From<Rav1dTxfmMode> for Dav1dTxfmMode {
+    fn from(value: Rav1dTxfmMode) -> Self {
+        value as Dav1dTxfmMode
+    }
+}
+
+impl TryFrom<Dav1dTxfmMode> for Rav1dTxfmMode {
+    type Error = ();
+
+    fn try_from(value: Dav1dTxfmMode) -> Result<Self, Self::Error> {
+        Self::from_repr(value as usize).ok_or(())
+    }
+}
 
 pub type Dav1dFilterMode = u8;
 pub const DAV1D_FILTER_SWITCHABLE: Dav1dFilterMode = 4;
@@ -2334,7 +2349,7 @@ impl From<Dav1dFrameHeader> for Rav1dFrameHeader {
             loopfilter: loopfilter.into(),
             cdef: cdef.into(),
             restoration: restoration.into(),
-            txfm_mode,
+            txfm_mode: txfm_mode.try_into().unwrap(),
             switchable_comp_refs,
             skip_mode: Rav1dFrameSkipMode {
                 allowed: skip_mode_allowed,
@@ -2451,7 +2466,7 @@ impl From<Rav1dFrameHeader> for Dav1dFrameHeader {
             loopfilter: loopfilter.into(),
             cdef: cdef.into(),
             restoration: restoration.into(),
-            txfm_mode,
+            txfm_mode: txfm_mode.into(),
             switchable_comp_refs,
             skip_mode_allowed,
             skip_mode_enabled,

diff --git a/src/decode.rs b/src/decode.rs
@@ -8,12 +8,12 @@ use crate::include::common::intops::iclip;
 use crate::include::common::intops::iclip_u8;
 use crate::include::common::intops::ulog2;
 use crate::include::dav1d::headers::Dav1dFilterMode;
-use crate::include::dav1d::headers::Dav1dTxfmMode;
 use crate::include::dav1d::headers::Rav1dFrameHeader;
 use crate::include::dav1d::headers::Rav1dFrameHeader_tiling;
 use crate::include::dav1d::headers::Rav1dPixelLayout;
 use crate::include::dav1d::headers::Rav1dRestorationType;
 use crate::include::dav1d::headers::Rav1dSequenceHeader;
+use crate::include::dav1d::headers::Rav1dTxfmMode;
 use crate::include::dav1d::headers::Rav1dWarpedMotionParams;
 use crate::include::dav1d::headers::RAV1D_FILTER_8TAP_REGULAR;
 use crate::include::dav1d::headers::RAV1D_FILTER_SWITCHABLE;
@@ -24,7 +24,6 @@ use crate::include::dav1d::headers::RAV1D_RESTORATION_NONE;
 use crate::include::dav1d::headers::RAV1D_RESTORATION_SGRPROJ;
 use crate::include::dav1d::headers::RAV1D_RESTORATION_SWITCHABLE;
 use crate::include::dav1d::headers::RAV1D_RESTORATION_WIENER;
-use crate::include::dav1d::headers::RAV1D_TX_SWITCHABLE;
 use crate::include::dav1d::headers::RAV1D_WM_TYPE_AFFINE;
 use crate::include::dav1d::headers::RAV1D_WM_TYPE_IDENTITY;
 use crate::include::dav1d::headers::RAV1D_WM_TYPE_TRANSLATION;
@@ -1042,14 +1041,14 @@ unsafe fn read_vartx_tree(
     let mut tx_split = [0u16; 2];
     *b.max_ytx_mut() = dav1d_max_txfm_size_for_bs[bs as usize][0];
     let frame_hdr = &***f.frame_hdr.as_ref().unwrap();
-    let txfm_mode = frame_hdr.txfm_mode as Dav1dTxfmMode;
+    let txfm_mode = frame_hdr.txfm_mode;
     if b.skip == 0
         && (frame_hdr.segmentation.lossless[b.seg_id as usize] != 0
             || b.max_ytx() as TxfmSize == TX_4X4)
     {
         b.uvtx = TX_4X4 as u8;
         *b.max_ytx_mut() = b.uvtx;
-        if txfm_mode == RAV1D_TX_SWITCHABLE {
+        if txfm_mode == Rav1dTxfmMode::Switchable {
             CaseSet::<32, false>::many(
                 [&mut t.l, &mut *t.a],
                 [bh4 as usize, bw4 as usize],
@@ -1059,8 +1058,8 @@ unsafe fn read_vartx_tree(
                 },
             );
         }
-    } else if txfm_mode != RAV1D_TX_SWITCHABLE || b.skip != 0 {
-        if txfm_mode == RAV1D_TX_SWITCHABLE {
+    } else if txfm_mode != Rav1dTxfmMode::Switchable || b.skip != 0 {
+        if txfm_mode == Rav1dTxfmMode::Switchable {
             CaseSet::<32, false>::many(
                 [(&mut t.l, 1), (&mut *t.a, 0)],
                 [bh4 as usize, bw4 as usize],
@@ -2140,7 +2139,7 @@ unsafe fn decode_b_inner(
             *b.tx_mut() = dav1d_max_txfm_size_for_bs[bs as usize][0];
             b.uvtx = dav1d_max_txfm_size_for_bs[bs as usize][f.cur.p.layout as usize];
             let mut t_dim = &dav1d_txfm_dimensions[b.tx() as usize];
-            if f.frame_hdr().txfm_mode == RAV1D_TX_SWITCHABLE && t_dim.max > TX_4X4 as u8 {
+            if f.frame_hdr().txfm_mode == Rav1dTxfmMode::Switchable && t_dim.max > TX_4X4 as u8 {
                 let tctx = get_tx_ctx(&*t.a, &t.l, &*t_dim, by4, bx4);
                 let tx_cdf = &mut ts.cdf.m.txsz[(t_dim.max - 1) as usize][tctx as usize];
                 let depth = rav1d_msac_decode_symbol_adapt4(

diff --git a/src/obu.rs b/src/obu.rs
@@ -39,6 +39,7 @@ use crate::include::dav1d::headers::Rav1dSequenceHeader;
 use crate::include::dav1d::headers::Rav1dSequenceHeaderOperatingParameterInfo;
 use crate::include::dav1d::headers::Rav1dSequenceHeaderOperatingPoint;
 use crate::include::dav1d::headers::Rav1dTransferCharacteristics;
+use crate::include::dav1d::headers::Rav1dTxfmMode;
 use crate::include::dav1d::headers::Rav1dWarpedMotionParams;
 use crate::include::dav1d::headers::RAV1D_ADAPTIVE;
 use crate::include::dav1d::headers::RAV1D_CHR_UNKNOWN;
@@ -56,9 +57,6 @@ use crate::include::dav1d::headers::RAV1D_REFS_PER_FRAME;
 use crate::include::dav1d::headers::RAV1D_RESTORATION_NONE;
 use crate::include::dav1d::headers::RAV1D_TRC_SRGB;
 use crate::include::dav1d::headers::RAV1D_TRC_UNKNOWN;
-use crate::include::dav1d::headers::RAV1D_TX_4X4_ONLY;
-use crate::include::dav1d::headers::RAV1D_TX_LARGEST;
-use crate::include::dav1d::headers::RAV1D_TX_SWITCHABLE;
 use crate::include::dav1d::headers::RAV1D_WM_TYPE_AFFINE;
 use crate::include::dav1d::headers::RAV1D_WM_TYPE_IDENTITY;
 use crate::include::dav1d::headers::RAV1D_WM_TYPE_ROT_ZOOM;
@@ -1959,11 +1957,11 @@ unsafe fn parse_frame_hdr(
     );
 
     let txfm_mode = if all_lossless != 0 {
-        RAV1D_TX_4X4_ONLY
+        Rav1dTxfmMode::Only4x4
     } else if gb.get_bit() {
-        RAV1D_TX_SWITCHABLE
+        Rav1dTxfmMode::Switchable
     } else {
-        RAV1D_TX_LARGEST
+        Rav1dTxfmMode::Largest
     };
     debug.post(gb, "txfmmode");
     let switchable_comp_refs = if frame_type.is_inter_or_switch() {