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?

6 posts - 4 participants

Read full topic