struct Av1Restoration: Put Av1RestorationUnit in a Cell (#766)

memorysafety · Feb 27, 2024 · 52b3c32 · 52b3c32
2 parents 1eb15d1 + 005e191
commit 52b3c32
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 11 deletions.
diff --git a/src/decode.rs b/src/decode.rs
@@ -3925,18 +3925,20 @@ unsafe fn setup_tile(
             if sb128x >= f.sr_sb128w {
                 continue;
             }
-            &mut (*f.lf.lr_mask.offset((sb_idx + sb128x) as isize)).lr[p][u_idx as usize]
+            &(*f.lf.lr_mask.offset((sb_idx + sb128x) as isize)).lr[p][u_idx as usize]
         } else {
-            &mut (*f.lf.lr_mask.offset(sb_idx as isize)).lr[p][unit_idx as usize]
+            &(*f.lf.lr_mask.offset(sb_idx as isize)).lr[p][unit_idx as usize]
         };
 
-        *lr_ref = Av1RestorationUnit {
+        let lr = lr_ref.get();
+        let lr = Av1RestorationUnit {
             filter_v: [3, -7, 15],
             filter_h: [3, -7, 15],
             sgr_weights: [-32, 31],
-            ..*lr_ref
+            ..lr
         };
-        ts.lr_ref[p] = *lr_ref;
+        lr_ref.set(lr);
+        ts.lr_ref[p] = lr;
     }
 
     if c.tc.len() > 1 {
@@ -4185,8 +4187,8 @@ pub(crate) unsafe fn rav1d_decode_tile_sbrow(
                     let px_x = x << unit_size_log2 + ss_hor;
                     let sb_idx = (t.by >> 5) * f.sr_sb128w + (px_x >> 7);
                     let unit_idx = ((t.by & 16) >> 3) + ((px_x & 64) >> 6);
-                    let lr =
-                        &mut (*(f.lf.lr_mask).offset(sb_idx as isize)).lr[p][unit_idx as usize];
+                    let lr = (*(f.lf.lr_mask).offset(sb_idx as isize)).lr[p][unit_idx as usize]
+                        .get_mut();
 
                     read_restoration_info(t, f, lr, p, frame_type);
                 }
@@ -4203,7 +4205,8 @@ pub(crate) unsafe fn rav1d_decode_tile_sbrow(
                 }
                 let sb_idx = (t.by >> 5) * f.sr_sb128w + (t.bx >> 5);
                 let unit_idx = ((t.by & 16) >> 3) + ((t.bx & 16) >> 4);
-                let lr = &mut (*(f.lf.lr_mask).offset(sb_idx as isize)).lr[p][unit_idx as usize];
+                let lr =
+                    (*(f.lf.lr_mask).offset(sb_idx as isize)).lr[p][unit_idx as usize].get_mut();
 
                 read_restoration_info(t, f, lr, p, frame_type);
             }

diff --git a/src/lf_mask.rs b/src/lf_mask.rs
@@ -11,6 +11,7 @@ use crate::src::levels::TX_4X4;
 use crate::src::tables::dav1d_block_dimensions;
 use crate::src::tables::dav1d_txfm_dimensions;
 use libc::ptrdiff_t;
+use std::cell::Cell;
 use std::cmp;
 use std::ffi::c_int;
 
@@ -42,7 +43,7 @@ pub struct Av1Filter {
 
 #[repr(C)]
 pub struct Av1Restoration {
-    pub lr: [[Av1RestorationUnit; 4]; 3],
+    pub lr: [[Cell<Av1RestorationUnit>; 4]; 3],
 }
 
 /// In `txa`, the array lengths represent from inner to outer:

diff --git a/src/lr_apply.rs b/src/lr_apply.rs
@@ -182,15 +182,16 @@ unsafe fn lr_sbrow<BD: BitDepth>(
     aligned_unit_pos <<= ss_ver;
     let sb_idx = (aligned_unit_pos >> 7) * f.sr_sb128w;
     let unit_idx = (aligned_unit_pos >> 6 & 1) << 1;
-    lr[0] = (*(f.lf.lr_mask).offset(sb_idx as isize)).lr[plane as usize][unit_idx as usize];
+    lr[0] = (*(f.lf.lr_mask).offset(sb_idx as isize)).lr[plane as usize][unit_idx as usize].get();
     let mut restore = lr[0].r#type != RAV1D_RESTORATION_NONE;
     let mut x = 0;
     let mut bit = false;
     while x + max_unit_size <= w {
         let next_x = x + unit_size;
         let next_u_idx = unit_idx + (next_x >> shift_hor - 1 & 1);
         lr[!bit as usize] = (*(f.lf.lr_mask).offset((sb_idx + (next_x >> shift_hor)) as isize)).lr
-            [plane as usize][next_u_idx as usize];
+            [plane as usize][next_u_idx as usize]
+            .get();
         let restore_next = lr[!bit as usize].r#type != RAV1D_RESTORATION_NONE;
         if restore_next {
             backup4xU::<BD>(