fn splat_dc: Use zerocopy to make the slice casting safe (#1200)

This is what I meant by #1105 (review). Also, I included a couple of other random little changes.
memorysafety · Jun 18, 2024 · d0a3332 · d0a3332
2 parents 231e81e + ea357e3
commit d0a3332
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 38 deletions.
diff --git a/src/cdf.rs b/src/cdf.rs
@@ -5097,10 +5097,11 @@ pub fn rav1d_cdf_thread_init_static(qidx: u8) -> CdfThreadContext {
 pub fn rav1d_cdf_thread_copy(src: &CdfThreadContext) -> CdfContext {
     match src {
         CdfThreadContext::Cdf(src) => src.cdf.try_read().unwrap().clone(),
-        CdfThreadContext::QCat(i) => CdfContext {
+        &CdfThreadContext::QCat(i) => CdfContext {
             m: Default::default(),
             kfym: default_kf_y_mode_cdf,
-            coef: av1_default_coef_cdf[*i as usize].clone(),
+            // `i` is the sum of 3 `bool`s
+            coef: av1_default_coef_cdf[i as usize & 3].clone(),
             mv: Default::default(),
         },
     }

diff --git a/src/ipred.rs b/src/ipred.rs
@@ -36,9 +36,10 @@ use libc::ptrdiff_t;
 use std::cmp;
 use std::ffi::c_int;
 use std::ffi::c_uint;
-use std::mem;
 use std::slice;
 use strum::FromRepr;
+use zerocopy::AsBytes;
+use zerocopy::FromBytes;
 
 #[cfg(all(
     feature = "asm",
@@ -214,49 +215,33 @@ pub struct Rav1dIntraPredDSPContext {
 
 #[inline(never)]
 unsafe fn splat_dc<BD: BitDepth>(
-    mut dst: *mut BD::Pixel,
+    dst: *mut BD::Pixel,
     stride: ptrdiff_t,
     width: c_int,
     height: c_int,
     dc: c_int,
     bd: BD,
 ) {
     let stride = BD::pxstride(stride);
+    let height = height as isize;
     let width = width as usize;
-    match BD::BPC {
-        BPC::BPC8 => {
-            assert!(dc <= 0xff);
-            if width > 4 {
-                let dcN = dc as u64 * 0x101010101010101;
-                for _ in 0..height {
-                    let slice =
-                        slice::from_raw_parts_mut(dst.cast::<u64>(), width / mem::size_of::<u64>());
-                    slice.fill(dcN);
-                    dst = dst.offset(stride);
-                }
-            } else {
-                let dcN = dc as u32 * 0x1010101;
-                for _ in 0..height {
-                    let slice =
-                        slice::from_raw_parts_mut(dst.cast::<u32>(), width / mem::size_of::<u32>());
-                    slice.fill(dcN);
-                    dst = dst.offset(stride);
-                }
-            };
+    assert!(dc <= bd.bitdepth_max().as_::<c_int>());
+    let dc = dc.as_::<BD::Pixel>();
+    if BD::BPC == BPC::BPC8 && width > 4 {
+        for y in 0..height {
+            let dst = dst.offset(y * stride);
+            let dst = slice::from_raw_parts_mut(dst, width);
+            let dst = FromBytes::mut_slice_from(AsBytes::as_bytes_mut(dst)).unwrap();
+            dst.fill([dc; 8]);
         }
-        BPC::BPC16 => {
-            assert!(dc <= bd.bitdepth_max().as_::<c_int>());
-            let dcN = dc as u64 * 0x1000100010001;
-            for _ in 0..height {
-                let slice = slice::from_raw_parts_mut(
-                    dst.cast::<u64>(),
-                    width / (mem::size_of::<u64>() >> 1),
-                );
-                slice.fill(dcN);
-                dst = dst.offset(stride);
-            }
+    } else {
+        for y in 0..height {
+            let dst = dst.offset(y * stride);
+            let dst = slice::from_raw_parts_mut(dst, width);
+            let dst = FromBytes::mut_slice_from(AsBytes::as_bytes_mut(dst)).unwrap();
+            dst.fill([dc; 4]);
         }
-    }
+    };
 }
 
 #[inline(never)]

diff --git a/src/lf_apply.rs b/src/lf_apply.rs
@@ -360,7 +360,7 @@ unsafe fn filter_plane_cols_y<BD: BitDepth>(
         if !have_left && x == 0 {
             continue;
         }
-        let mask = &mask[x];
+        let mask = &mask[x % mask.len()]; // To elide the bounds check.;
         let hmask = if starty4 == 0 {
             if endy4 > 16 {
                 mask.each_ref()
@@ -435,7 +435,7 @@ unsafe fn filter_plane_cols_uv<BD: BitDepth>(
         if !have_left && x == 0 {
             continue;
         }
-        let mask = &mask[x];
+        let mask = &mask[x % mask.len()]; // To elide the bounds check.;
         let hmask = if starty4 == 0 {
             if endy4 > 16 >> ss_ver {
                 mask.each_ref()