rustslice

How do I sort a sequence of slices in a vector?


I have a data structure where conceptually, I have a sequence of tuples of integers, all of the same size. However, this size is not known at compile time, so I can't use actual tuples or arrays or the like: those need their length to be specified in the type. This is similar to the chunks method on an iterator, except I need random access to the "tuples". So I'd like to implement this as a Vec of integers, accessing these "tuples" as consecutive slices of the Vec. For example, the sequence (1, 2, 4), (1, 2, 3), (2, 2, 1) would be represented as the vector [1, 2, 4, 1, 2, 3, 2, 2, 1], with the "chunk size" stored separately and accessing the ith slice as v[i * chunk_size .. (i+1) * chunk_size]. I have some of this working.

Now I'd like to sort the sequence of slices lexicographically, turning the above into (1, 2, 3), (1, 2, 4), (2, 2, 1) conceptually, meaning the vector would be turned into [1, 2, 3, 1, 2, 4, 2, 2, 1]. Of course I could implement my own simple sorting algorithm to do this, but if possible I'd like to wrap the slices up into some struct, implement Ord, and then call sort_unstable(). Something like this:

#[derive(PartialEq, Eq, Ord, PartialOrd)]
struct Chunk<'a>
{
    slice: &'a [usize],
}

impl<'a> Chunk<'a> {
    /// Consider chunks as a collection of slices of length chunk_size. Sort
    /// this collection in place, leaving each slice intact. This will panic
    /// if the length of `chunks` is not a multiple of `chunk_size`.
    pub fn sort_slice_of_chunks(mut chunks: &[usize], chunk_size: usize) {
        let n_slices = chunks.len() / chunk_size();
        let mut chunk_objects: Vec<_> = (0 .. n_slices)
           .map(|i| Chunk {
                slice: &chunks[i * chunk_size..(i + 1) * chunk_size],
            })
            .collect();
        chunk_objects.sort_unstable();
    }
}

However, this will merely sort the vector chunk_objects of my chunk objects, rather than the underlying slice. Is there a way to make sort_unstable() modify the underlying slice, or is this the wrong approach? Should I bite the bullet and implement introsort?


Solution

  • To sort in-place using the slice methods, the values you're sorting need to have a known size at compile time, because under the hood the sort happens as a sequence of swaps, which need to know the size of the things being swapping to emit the code to perform the swap. This means you must know the length of each subarray in order to do an in-place sort using the built-in slice sort methods. If Rust let you specialize the swap operation somehow then you could do this with a runtime-provided subarray length using a newtype, but this is not possible on stable Rust today.

    As far as I can determine, there's two ways you can somewhat-efficiently handle this without resorting to building your own in-place sort:

    1. You can collect a Vec of slices and sort those (as you do here) and then apply that sorting to the original Vec. This requires a low, constant number of allocations.
    2. You can consume the original array to build up a Vec<Vec<_>>, sort that, then flatten it back. This is easy and does not require unsafe code, but requires N+2 extra allocations.

    How would one tackle the first approach? Well, we can't just swap stuff around in the underlying array because the borrow checker won't let us. There is a way to handle this using a lot of unsafe code, but perhaps simpler would be to "make our own slices" by collecting all of the start indices of each subarray. We don't need the length because we already have it (chunk_size) so this also requires half the memory of using slices.

    We still need some unsafe code to do the actual swapping (unless we can constrain T: Copy -- more on that later).

    Let's build this function:

    pub fn sort_chunks<T: Ord>(chunks: &mut [T], chunk_size: usize);
    

    First, we need to collect the indices that start each subarray.

    let mut chunk_indices: Vec<_> = (0..chunks.len()).step_by(chunk_size).collect();
    

    Then sort those based on how the corresponding subsections of the original array would sort.

    chunk_indices.sort_unstable_by(|&a, &b| chunks[a..(a + chunk_size)].cmp(&chunks[b..(b + chunk_size)]));
    

    Ok... now what? We can't blindly move stuff around because of the classic problem where some of the rearrangements form a cycle, such as subslice 0 belongs in position 1 and subslice 1 belongs in position 0.

    Other than implementing a really complicated algorithm to handle this, the simplest approach is to create a new allocation to temporarily hold the elements in their new order, then write them back.

    We need a bit of unsafe code here. We are going to "steal" the contents of each T into the new array and then write them back. To ensure this is panic-safe, we want to preallocate space to hold the temporary copies before we start blitting data around. If we don't do this, a panic triggered by a failed reallocation could result in a double-free, and we can't have that.

    let mut new_order: Vec<T> = Vec::with_capacity(chunk_indices.len() * chunk_size);
    

    With this space allocated, we can blit the T's out of chunks and into this new Vec. We reorder them into the final order as we do this.

    new_order.extend(
        chunk_indices
            .into_iter()
            .flat_map(|i| (i..(i + chunk_size)).map(|i| unsafe { std::ptr::read(&chunks[i]) })),
    );
    

    Finally, we blit them back into the source slice.

    for (dest, src) in chunks.iter_mut().zip(new_order) {
        unsafe {
            std::ptr::write(dest, src);
        }
    }
    

    Note this consumes the temporary copies without dropping them, as we move the copies into std::ptr::write, which is documented not to drop the source value.

    All together:

    pub fn sort_chunks<T: Ord>(chunks: &mut [T], chunk_size: usize) {
        let mut chunk_indices: Vec<_> = (0..chunks.len()).step_by(chunk_size).collect();
    
        chunk_indices
            .sort_unstable_by(|&a, &b| chunks[a..(a + chunk_size)].cmp(&chunks[b..(b + chunk_size)]));
    
        // Allocate before we do this stuff to ensure we can't panic.
        let mut new_order: Vec<T> = Vec::with_capacity(chunk_indices.len() * chunk_size);
    
        new_order.extend(
            chunk_indices
                .into_iter()
                .flat_map(|i| (i..(i + chunk_size)).map(|i| unsafe { std::ptr::read(&chunks[i]) })),
        );
    
        for (dest, src) in chunks.iter_mut().zip(new_order) {
            unsafe {
                std::ptr::write(dest, src);
            }
        }
    }
    

    To my knowledge, this is sound even with non-Copy types. It's sound to blit non-Copy types (otherwise std::mem::swap would be unsound) and soundness issues from having two non-Copy values blitted from the same "source" are dependent on the implementation of T and things that might use a T. We never actually call into any implementation of T here (including its Drop implementation) or anything that uses T with knowledge of what it is.

    Note that we could "harden" the code against any unforeseen panics by changing new_order to be a Vec<ManuallyDrop<T>> instead.

    However, you can eliminate the unsafe code if you constrain T: Copy.

    pub fn sort_chunks<T: Copy + Ord>(chunks: &mut [T], chunk_size: usize) {
        let mut chunk_indices: Vec<_> = (0..chunks.len()).step_by(chunk_size).collect();
    
        chunk_indices
            .sort_unstable_by(|&a, &b| chunks[a..(a + chunk_size)].cmp(&chunks[b..(b + chunk_size)]));
    
        let new_order: Vec<T> = chunk_indices
            .into_iter()
            .flat_map(|i| (i..(i + chunk_size)).map(|i| chunks[i]))
            .collect();
    
        for (dest, src) in chunks.iter_mut().zip(new_order) {
            *dest = src;
        }
    }