refactor: clean up double buffer to fix a Miri warning

Author: JakubValtar

src/double_buffer.rs

@@ -0,0 +1,127 @@
+use core::{mem::MaybeUninit, slice};
+
+use alloc::{boxed::Box, vec::Vec};
+
+/// Double buffer. Wraps a mutable slice and allocates a scratch memory of the same size, so that
+/// elements can be freely scattered from buffer to buffer.
+///
+/// # Drop behavior
+///
+/// Drop ensures that the mutable slice this buffer was constructed with contains all the original
+/// elements.
+pub struct DoubleBuffer<'a, T> {
+    slice: &'a mut [MaybeUninit<T>],
+    scratch: Box<[MaybeUninit<T>]>,
+    slice_is_write: bool,
+}
+
+impl<'a, T> DoubleBuffer<'a, T> {
+    /// Creates a double buffer, allocating a scratch buffer of the same length as the input slice.
+    ///
+    /// The supplied slice becomes the read buffer, the scratch buffer becomes the write buffer.
+    pub fn new(slice: &'a mut [T]) -> Self {
+        // SAFETY: The Drop impl ensures that the slice is initialized.
+        let slice = unsafe { slice_as_uninit_mut(slice) };
+        let scratch = {
+            let mut v = Vec::with_capacity(slice.len());
+            // SAFETY: we just allocated this capacity and MaybeUninit can be garbage.
+            unsafe {
+                v.set_len(slice.len());
+            }
+            v.into_boxed_slice()
+        };
+        DoubleBuffer {
+            slice,
+            scratch,
+            slice_is_write: false,
+        }
+    }
+
+    /// Scatters the elements from the read buffer to the computed indices in
+    /// the write buffer. The read buffer is iterated from the beginning.
+    ///
+    /// Call `swap` after this function to commit the write buffer state.
+    pub fn scatter<F>(&mut self, mut indexer: F)
+    where
+        F: FnMut(&T) -> usize,
+    {
+        let (read, write) = self.as_read_write();
+
+        let len = write.len();
+
+        for t in read {
+            let index = indexer(t);
+            if index >= len {
+                return;
+            }
+            let write_ptr = write[index].as_mut_ptr();
+            unsafe {
+                // SAFETY: both pointers are valid for T, aligned, and nonoverlapping
+                write_ptr.copy_from_nonoverlapping(t as *const T, 1);
+            }
+        }
+    }
+
+    /// Returns the current read and write buffers.
+    fn as_read_write(&mut self) -> (&[T], &mut [MaybeUninit<T>]) {
+        let (read, write): (&[MaybeUninit<T>], &mut [MaybeUninit<T>]) = if self.slice_is_write {
+            (self.scratch.as_ref(), self.slice)
+        } else {
+            (self.slice, self.scratch.as_mut())
+        };
+
+        // SAFETY: The read buffer is always initialized.
+        let read = unsafe { slice_assume_init_ref(read) };
+
+        (read, write)
+    }
+
+    /// Swaps the read and write buffer, committing the write buffer state.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that every element of the write buffer was
+    /// written to before calling this function.
+    pub unsafe fn swap(&mut self) {
+        self.slice_is_write = !self.slice_is_write;
+    }
+}
+
+/// Ensures that the input slice contains all the original elements.
+impl<'a, T> Drop for DoubleBuffer<'a, T> {
+    fn drop(&mut self) {
+        if self.slice_is_write {
+            // The input slice is the write buffer, copy the consistent state from the read buffer
+            unsafe {
+                // SAFETY: `scratch` is the read buffer, it is initialized. The length is the same.
+                self.slice
+                    .as_mut_ptr()
+                    .copy_from_nonoverlapping(self.scratch.as_ptr(), self.slice.len());
+            }
+            self.slice_is_write = false;
+        }
+    }
+}
+
+/// Get a slice of the initialized items.
+///
+/// # Safety
+///
+/// The caller must ensure that all the items are initialized.
+#[inline(always)]
+pub unsafe fn slice_assume_init_ref<T>(slice: &[MaybeUninit<T>]) -> &[T] {
+    // SAFETY: `[MaybeUninit<T>]` and `[T]` have the same layout.
+    unsafe { slice::from_raw_parts(slice.as_ptr() as *const T, slice.len()) }
+}
+
+/// View the mutable slice of `T` as a slice of `MaybeUnint<T>`.
+///
+/// # Safety
+///
+/// The caller must ensure that all the items of the returned slice are
+/// initialized before dropping it.
+#[inline(always)]
+pub unsafe fn slice_as_uninit_mut<T>(slice: &mut [T]) -> &mut [MaybeUninit<T>] {
+    // SAFETY: `[MaybeUninit<T>]` and `[T]` have the same layout.
+    unsafe { slice::from_raw_parts_mut(slice.as_mut_ptr() as *mut MaybeUninit<T>, slice.len()) }
+}

src/lib.rs

@@ -94,6 +94,7 @@ extern crate alloc;
 
 use alloc::vec::Vec;
 
+mod double_buffer;
 mod scalar;
 mod sort;
 

src/sort.rs

@@ -1,9 +1,8 @@
 //! Implementations of radix keys and sorting functions.
 
-use alloc::vec::Vec;
 use core::mem;
 
-use crate::Key;
+use crate::{double_buffer::DoubleBuffer, Key};
 
 /// Unsigned integers used as sorting keys for radix sort.
 ///
@@ -139,41 +138,32 @@ macro_rules! sort_impl {
             // digit, our elements are sorted.
             for digit in 0..DIGIT_COUNT {
                 if !(digit_skip_enabled && skip_digit[digit]) {
-                    // Copy the offsets. We need the original later for a consistency check.
-                    // As we write elements into each bucket, we increment the bucket offset
-                    // so that it points to the next empty slot.
-                    let mut working_offsets: [$offset_type; BUCKET_COUNT] = offsets[digit];
-
-                    for r_pos in 0..len {
-                        let t: &T = unsafe {
-                            // This is safe, r_pos is in (0..len)
-                            buffer.read(r_pos)
-                        };
+                    // Initial offset of each bucket.
+                    let init_offsets = &offsets[digit];
+                    // Offset of the first empty index in each bucket.
+                    let mut working_offsets = *init_offsets;
 
+                    buffer.scatter(|t| {
                         let key = key_fn(t);
                         let bucket = extract_digit(key, digit);
 
                         let offset = &mut working_offsets[bucket];
 
-                        unsafe {
-                            // Make sure the offset is in bounds. An unreliable key function, which
-                            // returns different keys for the same element when called repeatedly,
-                            // can cause offsets to go out of bounds.
-                            let w_pos = usize::min(*offset as usize, len - 1);
-
-                            // This is safe, w_pos is in (0..len)
-                            buffer.write(w_pos, t);
-                        }
+                        // Make sure the offset is in bounds. An unreliable key function, which
+                        // returns different keys for the same element when called repeatedly,
+                        // can cause offsets to go out of bounds.
+                        let clamped_offset = usize::min(*offset as usize, len - 1);
 
                         // Increment the offset of the bucket. Use wrapping add in case the
                         // key function is unreliable and the bucket overflowed.
                         *offset = offset.wrapping_add(1);
-                    }
+
+                        clamped_offset
+                    });
 
                     // Check that each bucket had the same number of insertions as we expected.
-                    // If this is not true, then the key function is unreliable and the write buffer
-                    // is not consistent: some elements were overwritten, some were not written to
-                    // and contain garbage.
+                    // If this is not true, then the key function is unreliable and some elements
+                    // in the write buffer were not written to.
                     //
                     // If the key function is unreliable, but the sizes of buckets ended up being
                     // the same, it would not get detected. This is sound, the only consequence is
@@ -191,9 +181,7 @@ macro_rules! sort_impl {
                             // The bucket sizes do not match expected sizes, the key function is
                             // unreliable (programming mistake).
                             //
-                            // Drop impl of the double buffer will make sure that the input slice is
-                            // consistent. This would happen automatically, but I'm making it
-                            // explicit for clarity.
+                            // The Drop impl will copy the last completed buffer into the slice.
                             drop(buffer);
                             panic!(
                                 "The key function is not reliable: when called repeatedly, \
@@ -203,15 +191,13 @@ macro_rules! sort_impl {
                     }
 
                     unsafe {
-                        // This is safe, we just ensured that the write buffer is consistent.
+                        // SAFETY: we just ensured that every index was written to.
                         buffer.swap();
                     }
                 }
             }
 
-            // In case the result ended up in the temporary buffer, the Drop impl will copy it over
-            // to the input slice. This would happen automatically, but I'm making it explicit for
-            // clarity.
+            // The Drop impl will copy the last completed buffer into the slice.
             drop(buffer);
         }
     };
@@ -230,98 +216,3 @@ macro_rules! radix_key_impl {
 }
 
 radix_key_impl! { u8 u16 u32 u64 u128 }
-
-/// Double buffer. Allocates a temporary memory the size of the slice, so that
-/// elements can be freely reordered from buffer to buffer.
-///
-/// # Consistency
-///
-/// For the purposes of this struct, buffer in a consistent state contains a
-/// permutation of elements from the original slice. In other words, elements
-/// can be reordered, but not duplicated or lost.
-///
-/// `read_buf` is always consistent. Before calling `swap`, the caller must
-/// ensure that `write_buf` is also consistent.
-///
-/// # Drop behavior
-///
-/// Drop impl ensures that the slice this buffer was constructed with is left in
-/// a consistent state. If the input slice ended up as `write_buf`, the
-/// temporary memory (which is now `read_buf` and therefore consistent) is
-/// copied into the slice and the buffers are swapped.
-struct DoubleBuffer<'a, T> {
-    slice: &'a mut [T],
-    _aux: Vec<T>,
-
-    /// Read buffer is read-only and always consistent.
-    read_buf: *const T,
-
-    /// Write buffer is write-only. Elements can be present multiple times or
-    /// not at all. The caller must ensure that it is consistent before calling
-    /// `swap`.
-    write_buf: *mut T,
-}
-
-impl<'a, T> DoubleBuffer<'a, T> {
-    fn new(slice: &'a mut [T]) -> DoubleBuffer<'a, T> {
-        let mut aux = Vec::with_capacity(slice.len());
-        let read_buf = slice.as_ptr();
-        let write_buf = aux.as_mut_ptr();
-        DoubleBuffer {
-            // Hold on to the &mut slice to make sure it outlives the pointer
-            // and to prevent writes from the outside
-            slice,
-            // Hold on to the Vec to make sure it outlives the pointer
-            _aux: aux,
-            read_buf,
-            write_buf,
-        }
-    }
-
-    /// Returns a ref to an element from the read buffer.
-    ///
-    /// Caller must ensure that `index` is in (0..len).
-    #[inline(always)]
-    unsafe fn read(&self, index: usize) -> &T {
-        &*self.read_buf.add(index)
-    }
-
-    /// Copies the referenced element into the write buffer.
-    ///
-    /// Caller must ensure that `index` is in (0..len).
-    #[inline(always)]
-    unsafe fn write(&self, index: usize, t: &T) {
-        self.write_buf
-            .add(index)
-            .copy_from_nonoverlapping(t as *const T, 1);
-    }
-
-    /// Swaps the read and write buffers.
-    ///
-    /// Caller must ensure that the write buffer is consistent before calling
-    /// this function.
-    unsafe fn swap(&mut self) {
-        // The cast is ok, we have an exclusive access to both buffers
-        // (&mut [T] and Vec<T>). The caller guarantees that the write buffer is
-        // consistent and therefore it's safe to read from it and use it as a
-        // read buffer.
-        let temp = self.write_buf as *const T;
-        self.write_buf = self.read_buf as *mut T;
-        self.read_buf = temp;
-    }
-}
-
-impl<'a, T> Drop for DoubleBuffer<'a, T> {
-    fn drop(&mut self) {
-        let input_slice_is_write = self.write_buf as *const T == self.slice.as_ptr();
-        if input_slice_is_write {
-            // Input slice is the write buffer, copy the consistent state from the read buffer
-            unsafe {
-                // This is safe, `read_buf` is always consistent and the length is the same.
-                self.write_buf
-                    .copy_from_nonoverlapping(self.read_buf, self.slice.len());
-                self.swap();
-            }
-        }
-    }
-}