struct EdgeFlags: Make bitflags!

src/decode.rs

@@ -81,7 +81,6 @@ use crate::src::internal::ScalableMotionParams;
 use crate::src::intra_edge::EdgeFlags;
 use crate::src::intra_edge::EdgeIndex;
 use crate::src::intra_edge::IntraEdges;
-use crate::src::intra_edge::EDGE_I444_TOP_HAS_RIGHT;
 use crate::src::ipred::rav1d_intra_pred_dsp_init;
 use crate::src::levels::mv;
 use crate::src::levels::Av1Block;
@@ -466,7 +465,7 @@ unsafe fn find_matching_ref(
     let mut have_topright = cmp::max(bw4, bh4) < 32
         && have_top
         && t.bx + bw4 < (*t.ts).tiling.col_end
-        && intra_edge_flags & EDGE_I444_TOP_HAS_RIGHT != 0;
+        && intra_edge_flags.contains(EdgeFlags::I444_TOP_HAS_RIGHT);
 
     let bs = |rp: &refmvs_block| dav1d_block_dimensions[rp.0.bs as usize];
     let matches = |rp: &refmvs_block| rp.0.r#ref.r#ref[0] == r#ref + 1 && rp.0.r#ref.r#ref[1] == -1;

src/intra_edge.rs

@@ -1,22 +1,67 @@
+use crate::include::dav1d::headers::Rav1dPixelLayout;
 use crate::src::enum_map::DefaultValue;
 use crate::src::levels::BlockLevel;
 use crate::src::levels::BL_128X128;
 use crate::src::levels::BL_16X16;
 use crate::src::levels::BL_32X32;
 use crate::src::levels::BL_64X64;
+use bitflags::bitflags;
+use std::ops::Shr;
+
+bitflags! {
+    #[derive(Clone, Copy, PartialEq, Eq, Hash, Default)]
+    pub(crate) struct EdgeFlags: u8 {
+        const I444_TOP_HAS_RIGHT = 1 << 0;
+        const I422_TOP_HAS_RIGHT = 1 << 1;
+        const I420_TOP_HAS_RIGHT = 1 << 2;
+
+        const I444_LEFT_HAS_BOTTOM = 1 << 3;
+        const I422_LEFT_HAS_BOTTOM = 1 << 4;
+        const I420_LEFT_HAS_BOTTOM = 1 << 5;
+    }
+}
+
+impl EdgeFlags {
+    const LEFT_HAS_BOTTOM: Self = Self::union_all([
+        Self::I444_LEFT_HAS_BOTTOM,
+        Self::I422_LEFT_HAS_BOTTOM,
+        Self::I420_LEFT_HAS_BOTTOM,
+    ]);
+
+    const TOP_HAS_RIGHT: Self = Self::union_all([
+        Self::I444_TOP_HAS_RIGHT,
+        Self::I422_TOP_HAS_RIGHT,
+        Self::I420_TOP_HAS_RIGHT,
+    ]);
+
+    pub const fn union_all<const N: usize>(flags: [Self; N]) -> Self {
+        let mut i = 0;
+        let mut output = Self::empty();
 
-pub type EdgeFlags = u8;
-pub const EDGE_I420_LEFT_HAS_BOTTOM: EdgeFlags = 32;
-pub const EDGE_I422_LEFT_HAS_BOTTOM: EdgeFlags = 16;
-pub const EDGE_I444_LEFT_HAS_BOTTOM: EdgeFlags = 8;
-pub const EDGE_I420_TOP_HAS_RIGHT: EdgeFlags = 4;
-pub const EDGE_I422_TOP_HAS_RIGHT: EdgeFlags = 2;
-pub const EDGE_I444_TOP_HAS_RIGHT: EdgeFlags = 1;
+        while i < N {
+            output = output.union(flags[i]);
+            i += 1;
+        }
 
-pub const EDGE_LEFT_HAS_BOTTOM: EdgeFlags =
-    EDGE_I444_LEFT_HAS_BOTTOM | EDGE_I422_LEFT_HAS_BOTTOM | EDGE_I420_LEFT_HAS_BOTTOM;
-pub const EDGE_TOP_HAS_RIGHT: EdgeFlags =
-    EDGE_I444_TOP_HAS_RIGHT | EDGE_I422_TOP_HAS_RIGHT | EDGE_I420_TOP_HAS_RIGHT;
+        output
+    }
+
+    pub(crate) const fn select(&self, select: bool) -> Self {
+        if select {
+            *self
+        } else {
+            Self::empty()
+        }
+    }
+}
+
+impl Shr<Rav1dPixelLayout> for EdgeFlags {
+    type Output = Self;
+
+    fn shr(self, rhs: Rav1dPixelLayout) -> Self::Output {
+        Self::from_bits_retain(self.bits() >> (rhs as u32).wrapping_sub(1))
+    }
+}
 
 const B: usize = 4;
 
@@ -77,22 +122,33 @@ impl EdgeTip {
     const fn new(edge_flags: EdgeFlags) -> Self {
         let o = edge_flags;
         let h = [
-            edge_flags | EDGE_LEFT_HAS_BOTTOM,
-            edge_flags & (EDGE_LEFT_HAS_BOTTOM | EDGE_I420_TOP_HAS_RIGHT),
+            edge_flags.union(EdgeFlags::LEFT_HAS_BOTTOM),
+            edge_flags.intersection(EdgeFlags::union_all([
+                EdgeFlags::LEFT_HAS_BOTTOM,
+                EdgeFlags::I420_TOP_HAS_RIGHT,
+            ])),
         ];
         let v = [
-            edge_flags | EDGE_TOP_HAS_RIGHT,
-            edge_flags
-                & (EDGE_TOP_HAS_RIGHT | EDGE_I420_LEFT_HAS_BOTTOM | EDGE_I422_LEFT_HAS_BOTTOM),
+            edge_flags.union(EdgeFlags::TOP_HAS_RIGHT),
+            edge_flags.intersection(EdgeFlags::union_all([
+                EdgeFlags::TOP_HAS_RIGHT,
+                EdgeFlags::I420_LEFT_HAS_BOTTOM,
+                EdgeFlags::I422_LEFT_HAS_BOTTOM,
+            ])),
         ];
         let node = EdgeNode { o, h, v };
 
         let split = [
-            EDGE_TOP_HAS_RIGHT | EDGE_LEFT_HAS_BOTTOM,
-            (edge_flags & EDGE_TOP_HAS_RIGHT) | EDGE_I422_LEFT_HAS_BOTTOM,
-            edge_flags | EDGE_I444_TOP_HAS_RIGHT,
+            EdgeFlags::all(),
             edge_flags
-                & (EDGE_I420_TOP_HAS_RIGHT | EDGE_I420_LEFT_HAS_BOTTOM | EDGE_I422_LEFT_HAS_BOTTOM),
+                .intersection(EdgeFlags::TOP_HAS_RIGHT)
+                .union(EdgeFlags::I422_LEFT_HAS_BOTTOM),
+            edge_flags.union(EdgeFlags::I444_TOP_HAS_RIGHT),
+            edge_flags.intersection(EdgeFlags::union_all([
+                EdgeFlags::I420_TOP_HAS_RIGHT,
+                EdgeFlags::I420_LEFT_HAS_BOTTOM,
+                EdgeFlags::I422_LEFT_HAS_BOTTOM,
+            ])),
         ];
 
         Self { node, split }
@@ -103,58 +159,59 @@ impl EdgeBranch {
     const fn new(edge_flags: EdgeFlags, bl: BlockLevel) -> Self {
         let o = edge_flags;
         let h = [
-            edge_flags | EDGE_LEFT_HAS_BOTTOM,
-            edge_flags & EDGE_LEFT_HAS_BOTTOM,
+            edge_flags.union(EdgeFlags::LEFT_HAS_BOTTOM),
+            edge_flags.intersection(EdgeFlags::LEFT_HAS_BOTTOM),
         ];
         let v = [
-            edge_flags | EDGE_TOP_HAS_RIGHT,
-            edge_flags & EDGE_TOP_HAS_RIGHT,
+            edge_flags.union(EdgeFlags::TOP_HAS_RIGHT),
+            edge_flags.intersection(EdgeFlags::TOP_HAS_RIGHT),
         ];
         let node = EdgeNode { o, h, v };
 
         let h4 = [
-            edge_flags | EDGE_LEFT_HAS_BOTTOM,
-            EDGE_LEFT_HAS_BOTTOM
-                | (if bl == BL_16X16 {
-                    edge_flags & EDGE_I420_TOP_HAS_RIGHT
-                } else {
-                    0 as EdgeFlags
-                }),
-            EDGE_LEFT_HAS_BOTTOM,
-            edge_flags & EDGE_LEFT_HAS_BOTTOM,
+            edge_flags.union(EdgeFlags::LEFT_HAS_BOTTOM),
+            EdgeFlags::LEFT_HAS_BOTTOM.union(
+                edge_flags
+                    .intersection(EdgeFlags::I420_TOP_HAS_RIGHT)
+                    .select(bl == BL_16X16),
+            ),
+            EdgeFlags::LEFT_HAS_BOTTOM,
+            edge_flags.intersection(EdgeFlags::LEFT_HAS_BOTTOM),
         ];
 
         let v4 = [
-            edge_flags | EDGE_TOP_HAS_RIGHT,
-            EDGE_TOP_HAS_RIGHT
-                | (if bl == BL_16X16 {
-                    edge_flags & (EDGE_I420_LEFT_HAS_BOTTOM | EDGE_I422_LEFT_HAS_BOTTOM)
-                } else {
-                    0 as EdgeFlags
-                }),
-            EDGE_TOP_HAS_RIGHT,
-            edge_flags & EDGE_TOP_HAS_RIGHT,
+            edge_flags.union(EdgeFlags::TOP_HAS_RIGHT),
+            EdgeFlags::TOP_HAS_RIGHT.union(
+                edge_flags
+                    .intersection(EdgeFlags::union_all([
+                        EdgeFlags::I420_LEFT_HAS_BOTTOM,
+                        EdgeFlags::I422_LEFT_HAS_BOTTOM,
+                    ]))
+                    .select(bl == BL_16X16),
+            ),
+            EdgeFlags::TOP_HAS_RIGHT,
+            edge_flags.intersection(EdgeFlags::TOP_HAS_RIGHT),
         ];
 
         let tls = [
-            EDGE_TOP_HAS_RIGHT | EDGE_LEFT_HAS_BOTTOM,
-            edge_flags & EDGE_LEFT_HAS_BOTTOM,
-            edge_flags & EDGE_TOP_HAS_RIGHT,
+            EdgeFlags::all(),
+            edge_flags.intersection(EdgeFlags::LEFT_HAS_BOTTOM),
+            edge_flags.intersection(EdgeFlags::TOP_HAS_RIGHT),
         ];
         let trs = [
-            edge_flags | EDGE_TOP_HAS_RIGHT,
-            edge_flags | EDGE_LEFT_HAS_BOTTOM,
-            0 as EdgeFlags,
+            edge_flags.union(EdgeFlags::TOP_HAS_RIGHT),
+            edge_flags.union(EdgeFlags::LEFT_HAS_BOTTOM),
+            EdgeFlags::empty(),
         ];
         let tts = [
-            EDGE_TOP_HAS_RIGHT | EDGE_LEFT_HAS_BOTTOM,
-            edge_flags & EDGE_TOP_HAS_RIGHT,
-            edge_flags & EDGE_LEFT_HAS_BOTTOM,
+            EdgeFlags::all(),
+            edge_flags.intersection(EdgeFlags::TOP_HAS_RIGHT),
+            edge_flags.intersection(EdgeFlags::LEFT_HAS_BOTTOM),
         ];
         let tbs = [
-            edge_flags | EDGE_LEFT_HAS_BOTTOM,
-            edge_flags | EDGE_TOP_HAS_RIGHT,
-            0 as EdgeFlags,
+            edge_flags.union(EdgeFlags::LEFT_HAS_BOTTOM),
+            edge_flags.union(EdgeFlags::TOP_HAS_RIGHT),
+            EdgeFlags::empty(),
         ];
 
         let split = [EdgeIndex::root(); 4];
@@ -173,11 +230,11 @@ impl EdgeBranch {
 }
 
 impl DefaultValue for EdgeTip {
-    const DEFAULT: Self = Self::new(0 as EdgeFlags);
+    const DEFAULT: Self = Self::new(EdgeFlags::empty());
 }
 
 impl DefaultValue for EdgeBranch {
-    const DEFAULT: Self = Self::new(0 as EdgeFlags, 0 as BlockLevel);
+    const DEFAULT: Self = Self::new(EdgeFlags::empty(), 0 as BlockLevel);
 }
 
 struct EdgeIndices {
@@ -215,15 +272,10 @@ impl<const SB128: bool, const N_BRANCH: usize, const N_TIP: usize>
         left_has_bottom: bool,
     ) -> (Self, EdgeIndices) {
         let mut branch = EdgeBranch::new(
-            (if top_has_right {
-                EDGE_TOP_HAS_RIGHT
-            } else {
-                0 as EdgeFlags
-            }) | (if left_has_bottom {
-                EDGE_LEFT_HAS_BOTTOM
-            } else {
-                0 as EdgeFlags
-            }),
+            EdgeFlags::union_all([
+                EdgeFlags::TOP_HAS_RIGHT.select(top_has_right),
+                EdgeFlags::LEFT_HAS_BOTTOM.select(left_has_bottom),
+            ]),
             bl,
         );
         if bl == BL_16X16 {
@@ -232,16 +284,10 @@ impl<const SB128: bool, const N_BRANCH: usize, const N_TIP: usize>
                 let (tip, next) = indices.tip.pop_front();
                 indices.tip = next;
                 branch.split[n as usize] = tip;
-                let edge_flags = (if n == 3 || (n == 1 && !top_has_right) {
-                    0 as EdgeFlags
-                } else {
-                    EDGE_TOP_HAS_RIGHT
-                }) | (if !(n == 0 || (n == 2 && left_has_bottom)) {
-                    0 as EdgeFlags
-                } else {
-                    EDGE_LEFT_HAS_BOTTOM
-                });
-                self.tip[tip.index as usize] = EdgeTip::new(edge_flags);
+                self.tip[tip.index as usize] = EdgeTip::new(EdgeFlags::union_all([
+                    EdgeFlags::TOP_HAS_RIGHT.select(!(n == 3 || (n == 1 && !top_has_right))),
+                    EdgeFlags::LEFT_HAS_BOTTOM.select(n == 0 || (n == 2 && left_has_bottom)),
+                ]));
                 n += 1;
             }
         } else {

src/ipred_prepare.rs

@@ -3,8 +3,6 @@ use crate::include::common::bitdepth::BitDepth;
 use crate::src::const_fn::const_for;
 use crate::src::env::BlockContext;
 use crate::src::intra_edge::EdgeFlags;
-use crate::src::intra_edge::EDGE_I444_LEFT_HAS_BOTTOM;
-use crate::src::intra_edge::EDGE_I444_TOP_HAS_RIGHT;
 use crate::src::levels::IntraPredMode;
 use crate::src::levels::DC_128_PRED;
 use crate::src::levels::DC_PRED;
@@ -228,7 +226,7 @@ pub fn rav1d_prepare_intra_edges<BD: BitDepth>(
             let have_bottomleft = if !have_left || y + th >= h {
                 false
             } else {
-                (edge_flags & EDGE_I444_LEFT_HAS_BOTTOM) != 0
+                edge_flags.contains(EdgeFlags::I444_LEFT_HAS_BOTTOM)
             };
             if have_bottomleft {
                 let px_have = cmp::min(sz, (h - y - th << 2) as usize);
@@ -274,7 +272,7 @@ pub fn rav1d_prepare_intra_edges<BD: BitDepth>(
             let have_topright = if !have_top || x + tw >= w {
                 false
             } else {
-                (edge_flags & EDGE_I444_TOP_HAS_RIGHT) != 0
+                edge_flags.contains(EdgeFlags::I444_TOP_HAS_RIGHT)
             };
             if have_topright {
                 let top_right = &mut top[sz..];