Replacing an element of a vector with the result of a closure

help
1

Suppose I want to replace an element at a specific index in a vector with the result of a closure, i.e. the following

fn replace_with<T, F>(vec: &mut Vec<T>, idx: usize, f: F)
where
    F: FnOnce(T) -> T,
{
    let e = vec.remove(idx);
    vec.insert(idx, f(e))
}

Of course this is quite inefficient since it shifts all of the elements twice.

It would be great to do this swap in-place. What is the best way to do this?

(In practice, I actually want something like

fn replace_with_iter<T, F, I>(vec: &mut Vec<T>, idx: usize, f: F)
where
    I: IntoIterator<Item = T>,
    F: FnOnce(T) -> I,
{
    // todo!
}

but for purposes of this question the simpler version contains the key elements of this problem that I don't totally understand.)

Using my rudimentary understanding of unsafe I wrote the following. I understand that the main concern is that the closure might panic, in which case we need to avoid the double-free of the removed element.

fn replace_with<T, F>(vec: &mut Vec<T>, idx: usize, f: F)
where
    F: FnOnce(T) -> T,
{
    assert!(idx < vec.len());
    unsafe {
        let original_len = vec.len();

        // now the vector is only responsible for elements 0..idx
        vec.set_len(idx);

        // update the element
        let p = vec.as_mut_ptr().add(idx);
        let e = std::ptr::read(p);
        let new = f(e);
        std::ptr::write(p, new);

        // this is valid since the element at `idx` is once again initialized
        vec.set_len(original_len);

    };
}

However, if f panics, then it seems to me that destructor for the elements in the vector at indices larger than idx will not run. Is there a good way to avoid this? In general, is this something that I should try to handle, or is panic usually treated in such a way that memory leaks like this are considered fine?

2

Swap in the last element of the vector temporarily:

fn replace_with<T, F>(vec: &mut Vec<T>, idx: usize, f: F)
where
    F: FnOnce(T) -> T,
{
    let x = vec.swap_remove(idx);
    let last_idx = vec.len();
    vec.push(f(x));
    vec.swap(idx, last_idx);
}
4 Likes
3

Alice suggested the use of the take_mut crate for a use case like this one: Replacing element of vector - #4 by alice

1 Like
4

Thanks, this is a great point! Now I remember Nico Matsakis' article about this pattern.

I am still curious about the correct handling of leaked memory in general in the presence of panics.

5

Thanks for the linked thread about this!

6

replace_with is an alternative (take_mut is unmaintained I believe).

5 Likes
7

Don't forget you can simplify it a lot if T: Copy.

    let e = f(vec[idx]);
    vec[idx] = e;
1 Like