Feature item_find_many

[0xForerunner]

Proposal

Problem statement

This feature would allow finding, returning, and potentially mapping, multiple &muts within a single iterator using a convenient const generic array syntax

Motivation, use-cases

Credit to @cuviper on the rust internals forum for building out most of this implementation.

This is a very general addition and could support many use cases. Currently it is rather difficult and cumbersome to find multiple &muts from an iterator.

Solution sketches

This feature would introduce 2 new methods on Iterator<Item = &mut T>.

find_many takes an iterator and return multiple mutable references to the items which match the predicate.

Example

let mut v = vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)];
let [left, right] =
v.find_many([&2, &3], |item, key| &item.0 == key).unwrap();
assert_eq!(*left, (2, 3));
assert_eq!(*right, (3, 4));

find_map_many takes an iterator and returns multiple mutable references to the items which match the predicate. The items which match the predicate are then mapped to a different &mut. This is particularly useful when doing something similar to a key value search.

Example

let mut v = vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)];
let [left, right] =
v.find_map_many([&2, &3], |item, key| &item.0 == key, |item| &mut item.1).unwrap();
assert_eq!(*left, 3);
assert_eq!(*right, 4);

For both methods an option containing an array to the &muts would be returned. The None variant would represent cases where there are no matching items in the iterator for every key. Each key provided requires a unique Item.

This feature will make make handling mutable iterators much more ergonomic and should prevent programmers from rearranging code in unintuitive ways just to make the borrow checker happy. This should result in more natural layout of functions.

Reference-level explanation

These methods rely on two nightly features of rust to be enabled. array_try_map and inline_const.

#![feature(array_try_map, inline_const)]

pub fn find_many<'a, I, T, F, K, const LEN: usize>(collection: I, keys: [&K; LEN], mut predicate: 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];

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

    output.try_map(|opt| opt)
}

#![feature(array_try_map, inline_const)]

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

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

    output.try_map(|opt| opt)
}

[ajwock]

Re your first example, this is quite close to your find_many.

fn main() {
    let v = vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)];
    // This demonstrates how you can shorten the syntax, but it's returning a Vec<&mut (i32, i32)>
    let vec: Vec<_> = v.iter_mut().filter_map(|item| [2, 3].contains(&item.0).then_some(item)).collect();
    assert_eq!(*vec[0], (2, 3));
    assert_eq!(*vec[1], (3, 4));
}

Filter_map and bool::then_some can do this well, and it's only a couple of characters away from being like your second example and equivalent to find_map_many.

Another idiosyncracy this array-based interface has: What if one key has multiple matches? It's unclear what would happen there from a high level, but in your implementation it seems like the output value would simply get overwritten, which would be surprising behavior to me as a user.

The array interface removes flexibility in several ways. You could only use a compile-time-fixed amount of keys, and they would be checked with linear sweep search. There's no option to instead find the keys based on, say, a HashSet or a sorted Vec instead. Whereas filter_map already exposes that possibility due to the wonderful capabilities of closures:

use std::collections::HashSet;

fn main() {
    let mut v = vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)];
    let mut keys = HashSet::new();
    keys.insert(2);
    keys.insert(3);
    let vec: Vec<&mut i32> =
    v.iter_mut().filter_map(|item| keys.contains(&item.0).then_some(&mut item.1)).collect();
    assert_eq!(*vec[0], 3);
    assert_eq!(*vec[1], 4);
}

[the8472]

#![feature(iter_next_chunk)]
fn main() {

    let mut v = vec![1; 20];
    let it = v.iter_mut();
    
    let many = it.filter(|i| **i > 0).next_chunk::<5>().unwrap();
    *many[0] = 2;
    *many[1] = 3;
    dbg!(v);
}

prints

[src/main.rs:10] v = [
    2,
    3,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
]

[Amanieu]

As @the8472 has shown, there are already ways to achieve this using existing (although unstable) APIs. The proposed API is quite complex and very specific, so we would prefer to not have this and instead encourage the use of other, more general methods on Iterator.

[0xForerunner]

@Amanieu Neither of the solutions above actually solve the problem at hand.

The solution from @the8472 doesn't actually attempt to match to keys so isn't actually relevant.

And the solution from @ajwock seems to work, except the returned Vec<&mut> is not sorted correctly, so the the caller doesn't know which index corresponds to which key, which is sort of the whole point. I don't believe there is a reasonably convenient way to solve this problem with standard methods.

Thoughts @Amanieu ?

[0xForerunner]

I will add that in particular find_map_many is quite a busy api, thought I find find_map to be fairly concise.

[0xForerunner]

@ajwock I think the expected behaviour should be the same as find where a reference to the first matching value is returned. That is indeed the way my proposal works. Your point about have the keys as a Hashset is sort of valid. If we are using a const array then likely its length is not very large, in which case a hashset wouldn't be useful. As well I think HashMap already has a find_many implementation.

[the8472]

The solution from @the8472 doesn't actually attempt to match to keys so isn't actually relevant.

    let mut keys = HashSet::from([1, 2, 3])
    let many = it.filter(|i| keys.remove(**i)).next_chunk::<5>().unwrap();

[0xForerunner]

The solution from @the8472 doesn't actually attempt to match to keys so isn't actually relevant.

    let mut keys = HashSet::from([1, 2, 3])
    let many = it.filter(|i| keys.remove(**i)).next_chunk::<5>().unwrap();

Better, but it still doesn't solve the problem since the result is not sorted correctly, so you can't easily associated each ref to each key.

[0xForerunner]

To make it more clear, this is the exact problem you're trying to solve:

enum Fruit {
    Pineapple,
    Pear,
    Apple,
    Orange,
    Grapefruit,
}
fn main() {
    use Fruit::*;
    let mut v = vec![Pineapple, Grapefruit, Pear, Orange, Apple];

    let pineapple_ref: &mut Fruit = todo!();
    let pear_ref: &mut Fruit = todo!();
    assert!(matches!(*pineapple_ref, Pineapple));
    assert!(matches!(*pear_ref, Pear));
}