CODEWITHSUNDEEP
rustborrow-checker
xixixao
xixixao

Reputation: 563

Conditionally sort a Vec in Rust

Let's say I want to sort a Vec of non-Clone items - but only maybe (this is a boiled down example of an issue in my code).

My attempt would be something like:

fn maybe_sort<T>(x: Vec<T>) -> Vec<T>
where
    T: std::cmp::Ord,
{
    // First, I need a copy of the vector - but only the vector, not the items inside
    let mut copied = x.iter().collect::<Vec<_>>();
    copied.sort();
    // In my actual code the line below depends on the sorted vec
    if rand::random() {
        return copied.into_iter().map(|x| *x).collect::<Vec<_>>();
    } else {
        return x;
    }
}

Alas the borrow checker isn't happy. I have a shared reference to each item in the Vec, and although I am not ever returning 2 references to the same item, Rust can't tell.

Is there a way to do this without unsafe? (and if not, what's the cleanest way to do it with unsafe.

Upvotes: 0

Views: 737

Answers (2)

Jmb
Jmb

Reputation: 23414

If T implements Default, you can do it with a single sort and without unsafe like this:

fn maybe_sort<T: Ord + Default> (mut x: Vec<T>) -> Vec<T> {
    let mut idx = (0..x.len()).collect::<Vec<_>>();
    idx.sort_by_key (|&i| &x[i]);
    if rand::random() {
        return x;
    } else {
        let mut r = Vec::new();
        r.resize_with (x.len(), Default::default);
        for (i, v) in idx.into_iter().zip (x.drain(..)) {
            r[i] = v;
        }
        return r;
    }
}

Playground

If T does not implement Default, the same thing can be done with MaybeUninit:

use std::mem::{self, MaybeUninit};

fn maybe_sort<T: Ord> (mut x: Vec<T>) -> Vec<T> {
    let mut idx = (0..x.len()).collect::<Vec<_>>();
    idx.sort_by_key (|&i| &x[i]);
    if rand::random() {
        return x;
    } else {
        let mut r = Vec::new();
        r.resize_with (x.len(), || unsafe { MaybeUninit::uninit().assume_init() });
        for (i, v) in idx.into_iter().zip (x.drain(..)) {
            r[i] = MaybeUninit::new (v);
        }
        return unsafe { mem::transmute::<_, Vec<T>> (r) };
    }
}

Playground

Finally, here's a safe solution which doesn't require T to implement Default, but allocates an extra buffer (there is theoretically a way to reorder the indices in place, but I'll leave it as an exercise to the reader ☺):

fn maybe_sort<T: Ord> (mut x: Vec<T>) -> Vec<T> {
    let mut idx = (0..x.len()).collect::<Vec<_>>();
    idx.sort_by_key (|&i| &x[i]);
    if rand::random() {
        let mut rev = vec![0; x.len()];
        for (i, &j) in idx.iter().enumerate() {
            rev[j] = i;
        }
        for i in 0..x.len() {
            while rev[i] != i {
                let j = rev[i];
                x.swap (j, i);
                rev.swap (j, i);
            }
        }
    }
    x
}

Playground

Upvotes: 1

kmdreko
kmdreko

Reputation: 60492

You can .enumerate() the values to keep their original index. You can sort this based on its value T and decide whether to return the sorted version, or reverse the sort by sorting by original index.

fn maybe_sort<T: Ord>(x: Vec<T>) -> Vec<T> {
    let mut items: Vec<_> = x.into_iter().enumerate().collect();
    items.sort_by(|(_, a), (_, b)| a.cmp(b));
    
    if rand::random() {
        // return items in current order
    }
    else {
        // undo the sort
        items.sort_by_key(|(index, _)| *index);
    }
    
    items.into_iter().map(|(_, value)| value).collect()
}

Upvotes: 3

Related Questions