Rav1dFrameContext_frame_thread::cf: Make into AlignedVec64

memorysafety · Feb 17, 2024 · c4ad9be · c4ad9be
1 parent 2014feb
commit c4ad9be
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 32 deletions.
diff --git a/src/decode.rs b/src/decode.rs
@@ -3868,14 +3868,11 @@ unsafe fn setup_tile(
         } else {
             ptr::null_mut()
         };
-        ts.frame_thread[p].cf = if !f.frame_thread.cf.is_null() {
-            f.frame_thread
-                .cf
-                .cast::<u8>()
-                .offset(
-                    (tile_start_off * size_mul[0] as usize >> (seq_hdr.hbd == 0) as c_int) as isize,
-                )
-                .cast::<DynCoef>()
+        ts.frame_thread[p].cf = if !f.frame_thread.cf.is_empty() {
+            f.frame_thread.cf
+                [(tile_start_off * size_mul[0] as usize >> (seq_hdr.hbd == 0) as c_int) as usize..]
+                .as_ptr()
+                .cast::<DynCoef>() as *mut _
         } else {
             ptr::null_mut()
         };
@@ -4371,20 +4368,12 @@ pub(crate) unsafe fn rav1d_decode_frame_init(
         }
 
         let cf_sz = (num_sb128 * size_mul[0] as c_int) << hbd;
-        if cf_sz != f.frame_thread.cf_sz {
-            rav1d_freep_aligned(&mut f.frame_thread.cf as *mut *mut DynCoef as *mut c_void);
-            f.frame_thread.cf =
-                rav1d_alloc_aligned(cf_sz as usize * 128 * 128 / 2, 64) as *mut DynCoef;
-            if f.frame_thread.cf.is_null() {
-                f.frame_thread.cf_sz = 0;
-                return Err(ENOMEM);
-            }
-            slice::from_raw_parts_mut(
-                f.frame_thread.cf.cast::<u8>(),
-                usize::try_from(cf_sz).unwrap() * 128 * 128 / 2,
-            )
-            .fill(0);
-            f.frame_thread.cf_sz = cf_sz;
+        if cf_sz as usize != f.frame_thread.cf_sz() {
+            f.frame_thread.cf.clear();
+            // TODO: Fallible allocation
+            f.frame_thread
+                .cf
+                .resize_with(cf_sz as usize * 128 * 128 / 2, Default::default);
         }
 
         if frame_hdr.allow_screen_content_tools != 0 {
@@ -4861,12 +4850,8 @@ pub(crate) unsafe fn rav1d_decode_frame_exit(
     if !f.sr_cur.p.data.data[0].is_null() {
         f.task_thread.error = AtomicI32::new(0);
     }
-    if c.n_fc > 1 && retval.is_err() && !f.frame_thread.cf.is_null() {
-        slice::from_raw_parts_mut(
-            f.frame_thread.cf.cast::<u8>(),
-            usize::try_from(f.frame_thread.cf_sz).unwrap() * 128 * 128 / 2,
-        )
-        .fill(0);
+    if c.n_fc > 1 && retval.is_err() && !f.frame_thread.cf.is_empty() {
+        f.frame_thread.cf.fill_with(Default::default)
     }
     // TODO(kkysen) use array::zip when stable
     for i in 0..7 {

diff --git a/src/internal.rs b/src/internal.rs
@@ -412,8 +412,7 @@ pub struct Rav1dFrameContext_frame_thread {
     pub pal: AlignedVec64<[[u16; 8]; 3]>, /* [3 plane][8 idx] */
     // iterated over inside tile state
     pub pal_idx: AlignedVec64<u8>,
-    pub cf: *mut DynCoef,
-    pub cf_sz: c_int,
+    pub cf: AlignedVec64<u8>, // *mut DynCoef
     // start offsets per tile
     pub tile_start_off: *mut c_int,
 }
@@ -422,6 +421,10 @@ impl Rav1dFrameContext_frame_thread {
     pub fn pal_sz(&self) -> usize {
         self.pal.len() / (16 * 16)
     }
+
+    pub fn cf_sz(&self) -> usize {
+        self.cf.len() / (128 * 128 / 2)
+    }
 }
 
 /// loopfilter

diff --git a/src/lib.rs b/src/lib.rs
@@ -1,6 +1,5 @@
 use crate::include::common::bitdepth::BitDepth16;
 use crate::include::common::bitdepth::BitDepth8;
-use crate::include::common::bitdepth::DynCoef;
 use crate::include::common::validate::validate_input;
 use crate::include::dav1d::common::Dav1dDataProps;
 use crate::include::dav1d::common::Rav1dDataProps;
@@ -907,7 +906,7 @@ impl Drop for Rav1dContext {
                     let _ = mem::take(&mut f.frame_thread.b); // TODO: remove when context is owned
                     let _ =
                         mem::replace(&mut f.frame_thread.pal_idx, Vec::new_in(AlignedAlloc::<64>)); // TODO: remove when context is owned
-                    rav1d_freep_aligned(&mut f.frame_thread.cf as *mut *mut DynCoef as *mut c_void);
+                    let _ = mem::replace(&mut f.frame_thread.cf, Vec::new_in(AlignedAlloc::<64>)); // TODO: remove when context is owned
                     freep(&mut f.frame_thread.tile_start_off as *mut *mut c_int as *mut c_void);
                     let _ = mem::replace(&mut f.frame_thread.pal, Vec::new_in(AlignedAlloc::<64>)); // TODO: remove when context is owned
                     let _ = mem::take(&mut f.frame_thread.cbi); // TODO: remove when context is owned