From b42479b779eb8bd48cdc691d01e685d706be9ba7 Mon Sep 17 00:00:00 2001
From: Khyber Sen <kkysen@gmail.com>
Date: Mon, 23 Sep 2024 02:21:01 -0700
Subject: [PATCH] `fn DisjointMutArcSlice::new_zeroed_slice`: Add constructor
 that uses zeroed allocations to optimize out initialization, and use this on
 `f.mvs`.

---
 lib.rs              |  1 +
 src/decode.rs       |  9 ++++-----
 src/disjoint_mut.rs | 39 ++++++++++++++++++++++++++++++++++++---
 src/refmvs.rs       |  2 +-
 4 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/lib.rs b/lib.rs
index 3bc3ec59c..c90270bf4 100644
--- a/lib.rs
+++ b/lib.rs
@@ -4,6 +4,7 @@
     any(target_arch = "riscv32", target_arch = "riscv64"),
     feature(stdarch_riscv_feature_detection)
 )]
+#![feature(new_uninit)]
 #![deny(unsafe_op_in_unsafe_fn)]
 #![allow(clippy::all)]
 #![deny(clippy::undocumented_unsafe_blocks)]
diff --git a/src/decode.rs b/src/decode.rs
index 6695e9798..2f29ab6fc 100644
--- a/src/decode.rs
+++ b/src/decode.rs
@@ -29,6 +29,7 @@ use crate::src::cdf::CdfThreadContext;
 use crate::src::ctx::CaseSet;
 use crate::src::dequant_tables::dav1d_dq_tbl;
 use crate::src::disjoint_mut::DisjointMut;
+use crate::src::disjoint_mut::DisjointMutArcSlice;
 use crate::src::disjoint_mut::DisjointMutSlice;
 use crate::src::enum_map::enum_map;
 use crate::src::enum_map::enum_map_ty;
@@ -5220,11 +5221,9 @@ pub fn rav1d_submit_frame(c: &Rav1dContext, state: &mut Rav1dState) -> Rav1dResu
     // ref_mvs
     if frame_hdr.frame_type.is_inter_or_switch() || frame_hdr.allow_intrabc {
         // TODO fallible allocation
-        f.mvs = Some(
-            (0..f.sb128h as usize * 16 * (f.b4_stride >> 1) as usize)
-                .map(|_| Default::default())
-                .collect(),
-        );
+        f.mvs = Some(DisjointMutArcSlice::new_zeroed_slice(
+            f.sb128h as usize * 16 * (f.b4_stride >> 1) as usize,
+        ));
         if !frame_hdr.allow_intrabc {
             for i in 0..7 {
                 f.refpoc[i] = f.refp[i].p.frame_hdr.as_ref().unwrap().frame_offset as c_uint;
diff --git a/src/disjoint_mut.rs b/src/disjoint_mut.rs
index ea3e65f50..17b72a9fa 100644
--- a/src/disjoint_mut.rs
+++ b/src/disjoint_mut.rs
@@ -13,6 +13,7 @@ use std::fmt::Formatter;
 use std::marker::PhantomData;
 use std::mem;
 use std::mem::ManuallyDrop;
+use std::mem::MaybeUninit;
 use std::ops::Deref;
 use std::ops::DerefMut;
 use std::ops::Index;
@@ -27,6 +28,7 @@ use std::ptr::addr_of_mut;
 use std::sync::Arc;
 use zerocopy::AsBytes;
 use zerocopy::FromBytes;
+use zerocopy::FromZeroes;
 
 /// Wraps an indexable collection to allow unchecked concurrent mutable borrows.
 ///
@@ -1185,14 +1187,14 @@ impl<T> FromIterator<T> for DisjointMutArcSlice<T> {
     fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
         #[cfg(debug_assertions)]
         let inner = {
-            let box_slice = iter.into_iter().collect::<Box<[_]>>();
+            let box_slice = iter.into_iter().collect::<Box<[T]>>();
             Arc::new(DisjointMut::new(box_slice))
         };
         #[cfg(not(debug_assertions))]
         let inner = {
             use std::mem;
 
-            let arc_slice = iter.into_iter().collect::<Arc<[_]>>();
+            let arc_slice = iter.into_iter().collect::<Arc<[T]>>();
 
             // Do our best to check that `DisjointMut` is in fact `#[repr(transparent)]`.
             type A = Vec<u8>; // Some concrete sized type.
@@ -1201,7 +1203,7 @@ impl<T> FromIterator<T> for DisjointMutArcSlice<T> {
 
             // SAFETY: When `#[cfg(not(debug_assertions))]`, `DisjointMut` is `#[repr(transparent)]`,
             // containing only an `UnsafeCell`, which is also `#[repr(transparent)]`.
-            unsafe { mem::transmute::<Arc<[_]>, Arc<DisjointMut<[_]>>>(arc_slice) }
+            unsafe { mem::transmute::<Arc<[T]>, Arc<DisjointMut<[T]>>>(arc_slice) }
         };
         Self { inner }
     }
@@ -1212,3 +1214,34 @@ impl<T> Default for DisjointMutArcSlice<T> {
         [].into_iter().collect()
     }
 }
+
+impl<T: FromZeroes> DisjointMutArcSlice<T> {
+    pub fn new_zeroed_slice(len: usize) -> Self {
+        #[cfg(debug_assertions)]
+        let inner = {
+            let box_slice = Box::new_zeroed_slice(len);
+            Arc::new(DisjointMut::new(box_slice))
+        };
+        #[cfg(not(debug_assertions))]
+        let inner = {
+            use std::mem;
+
+            let arc_slice = Arc::new_zeroed_slice(len);
+
+            // Do our best to check that `DisjointMut` is in fact `#[repr(transparent)]`.
+            type A = Vec<u8>; // Some concrete sized type.
+            const _: () = assert!(mem::size_of::<DisjointMut<A>>() == mem::size_of::<A>());
+            const _: () = assert!(mem::align_of::<DisjointMut<A>>() == mem::align_of::<A>());
+
+            // SAFETY: When `#[cfg(not(debug_assertions))]`, `DisjointMut` is `#[repr(transparent)]`,
+            // containing only an `UnsafeCell`, which is also `#[repr(transparent)]`.
+            unsafe { mem::transmute::<Arc<[T]>, Arc<DisjointMut<[T]>>>(arc_slice) }
+        };
+        // SAFETY: `T: FromZeroes`, and the `MaybeUninit<T>` is all zeros,
+        // since it is allocated with `new_zeroed_slice`.
+        let inner = unsafe {
+            mem::transmute::<Arc<DisjointMutSlice<MaybeUninit<T>>>, Arc<DisjointMutSlice<T>>>(inner)
+        };
+        Self { inner }
+    }
+}
diff --git a/src/refmvs.rs b/src/refmvs.rs
index fc2c2bc36..4ffec2c84 100644
--- a/src/refmvs.rs
+++ b/src/refmvs.rs
@@ -30,7 +30,7 @@ use std::ptr;
 use std::slice;
 use zerocopy::FromZeroes;
 
-#[derive(Clone, Copy, Default, PartialEq, Eq)]
+#[derive(Clone, Copy, Default, PartialEq, Eq, FromZeroes)]
 #[repr(C, packed)]
 pub struct RefMvsTemporalBlock {
     pub mv: Mv,