Is there a nicer way to create a vector from indexed data in rust?

Question

I have the problem, that I get some data from an iterator (specifically a bevy Query) and part of that data is an index. I now want to create a vector with this data, where every element is placed at the index it is intended to go.

To illustrate the problem, this is how I might solve this in python:

def iterator():
    # The iterator doesn't necessarily give back the data in the correct order
    # However, the indices are always dense from 0...n
    yield ("A", 0)
    yield ("C", 2)
    yield ("B", 1)

def some_computation(x):
    return x + " + some other data"

data_list = [0] * 3 # I know the length
for data, index in iterator():
    data_list[index] = some_computation(data)

print(data_list)

Below is how I currently implemented this in rust, but I'm not super happy about it. Firstly, it takes O(n log n) time. Secondly, its feels difficult to read, especially the conversion to get rid of the index in the data. Thirdly, I haven't yet figured out how to extract this into its own function without violating lifetimes.

Open in the Rust Playground

// I have so far avoided deriving Clone on this struct.
#[derive(Debug)]
struct OwnedDataWithReference<'a> {
    data: &'a str,
}

fn do_something_with_data_list_and_drop_it(x: Vec<OwnedDataWithReference>) {
    println!("{:?}", x);
}

fn main() {
    // In my application this is a bevy query, which can't be indexed.
    let iterator = vec![("A", 0), ("C", 2), ("B", 1)];

    // First the data is created and stored in a vector with a index associated
    // and only subsequently sorted, which costs runtime, but has nice ownership.
    let mut data_list_with_index = Vec::new();
    for (data_ref, index) in iterator.iter() {
        // In my application I first have to fetch data using data_ref
        let new_data = OwnedDataWithReference {
            data: data_ref,
        };
        data_list_with_index.push((new_data, index));
    }
    // This reverse sort and stack based map is needed because I don't want to
    // make `OwnedDataWithReference` Clone. This way, I can transfer ownership.
    data_list_with_index.sort_by_key(|(_, i)| -*i);
    let mut data_list = Vec::new();
    while !data_list_with_index.is_empty() {
        data_list.push(data_list_with_index.pop().unwrap().0);
    }

    // In my application I then serialize this with serde.
    do_something_with_data_list_and_drop_it(data_list);
}

I have tried an approach with using Vec<Option<OwnedDataWithReference>>, but this complained that OwnedDataWithReference is not Clone, which I've tried to avoid so far.

Please let me know if you can find a nicer solution to this problem, or if you think I should derive Clone. (The problem is still not solved then, but I think it gets easier)

Answer 1

I have tried an approach with using Vec<Option<OwnedDataWithReference>>, but this complained that OwnedDataWithReference is not Clone, which I've tried to avoid so far.

You probably tried creating a Vec of None values using the vec! macro which requires T: Clone. You can instead create it using another way, for example (0..n).map(|_| None).collect(). With that, this can be done like this:

#[derive(Debug)]
struct OwnedDataWithReference<'a> {
    data: &'a str,
    other_stuff: u32,
}

fn do_something_with_data_list_and_drop_it(x: Vec<OwnedDataWithReference>) {
    println!("{:?}", x);
}

fn main() {
    let iterator = vec![("A", 0), ("C", 2), ("B", 1)];

    let mut data_list = (0..iterator.len()).map(|_| None).collect::<Vec<_>>();

    for (data_ref, index) in iterator.iter() {
        let new_data = OwnedDataWithReference {
            data: data_ref,
            other_stuff: 42,
        };
        data_list[*index] = Some(new_data);
    }
    
    // This will panic if any element is still `None` at this point.
    let data_list = data_list.into_iter().map(Option::unwrap).collect();
    
    do_something_with_data_list_and_drop_it(data_list);
}

Playground