struct Rav1dContext_refs::{,cur_,prev_}segmap: Arcify with an `Op…

…tion<DisjointMutArcSlice>`.

src/decode.rs

@@ -31,6 +31,7 @@ use crate::src::cdf::CdfMvContext;
 use crate::src::ctx::CaseSet;
 use crate::src::dequant_tables::dav1d_dq_tbl;
 use crate::src::disjoint_mut::DisjointMut;
+use crate::src::disjoint_mut::DisjointMutSlice;
 use crate::src::enum_map::enum_map;
 use crate::src::enum_map::enum_map_ty;
 use crate::src::enum_map::DefaultValue;
@@ -851,16 +852,12 @@ unsafe fn read_vartx_tree(
 }
 
 #[inline]
-unsafe fn get_prev_frame_segid(
+fn get_prev_frame_segid(
     frame_hdr: &Rav1dFrameHeader,
     b: Bxy,
     w4: c_int,
     h4: c_int,
-    // It's very difficult to make this safe (a slice),
-    // as it comes from [`Dav1dFrameContext::prev_segmap`],
-    // which is set to [`Dav1dFrameContext::prev_segmap_ref`],
-    // which is a [`Dav1dRef`], which has no size and is refcounted.
-    ref_seg_map: *const u8,
+    ref_seg_map: &DisjointMutSlice<u8>,
     stride: ptrdiff_t,
 ) -> u8 {
     assert!(frame_hdr.primary_ref_frame != RAV1D_PRIMARY_REF_NONE);
@@ -872,10 +869,9 @@ unsafe fn get_prev_frame_segid(
     let stride = usize::try_from(stride).unwrap();
 
     let mut prev_seg_id = 8;
-    let ref_seg_map = std::slice::from_raw_parts(
-        ref_seg_map.offset(b.y as isize * stride as isize + b.x as isize),
-        h4 * stride,
-    );
+    let offset = b.y as usize * stride as usize + b.x as usize;
+    let len = h4 as usize * stride;
+    let ref_seg_map = unsafe { ref_seg_map.index(offset..offset + len) };
 
     assert!(w4 <= stride);
     for ref_seg_map in ref_seg_map.chunks_exact(stride) {
@@ -1343,9 +1339,9 @@ unsafe fn decode_b_inner(
     let frame_hdr: &Rav1dFrameHeader = &f.frame_hdr.as_ref().unwrap();
     if frame_hdr.segmentation.enabled != 0 {
         if frame_hdr.segmentation.update_map == 0 {
-            if !(f.prev_segmap).is_null() {
+            if let Some(prev_segmap) = f.prev_segmap.as_ref() {
                 let seg_id =
-                    get_prev_frame_segid(frame_hdr, t.b, w4, h4, f.prev_segmap, f.b4_stride);
+                    get_prev_frame_segid(frame_hdr, t.b, w4, h4, &prev_segmap.inner, f.b4_stride);
                 if seg_id >= RAV1D_MAX_SEGMENTS.into() {
                     return Err(());
                 }
@@ -1364,9 +1360,15 @@ unsafe fn decode_b_inner(
                 seg_pred
             } {
                 // temporal predicted seg_id
-                if !(f.prev_segmap).is_null() {
-                    let seg_id =
-                        get_prev_frame_segid(frame_hdr, t.b, w4, h4, f.prev_segmap, f.b4_stride);
+                if let Some(prev_segmap) = f.prev_segmap.as_ref() {
+                    let seg_id = get_prev_frame_segid(
+                        frame_hdr,
+                        t.b,
+                        w4,
+                        h4,
+                        &prev_segmap.inner,
+                        f.b4_stride,
+                    );
                     if seg_id >= RAV1D_MAX_SEGMENTS.into() {
                         return Err(());
                     }
@@ -1379,7 +1381,7 @@ unsafe fn decode_b_inner(
                     t.b,
                     have_top,
                     have_left,
-                    f.cur_segmap,
+                    &f.cur_segmap.as_ref().unwrap().inner,
                     f.b4_stride as usize,
                 );
                 let diff = rav1d_msac_decode_symbol_adapt8(
@@ -1453,9 +1455,9 @@ unsafe fn decode_b_inner(
             seg_pred
         } {
             // temporal predicted seg_id
-            if !(f.prev_segmap).is_null() {
+            if let Some(prev_segmap) = f.prev_segmap.as_ref() {
                 let seg_id =
-                    get_prev_frame_segid(frame_hdr, t.b, w4, h4, f.prev_segmap, f.b4_stride);
+                    get_prev_frame_segid(frame_hdr, t.b, w4, h4, &prev_segmap.inner, f.b4_stride);
                 if seg_id >= RAV1D_MAX_SEGMENTS.into() {
                     return Err(());
                 }
@@ -1464,8 +1466,13 @@ unsafe fn decode_b_inner(
                 b.seg_id = 0;
             }
         } else {
-            let (pred_seg_id, seg_ctx) =
-                get_cur_frame_segid(t.b, have_top, have_left, f.cur_segmap, f.b4_stride as usize);
+            let (pred_seg_id, seg_ctx) = get_cur_frame_segid(
+                t.b,
+                have_top,
+                have_left,
+                &f.cur_segmap.as_ref().unwrap().inner,
+                f.b4_stride as usize,
+            );
             if b.skip != 0 {
                 b.seg_id = pred_seg_id as u8;
             } else {
@@ -2993,18 +3000,12 @@ unsafe fn decode_b_inner(
         // Need checked casts here because we're using `from_raw_parts_mut` and an overflow would be UB.
         let [by, bx, bh4, bw4] = [t.b.y, t.b.x, bh4, bw4].map(|it| usize::try_from(it).unwrap());
         let b4_stride = usize::try_from(f.b4_stride).unwrap();
-        let cur_segmap_len = (by * b4_stride + bx)
-            + if bh4 == 0 {
-                0
-            } else {
-                (b4_stride * (bh4 - 1)) + bw4
-            };
-        let cur_segmap = std::slice::from_raw_parts_mut(f.cur_segmap, cur_segmap_len);
-        let seg_ptr = &mut cur_segmap[by * b4_stride + bx..];
-
+        let cur_segmap = &f.cur_segmap.as_ref().unwrap().inner;
+        let offset = by * b4_stride + bx;
         CaseSet::<32, false>::one((), bw4, 0, |case, ()| {
-            for seg_ptr in seg_ptr.chunks_mut(b4_stride).take(bh4) {
-                case.set(seg_ptr, b.seg_id);
+            for i in 0..bh4 {
+                let i = offset + i * b4_stride;
+                case.set(&mut cur_segmap.index_mut(i..i + bw4), b.seg_id);
             }
         });
     }
@@ -4586,8 +4587,8 @@ pub(crate) unsafe fn rav1d_decode_frame_exit(
         }
     }
 
-    rav1d_ref_dec(&mut f.cur_segmap_ref);
-    rav1d_ref_dec(&mut f.prev_segmap_ref);
+    let _ = mem::take(&mut f.cur_segmap);
+    let _ = mem::take(&mut f.prev_segmap);
     rav1d_ref_dec(&mut f.mvs_ref);
     let _ = mem::take(&mut f.seq_hdr);
     let _ = mem::take(&mut f.frame_hdr);
@@ -4947,8 +4948,7 @@ pub unsafe fn rav1d_submit_frame(c: &mut Rav1dContext) -> Rav1dResult {
     // segmap
     if frame_hdr.segmentation.enabled != 0 {
         // By default, the previous segmentation map is not initialised.
-        f.prev_segmap_ref = ptr::null_mut();
-        f.prev_segmap = ptr::null();
+        f.prev_segmap = None;
 
         // We might need a previous frame's segmentation map.
         // This happens if there is either no update or a temporal update.
@@ -4958,50 +4958,36 @@ pub unsafe fn rav1d_submit_frame(c: &mut Rav1dContext) -> Rav1dResult {
             let ref_w = (ref_coded_width[pri_ref] + 7 >> 3) << 1;
             let ref_h = (f.refp[pri_ref].p.p.h + 7 >> 3) << 1;
             if ref_w == f.bw && ref_h == f.bh {
-                f.prev_segmap_ref = c.refs[frame_hdr.refidx[pri_ref] as usize].segmap;
-                if !f.prev_segmap_ref.is_null() {
-                    rav1d_ref_inc(f.prev_segmap_ref);
-                    f.prev_segmap = (*f.prev_segmap_ref).data.cast::<u8>();
-                }
+                f.prev_segmap = c.refs[frame_hdr.refidx[pri_ref] as usize].segmap.clone();
             }
         }
 
-        if frame_hdr.segmentation.update_map != 0 {
-            // We're updating an existing map,
-            // but need somewhere to put the new values.
-            // Allocate them here (the data actually gets set elsewhere).
-            f.cur_segmap_ref = rav1d_ref_create_using_pool(
-                c.segmap_pool,
-                ::core::mem::size_of::<u8>() * f.b4_stride as usize * 32 * f.sb128h as usize,
-            );
-            if f.cur_segmap_ref.is_null() {
-                rav1d_ref_dec(&mut f.prev_segmap_ref);
-                on_error(f, c, out);
-                return Err(ENOMEM);
-            }
-            f.cur_segmap = (*f.cur_segmap_ref).data.cast::<u8>();
-        } else if !f.prev_segmap_ref.is_null() {
-            // We're not updating an existing map,
-            // and we have a valid reference. Use that.
-            f.cur_segmap_ref = f.prev_segmap_ref;
-            rav1d_ref_inc(f.cur_segmap_ref);
-            f.cur_segmap = (*f.prev_segmap_ref).data.cast::<u8>();
-        } else {
-            // We need to make a new map. Allocate one here and zero it out.
-            let segmap_size =
-                ::core::mem::size_of::<u8>() * f.b4_stride as usize * 32 * f.sb128h as usize;
-            f.cur_segmap_ref = rav1d_ref_create_using_pool(c.segmap_pool, segmap_size);
-            if f.cur_segmap_ref.is_null() {
-                on_error(f, c, out);
-                return Err(ENOMEM);
-            }
-            f.cur_segmap = (*f.cur_segmap_ref).data.cast::<u8>();
-            slice::from_raw_parts_mut(f.cur_segmap, segmap_size).fill(0);
-        }
+        f.cur_segmap = Some(
+            match (
+                frame_hdr.segmentation.update_map != 0,
+                f.prev_segmap.as_mut(),
+            ) {
+                (true, _) | (false, None) => {
+                    // If we're updating an existing map,
+                    // we need somewhere to put the new values.
+                    // Allocate them here (the data actually gets set elsewhere).
+                    // Since this is Rust, we have to initialize it anyways.
+
+                    // Otherwise if there's no previous, we need to make a new map.
+                    // Allocate one here and zero it out.
+                    let segmap_size = f.b4_stride as usize * 32 * f.sb128h as usize;
+                    (0..segmap_size).map(|_| 0u8).collect()
+                }
+                (_, Some(prev_segmap)) => {
+                    // We're not updating an existing map,
+                    // and we have a valid reference. Use that.
+                    prev_segmap.clone()
+                }
+            },
+        );
     } else {
-        f.cur_segmap = ptr::null_mut();
-        f.cur_segmap_ref = ptr::null_mut();
-        f.prev_segmap_ref = ptr::null_mut();
+        f.cur_segmap = None;
+        f.prev_segmap = None;
     }
 
     // update references etc.
@@ -5019,11 +5005,7 @@ pub unsafe fn rav1d_submit_frame(c: &mut Rav1dContext) -> Rav1dResult {
                 c.cdf[i] = f.in_cdf.clone();
             }
 
-            rav1d_ref_dec(&mut c.refs[i].segmap);
-            c.refs[i].segmap = f.cur_segmap_ref;
-            if !f.cur_segmap_ref.is_null() {
-                rav1d_ref_inc(f.cur_segmap_ref);
-            }
+            c.refs[i].segmap = f.cur_segmap.clone();
             rav1d_ref_dec(&mut c.refs[i].refmvs);
             if !frame_hdr.allow_intrabc {
                 c.refs[i].refmvs = f.mvs_ref;
@@ -5045,7 +5027,7 @@ pub unsafe fn rav1d_submit_frame(c: &mut Rav1dContext) -> Rav1dResult {
                         rav1d_thread_picture_unref(&mut c.refs[i].p);
                     }
                     let _ = mem::take(&mut c.cdf[i]);
-                    rav1d_ref_dec(&mut c.refs[i].segmap);
+                    let _ = mem::take(&mut c.refs[i].segmap);
                     rav1d_ref_dec(&mut c.refs[i].refmvs);
                 }
             }

src/env.rs

@@ -4,6 +4,7 @@ use crate::include::dav1d::headers::Rav1dFrameHeader;
 use crate::include::dav1d::headers::Rav1dWarpedMotionParams;
 use crate::include::dav1d::headers::Rav1dWarpedMotionType;
 use crate::src::align::Align8;
+use crate::src::disjoint_mut::DisjointMutSlice;
 use crate::src::internal::Bxy;
 use crate::src::levels::mv;
 use crate::src::levels::BlockLevel;
@@ -25,7 +26,6 @@ use std::cmp;
 use std::cmp::Ordering;
 use std::ffi::c_int;
 use std::ffi::c_uint;
-use std::slice;
 
 #[derive(Default)]
 #[repr(C)]
@@ -593,31 +593,20 @@ pub fn get_drl_context(ref_mv_stack: &[refmvs_candidate; 8], ref_idx: usize) ->
 }
 
 #[inline]
-pub unsafe fn get_cur_frame_segid(
+pub fn get_cur_frame_segid(
     b: Bxy,
     have_top: bool,
     have_left: bool,
-    // It's very difficult to make this safe (a slice),
-    // as it is negatively indexed
-    // and it comes from [`Dav1dFrameContext::cur_segmap`],
-    // which is set to [`Dav1dFrameContext::cur_segmap_ref`] and [`Dav1dFrameContext::prev_segmap_ref`],
-    // which are [`Dav1dRef`]s, which have no size and are refcounted.
-    cur_seg_map: *const u8,
+    cur_seg_map: &DisjointMutSlice<u8>,
     stride: usize,
 ) -> (u8, u8) {
     let negative_adjustment = have_left as usize + have_top as usize * stride;
     let offset = b.x as usize + b.y as usize * stride - negative_adjustment;
-    let len = match (have_left, have_top) {
-        (true, true) => stride + 1,
-        (true, false) | (false, true) => 1,
-        (false, false) => 0,
-    };
-    let cur_seg_map = &slice::from_raw_parts(cur_seg_map, offset + len)[offset..];
     match (have_left, have_top) {
         (true, true) => {
-            let l = cur_seg_map[stride];
-            let a = cur_seg_map[1];
-            let al = cur_seg_map[0];
+            let l = *unsafe { cur_seg_map.index(offset + stride) };
+            let a = *unsafe { cur_seg_map.index(offset + 1) };
+            let al = *unsafe { cur_seg_map.index(offset) };
             let seg_ctx = if l == a && al == l {
                 2
             } else if l == a || al == l || a == al {
@@ -628,7 +617,7 @@ pub unsafe fn get_cur_frame_segid(
             let seg_id = if a == al { a } else { l };
             (seg_id, seg_ctx)
         }
-        (true, false) | (false, true) => (cur_seg_map[0], 0),
+        (true, false) | (false, true) => (*unsafe { cur_seg_map.index(offset) }, 0),
         (false, false) => (0, 0),
     }
 }

src/internal.rs

@@ -28,6 +28,7 @@ use crate::src::cdef::Rav1dCdefDSPContext;
 use crate::src::cdf::CdfContext;
 use crate::src::cdf::CdfThreadContext;
 use crate::src::disjoint_mut::DisjointMut;
+use crate::src::disjoint_mut::DisjointMutArcSlice;
 use crate::src::env::BlockContext;
 use crate::src::error::Rav1dResult;
 use crate::src::filmgrain::Rav1dFilmGrainDSPContext;
@@ -221,7 +222,7 @@ pub(crate) struct TaskThreadData {
 #[repr(C)]
 pub(crate) struct Rav1dContext_refs {
     pub p: Rav1dThreadPicture,
-    pub segmap: *mut Rav1dRef,
+    pub segmap: Option<DisjointMutArcSlice<u8>>,
     pub refmvs: *mut Rav1dRef,
     pub refpoc: [c_uint; 7],
 }
@@ -716,10 +717,8 @@ pub(crate) struct Rav1dFrameData {
     pub mvs: *mut refmvs_temporal_block,
     pub ref_mvs: [*mut refmvs_temporal_block; 7],
     pub ref_mvs_ref: [*mut Rav1dRef; 7],
-    pub cur_segmap_ref: *mut Rav1dRef,
-    pub prev_segmap_ref: *mut Rav1dRef,
-    pub cur_segmap: *mut u8,
-    pub prev_segmap: *const u8,
+    pub cur_segmap: Option<DisjointMutArcSlice<u8>>,
+    pub prev_segmap: Option<DisjointMutArcSlice<u8>>,
     pub refpoc: [c_uint; 7],
     pub refrefpoc: [[c_uint; 7]; 7],
     pub gmv_warp_allowed: [u8; 7],

src/lib.rs

@@ -681,7 +681,7 @@ pub(crate) unsafe fn rav1d_flush(c: *mut Rav1dContext) {
         if (*c).refs[i as usize].p.p.frame_hdr.is_some() {
             rav1d_thread_picture_unref(&mut (*((*c).refs).as_mut_ptr().offset(i as isize)).p);
         }
-        rav1d_ref_dec(&mut (*((*c).refs).as_mut_ptr().offset(i as isize)).segmap);
+        let _ = mem::take(&mut (*c).refs[i as usize].segmap);
         rav1d_ref_dec(&mut (*((*c).refs).as_mut_ptr().offset(i as isize)).refmvs);
         let _ = mem::take(&mut (*c).cdf[i]);
         i += 1;
@@ -854,7 +854,7 @@ impl Drop for Rav1dContext {
                     );
                 }
                 rav1d_ref_dec(&mut (*(self.refs).as_mut_ptr().offset(n_4 as isize)).refmvs);
-                rav1d_ref_dec(&mut (*(self.refs).as_mut_ptr().offset(n_4 as isize)).segmap);
+                let _ = mem::take(&mut self.refs[n_4 as usize].segmap);
                 n_4 += 1;
             }
             let _ = mem::take(&mut self.seq_hdr);

src/obu.rs

@@ -68,7 +68,6 @@ use crate::src::picture::rav1d_thread_picture_ref;
 use crate::src::picture::rav1d_thread_picture_unref;
 use crate::src::picture::PictureFlags;
 use crate::src::r#ref::rav1d_ref_dec;
-use crate::src::r#ref::rav1d_ref_inc;
 use crate::src::thread_task::FRAME_ERROR;
 use std::array;
 use std::cmp;
@@ -2296,7 +2295,7 @@ unsafe fn parse_obus(
                         if c.refs[i as usize].p.p.frame_hdr.is_some() {
                             rav1d_thread_picture_unref(&mut c.refs[i as usize].p);
                         }
-                        rav1d_ref_dec(&mut c.refs[i as usize].segmap);
+                        let _ = mem::take(&mut c.refs[i as usize].segmap);
                         rav1d_ref_dec(&mut c.refs[i as usize].refmvs);
                         let _ = mem::take(&mut c.cdf[i]);
                     }
@@ -2623,11 +2622,7 @@ unsafe fn parse_obus(
 
                     c.cdf[i as usize] = c.cdf[r as usize].clone();
 
-                    rav1d_ref_dec(&mut c.refs[i as usize].segmap);
-                    c.refs[i as usize].segmap = c.refs[r as usize].segmap;
-                    if !c.refs[r as usize].segmap.is_null() {
-                        rav1d_ref_inc(c.refs[r as usize].segmap);
-                    }
+                    c.refs[i as usize].segmap = c.refs[r as usize].segmap.clone();
                     rav1d_ref_dec(&mut c.refs[i as usize].refmvs);
                 }
             }