From cee12e492853e71e53a2cc120a71cc51e4fbde57 Mon Sep 17 00:00:00 2001
From: Nicole LeGare <dlegare.1001@gmail.com>
Date: Fri, 9 Feb 2024 14:36:15 -0800
Subject: [PATCH] `Rav1dFrameContext_frame_thread::cf`: Make into
 `AlignedVec64`

---
 src/align.rs    |  5 +++++
 src/decode.rs   | 39 +++++++++++----------------------------
 src/internal.rs |  9 +++++++--
 src/lib.rs      |  3 +--
 4 files changed, 24 insertions(+), 32 deletions(-)

diff --git a/src/align.rs b/src/align.rs
index 0cc79a08d..b3e260959 100644
--- a/src/align.rs
+++ b/src/align.rs
@@ -168,6 +168,11 @@ impl<T: Copy> AlignedVec64<T> {
 
         self.len = new_len;
     }
+
+    pub fn clear(&mut self) {
+        self.inner.clear();
+        self.len = 0;
+    }
 }
 
 impl<T: Copy> Deref for AlignedVec64<T> {
diff --git a/src/decode.rs b/src/decode.rs
index acfd722fa..8083182c8 100644
--- a/src/decode.rs
+++ b/src/decode.rs
@@ -3865,14 +3865,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()
         };
@@ -4364,20 +4361,10 @@ 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(cf_sz as usize * 128 * 128 / 2, 0);
         }
 
         if frame_hdr.allow_screen_content_tools != 0 {
@@ -4852,12 +4839,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
index c9a1fc192..4a52aee44 100644
--- a/src/internal.rs
+++ b/src/internal.rs
@@ -422,12 +422,17 @@ 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 u32,
 }
 
+impl Rav1dFrameContext_frame_thread {
+    pub fn cf_sz(&self) -> usize {
+        self.cf.len() / (128 * 128 / 2)
+    }
+}
+
 /// loopfilter
 #[repr(C)]
 pub struct Rav1dFrameContext_lf {
diff --git a/src/lib.rs b/src/lib.rs
index 3fb2274cb..d3d06a73c 100644
--- 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;
@@ -905,7 +904,7 @@ impl Drop for Rav1dContext {
                     );
                     let _ = mem::take(&mut f.frame_thread.b); // TODO: remove when context is owned
                     let _ = mem::take(&mut f.frame_thread.pal_idx); // TODO: remove when context is owned
-                    rav1d_freep_aligned(&mut f.frame_thread.cf as *mut *mut DynCoef as *mut c_void);
+                    let _ = mem::take(&mut f.frame_thread.cf); // TODO: remove when context is owned
                     freep(&mut f.frame_thread.tile_start_off as *mut *mut u32 as *mut c_void);
                     let _ = mem::take(&mut f.frame_thread.pal); // TODO: remove when context is owned
                     let _ = mem::take(&mut f.frame_thread.cbi); // TODO: remove when context is owned