Get multiple &mut from a vector, or any iterator

12

Okay, you nerd-sniped me. I used a couple nightly features, but no unsafe:

#![feature(array_try_map, inline_const)]

pub fn find_many<'a, I, T, F, K, const LEN: usize>(
    collection: I,
    keys: [&K; LEN],
    mut f: F,
) -> Option<[&'a mut T; LEN]>
where
    I: IntoIterator<Item = &'a mut T>,
    F: FnMut(&T, &K) -> bool,
{
    let mut remaining = LEN;
    let mut output = [const { None::<&'a mut T> }; LEN];

    if LEN > 0 {
        'collection: for elem in collection {
            for (key, out) in std::iter::zip(&keys, &mut output) {
                if out.is_none() && f(elem, key) {
                    *out = Some(elem);
                    remaining -= 1;
                    if remaining == 0 {
                        break 'collection;
                    }
                    break;
                }
            }
        }
    }

    // map [Option<&mut T>; LEN] to Option<[&mut T; LEN]>
    output.try_map(|opt| opt)
}