struct Rav1dFrameData::ipred_edge: Make into DisjointMut<AlignedVec<u8>>

src/decode.rs

@@ -1,5 +1,4 @@
 use crate::include::common::attributes::ctz;
-use crate::include::common::bitdepth::DynPixel;
 use crate::include::common::bitdepth::BPC;
 use crate::include::common::intops::apply_sign64;
 use crate::include::common::intops::iclip;
@@ -115,7 +114,6 @@ use crate::src::lf_mask::Av1RestorationUnit;
 use crate::src::log::Rav1dLog as _;
 use crate::src::mem::rav1d_alloc_aligned;
 use crate::src::mem::rav1d_free_aligned;
-use crate::src::mem::rav1d_freep_aligned;
 use crate::src::msac::rav1d_msac_decode_bool;
 use crate::src::msac::rav1d_msac_decode_bool_adapt;
 use crate::src::msac::rav1d_msac_decode_bool_equi;
@@ -4318,21 +4316,9 @@ pub(crate) unsafe fn rav1d_decode_frame_init(
     rav1d_calc_lf_values(&mut f.lf.lvl, &frame_hdr, &[0, 0, 0, 0]);
 
     let ipred_edge_sz = f.sbh * f.sb128w << hbd;
-    if ipred_edge_sz != f.ipred_edge_sz {
-        rav1d_freep_aligned(
-            &mut *f.ipred_edge.as_mut_ptr().offset(0) as *mut *mut DynPixel as *mut c_void,
-        );
-        f.ipred_edge[0] =
-            rav1d_alloc_aligned(ipred_edge_sz as usize * 128 * 3, 64) as *mut DynPixel;
-        let ptr = f.ipred_edge[0] as *mut u8;
-        if ptr.is_null() {
-            f.ipred_edge_sz = 0;
-            return Err(ENOMEM);
-        }
-        f.ipred_edge[1] = ptr.offset(ipred_edge_sz as isize * 128 * 1) as *mut DynPixel;
-        f.ipred_edge[2] = ptr.offset(ipred_edge_sz as isize * 128 * 2) as *mut DynPixel;
-        f.ipred_edge_sz = ipred_edge_sz;
-    }
+    // TODO: Fallible allocation
+    f.ipred_edge.resize(ipred_edge_sz as usize * 128 * 3, 0);
+    f.ipred_edge_off = bpc.pxstride(ipred_edge_sz as usize * 128);
 
     let re_sz = f.sb128h * frame_hdr.tiling.cols;
     // TODO: Fallible allocation

src/internal.rs

@@ -3,7 +3,6 @@ use crate::include::common::bitdepth::BitDepth16;
 use crate::include::common::bitdepth::BitDepth8;
 use crate::include::common::bitdepth::BitDepthDependentType;
 use crate::include::common::bitdepth::BitDepthUnion;
-use crate::include::common::bitdepth::DynPixel;
 use crate::include::common::bitdepth::BPC;
 use crate::include::dav1d::common::Rav1dDataProps;
 use crate::include::dav1d::data::Rav1dData;
@@ -907,8 +906,13 @@ pub(crate) struct Rav1dFrameData {
     pub n_ts: c_int,
     pub dsp: &'static Rav1dDSPContext,
 
-    pub ipred_edge_sz: c_int,
-    pub ipred_edge: [*mut DynPixel; 3],
+    // `ipred_edge` contains 3 arrays of size `ipred_edge_off`. Use `index *
+    // ipred_edge_off` to access one of the sub-arrays. Note that `ipred_edge_off`
+    // is in pixel units (not bytes), so use `slice_as`/`mut_slice_as` and an offset
+    // in pixel units when slicing.
+    pub ipred_edge: DisjointMut<AlignedVec64<u8>>, // DynPixel
+    pub ipred_edge_off: usize,
+
     pub b4_stride: ptrdiff_t,
     pub w4: c_int,
     pub h4: c_int,
@@ -940,6 +944,7 @@ impl Rav1dFrameData {
         addr_of_mut!((*data.as_mut_ptr()).sr_cur).write(Default::default());
         addr_of_mut!((*data.as_mut_ptr()).out_cdf).write(Default::default());
         addr_of_mut!((*data.as_mut_ptr()).tiles).write(vec![]);
+        addr_of_mut!((*data.as_mut_ptr()).ipred_edge).write(Default::default());
         addr_of_mut!((*data.as_mut_ptr()).a).write(vec![]);
         addr_of_mut!((*data.as_mut_ptr()).rf).write(Default::default());
         addr_of_mut!((*data.as_mut_ptr()).lowest_pixel_mem).write(Default::default());

src/ipred_prepare.rs

@@ -1,6 +1,8 @@
 use crate::include::common::bitdepth::AsPrimitive;
 use crate::include::common::bitdepth::BitDepth;
+use crate::src::align::AlignedVec64;
 use crate::src::const_fn::const_for;
+use crate::src::disjoint_mut::DisjointMut;
 use crate::src::env::BlockContext;
 use crate::src::intra_edge::EdgeFlags;
 use crate::src::levels::IntraPredMode;
@@ -125,7 +127,9 @@ pub fn rav1d_prepare_intra_edges<BD: BitDepth>(
     edge_flags: EdgeFlags,
     dst: &[BD::Pixel], // contains 4*h first rows of picture, last row in slice contains 4*w samples
     stride: ptrdiff_t,
-    prefilter_toplevel_sb_edge: Option<&[BD::Pixel]>,
+    // Buffer and offset pair. `isize` value is the base offset that should be used
+    // when indexing into the buffer.
+    prefilter_toplevel_sb_edge: Option<(&DisjointMut<AlignedVec64<u8>>, isize)>,
     mut mode: IntraPredMode,
     angle: &mut c_int,
     tw: c_int,
@@ -169,6 +173,7 @@ pub fn rav1d_prepare_intra_edges<BD: BitDepth>(
     }
 
     // `dst_top` starts with either the top or top-left sample depending on whether have_left is true
+    let edge_buf_guard;
     let dst_top = if have_top
         && (av1_intra_prediction_edges[mode as usize]
             .needs
@@ -183,9 +188,10 @@ pub fn rav1d_prepare_intra_edges<BD: BitDepth>(
     {
         let px_have = cmp::min(8 * tw, 4 * (w - x)) as usize;
         let n = px_have + have_left as usize;
-        if let Some(prefilter_toplevel_sb_edge) = prefilter_toplevel_sb_edge {
-            let offset = (x * 4) as usize - have_left as usize;
-            &prefilter_toplevel_sb_edge[offset..][..n]
+        if let Some((edge_buf, base)) = prefilter_toplevel_sb_edge {
+            let offset = ((x * 4) as usize - have_left as usize).wrapping_add_signed(base);
+            edge_buf_guard = edge_buf.slice_as(offset..offset + n);
+            &edge_buf_guard
         } else {
             &dst[(dst_offset as isize - stride) as usize - have_left as usize..][..n]
         }

src/lib.rs

@@ -766,7 +766,7 @@ impl Drop for Rav1dContext {
                 let _ = mem::take(&mut f.frame_thread); // TODO: remove when context is owned
                 mem::take(&mut fc.task_thread.tasks); // TODO: remove when context is owned
                 rav1d_free_aligned(f.ts as *mut c_void);
-                rav1d_free_aligned(f.ipred_edge[0] as *mut c_void);
+                let _ = mem::take(&mut f.ipred_edge); // TODO: remove when context is owned
                 let _ = mem::take(&mut f.a); // TODO: remove when context is owned
                 let _ = mem::take(&mut f.out_cdf); // TODO: remove when context is owned
                 let _ = mem::take(&mut f.tiles);

src/recon.rs

@@ -2630,11 +2630,10 @@ pub(crate) unsafe fn rav1d_recon_b_intra<BD: BitDepth>(
                             ),
                         ]);
                         let top_sb_edge_slice = if t.b.y & f.sb_step - 1 == 0 {
-                            let mut top_sb_edge: *const BD::Pixel =
-                                f.ipred_edge[0] as *mut BD::Pixel;
                             let sby = t.b.y >> f.sb_shift;
-                            top_sb_edge = top_sb_edge.offset((f.sb128w * 128 * (sby - 1)) as isize);
-                            Some(slice::from_raw_parts(top_sb_edge, f.sb128w as usize * 128))
+                            let offset = f.ipred_edge_off as isize * 0
+                                + (f.sb128w * 128 * (sby - 1)) as isize;
+                            Some((&f.ipred_edge, offset))
                         } else {
                             None
                         };
@@ -2867,12 +2866,10 @@ pub(crate) unsafe fn rav1d_recon_b_intra<BD: BitDepth>(
                         if !(b.c2rust_unnamed.c2rust_unnamed.cfl_alpha[pl as usize] == 0) {
                             let mut angle = 0;
                             let top_sb_edge_slice = if t.b.y & !ss_ver & f.sb_step - 1 == 0 {
-                                let mut top_sb_edge: *const BD::Pixel =
-                                    f.ipred_edge[(pl + 1) as usize] as *mut BD::Pixel;
                                 let sby = t.b.y >> f.sb_shift;
-                                top_sb_edge =
-                                    top_sb_edge.offset((f.sb128w * 128 * (sby - 1)) as isize);
-                                Some(slice::from_raw_parts(top_sb_edge, f.sb128w as usize * 128))
+                                let offset = (f.ipred_edge_off * (pl + 1)) as isize
+                                    + (f.sb128w * 128 * (sby - 1)) as isize;
+                                Some((&f.ipred_edge, offset))
                             } else {
                                 None
                             };
@@ -3069,15 +3066,10 @@ pub(crate) unsafe fn rav1d_recon_b_intra<BD: BitDepth>(
                                     EdgeFlags::I444_LEFT_HAS_BOTTOM
                                 });
                                 let top_sb_edge_slice = if t.b.y & !ss_ver & f.sb_step - 1 == 0 {
-                                    let mut top_sb_edge: *const BD::Pixel =
-                                        f.ipred_edge[(1 + pl) as usize] as *const BD::Pixel;
                                     let sby = t.b.y >> f.sb_shift;
-                                    top_sb_edge =
-                                        top_sb_edge.offset((f.sb128w * 128 * (sby - 1)) as isize);
-                                    Some(slice::from_raw_parts(
-                                        top_sb_edge,
-                                        f.sb128w as usize * 128,
-                                    ))
+                                    let offset = (f.ipred_edge_off * (1 + pl)) as isize
+                                        + (f.sb128w * 128 * (sby - 1)) as isize;
+                                    Some((&f.ipred_edge, offset))
                                 } else {
                                     None
                                 };
@@ -3628,10 +3620,10 @@ pub(crate) unsafe fn rav1d_recon_b_inter<BD: BitDepth>(
             };
             let mut angle = 0;
             let top_sb_edge_slice = if t.b.y & f.sb_step - 1 == 0 {
-                let mut top_sb_edge = f.ipred_edge[0] as *const BD::Pixel;
                 let sby = t.b.y >> f.sb_shift;
-                top_sb_edge = top_sb_edge.offset((f.sb128w * 128 * (sby - 1)) as isize);
-                Some(slice::from_raw_parts(top_sb_edge, f.sb128w as usize * 128))
+                let offset =
+                    (f.ipred_edge_off * 0) as isize + (f.sb128w * 128 * (sby - 1)) as isize;
+                Some((&f.ipred_edge, offset))
             } else {
                 None
             };
@@ -3945,11 +3937,10 @@ pub(crate) unsafe fn rav1d_recon_b_inter<BD: BitDepth>(
                             as *mut BD::Pixel)
                             .offset(uvdstoff as isize);
                         let top_sb_edge_slice = if t.b.y & f.sb_step - 1 == 0 {
-                            let mut top_sb_edge =
-                                f.ipred_edge[(pl + 1) as usize] as *const BD::Pixel;
                             let sby = t.b.y >> f.sb_shift;
-                            top_sb_edge = top_sb_edge.offset((f.sb128w * 128 * (sby - 1)) as isize);
-                            Some(slice::from_raw_parts(top_sb_edge, f.sb128w as usize * 128))
+                            let offset = (f.ipred_edge_off * (pl + 1)) as isize
+                                + (f.sb128w * 128 * (sby - 1)) as isize;
+                            Some((&f.ipred_edge, offset))
                         } else {
                             None
                         };
@@ -4529,15 +4520,14 @@ pub(crate) unsafe fn rav1d_backup_ipred_edge<BD: BitDepth>(
     let y: *const BD::Pixel = (f.cur.data.as_ref().unwrap().data[0] as *const BD::Pixel)
         .offset((x_off * 4) as isize)
         .offset((((t.b.y + f.sb_step) * 4 - 1) as isize * BD::pxstride(f.cur.stride[0])) as isize);
+    let ipred_edge_off = (f.ipred_edge_off * 0) + (sby_off + x_off * 4) as usize;
+    let n = 4 * (ts.tiling.col_end - x_off) as usize;
     BD::pixel_copy(
-        &mut slice::from_raw_parts_mut(
-            f.ipred_edge[0].cast(),
-            (sby_off + x_off * 4 + (4 * (ts.tiling.col_end - x_off)))
-                .try_into()
-                .unwrap(),
-        )[(sby_off + x_off * 4).try_into().unwrap()..],
-        slice::from_raw_parts(y, (4 * (ts.tiling.col_end - x_off)).try_into().unwrap()),
-        (4 * (ts.tiling.col_end - x_off)).try_into().unwrap(),
+        &mut f
+            .ipred_edge
+            .mut_slice_as(ipred_edge_off..ipred_edge_off + n),
+        slice::from_raw_parts(y, n),
+        n,
     );
     if f.cur.p.layout as c_uint != Rav1dPixelLayout::I400 as c_int as c_uint {
         let ss_ver =
@@ -4548,22 +4538,20 @@ pub(crate) unsafe fn rav1d_backup_ipred_edge<BD: BitDepth>(
             + (((t.b.y + f.sb_step) * 4 >> ss_ver) - 1) as isize * BD::pxstride(f.cur.stride[1]);
         let mut pl = 1;
         while pl <= 2 {
+            let ipred_edge_off =
+                (f.ipred_edge_off * pl) + (sby_off + (x_off * 4 >> ss_hor)) as usize;
+            let n = 4 * (ts.tiling.col_end - x_off) as usize >> ss_hor;
             BD::pixel_copy(
-                &mut slice::from_raw_parts_mut(
-                    f.ipred_edge[pl as usize].cast(),
-                    (sby_off + (x_off * 4 >> ss_hor) + (4 * (ts.tiling.col_end - x_off) >> ss_hor))
-                        .try_into()
-                        .unwrap(),
-                )[(sby_off + (x_off * 4 >> ss_hor)).try_into().unwrap()..],
+                &mut f
+                    .ipred_edge
+                    .mut_slice_as(ipred_edge_off..ipred_edge_off + n),
                 &slice::from_raw_parts(
-                    f.cur.data.as_ref().unwrap().data[pl as usize].cast(),
-                    (uv_off + (4 * (ts.tiling.col_end - x_off) >> ss_hor) as isize)
-                        .try_into()
-                        .unwrap(),
-                )[uv_off.try_into().unwrap()..],
-                (4 * (ts.tiling.col_end - x_off) >> ss_hor)
-                    .try_into()
-                    .unwrap(),
+                    f.cur.data.as_ref().unwrap().data[pl]
+                        .cast::<BD::Pixel>()
+                        .offset(uv_off),
+                    n,
+                ),
+                n,
             );
             pl += 1;
         }