fn rav1d_apply_grain_row: Make safe (#1127)

memorysafety · May 28, 2024 · decb534 · decb534
2 parents 9be05a2 + 3093ff8
commit decb534
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 27 deletions.
diff --git a/src/fg_apply.rs b/src/fg_apply.rs
@@ -129,7 +129,7 @@ pub(crate) unsafe fn rav1d_prep_grain<BD: BitDepth>(
     }
 }
 
-pub(crate) unsafe fn rav1d_apply_grain_row<BD: BitDepth>(
+pub(crate) fn rav1d_apply_grain_row<BD: BitDepth>(
     dsp: &Rav1dFilmGrainDSPContext,
     out: &Rav1dPicture,
     r#in: &Rav1dPicture,
@@ -151,9 +151,6 @@ pub(crate) unsafe fn rav1d_apply_grain_row<BD: BitDepth>(
     let ss_x = (r#in.p.layout != Rav1dPixelLayout::I444) as usize;
     let cpw = w + ss_x >> ss_x;
     let is_id = seq_hdr.mtrx == Rav1dMatrixCoefficients::IDENTITY;
-    let luma_src = in_data[0]
-        .as_strided_mut_ptr::<BD>()
-        .offset((row * BLOCK_SIZE) as isize * BD::pxstride(r#in.stride[0]));
     let bitdepth_max = (1 << out.p.bpc) - 1;
     let bd = BD::from_c(bitdepth_max);
 
@@ -180,10 +177,14 @@ pub(crate) unsafe fn rav1d_apply_grain_row<BD: BitDepth>(
 
     // extend padding pixels
     if out.p.w as usize & ss_x != 0 {
-        let mut ptr = luma_src;
+        let luma = &in_data[0];
+        let mut offset = luma
+            .pixel_offset::<BD>()
+            .wrapping_add_signed((row * BLOCK_SIZE) as isize * luma.pixel_stride::<BD>());
         for _ in 0..bh {
-            *ptr.add(out.p.w as usize) = *ptr.add((out.p.w - 1) as usize);
-            ptr = ptr.offset(BD::pxstride(r#in.stride[0]) << ss_y);
+            let padding = &mut *luma.slice_mut::<BD, _>((offset + out.p.w as usize - 1.., ..2));
+            padding[1] = padding[0];
+            offset = offset.wrapping_add_signed(luma.pixel_stride::<BD>() << ss_y);
         }
     }
 

diff --git a/src/thread_task.rs b/src/thread_task.rs
@@ -751,29 +751,23 @@ fn delayed_fg_task<'l, 'ttd: 'l>(
             match &delayed_fg.grain {
                 #[cfg(feature = "bitdepth_8")]
                 Grain::Bpc8(grain) => {
-                    // SAFETY: TODO make safe
-                    unsafe {
-                        rav1d_apply_grain_row::<BitDepth8>(
-                            dsp,
-                            &delayed_fg.out,
-                            &delayed_fg.in_0,
-                            grain,
-                            row as usize,
-                        );
-                    }
+                    rav1d_apply_grain_row::<BitDepth8>(
+                        dsp,
+                        &delayed_fg.out,
+                        &delayed_fg.in_0,
+                        grain,
+                        row as usize,
+                    );
                 }
                 #[cfg(feature = "bitdepth_16")]
                 Grain::Bpc16(grain) => {
-                    // SAFETY: TODO make safe
-                    unsafe {
-                        rav1d_apply_grain_row::<BitDepth16>(
-                            dsp,
-                            &delayed_fg.out,
-                            &delayed_fg.in_0,
-                            grain,
-                            row as usize,
-                        );
-                    }
+                    rav1d_apply_grain_row::<BitDepth16>(
+                        dsp,
+                        &delayed_fg.out,
+                        &delayed_fg.in_0,
+                        grain,
+                        row as usize,
+                    );
                 }
             }
         }