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)

Obviously Dogberts solution is superior, but I figured out that the while loop is not necessary when using `into_iter()`, which transfer ownership. See [Playground](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=43fd455fedff4f51c89c21f598f0fe5f) — Cyrill, Jun 13 '22 at 20:47
I'm pretty sure OwnedDataWithReference is cheaper to clone than to construct, as its trivially clonable. struct with reference in it copy the reference on clone, not the data being it. so its basicly just cloning a reference, its as cheap as copying a usize. — Bamontan, Jun 14 '22 at 01:25
@Bamontan that’s probably true in this case, but in my case I actually have way more references as well as owned data in that struct. Thats why I was hesitant. Also, it seems hacky to clone something that only exists once. — Cyrill, Jun 15 '22 at 06:47

score 2 · Accepted Answer · answered Jun 13 '22 at 16:12

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

Yessss this was exactly it! I had tried it with the `vec!` macro and I didn't know about this restriction. Thanks a lot :) — Cyrill, Jun 13 '22 at 16:37

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

1 Answers1