rust-itertools / itertools

Extra iterator adaptors, iterator methods, free functions, and macros.
https://docs.rs/itertools/
Apache License 2.0
2.72k stars 309 forks source link

Feature item_find_many #779

Open ewoolsey opened 1 year ago

ewoolsey commented 1 year ago

Copy and pasted from the libs-team api change proposal which likely isn't going through. Would this crate be a good place for an addition like this?

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)
}
Philippe-Cholet commented 1 year ago

Currently, this crate uses "stable Rust" only and our MSRV does not allow us to use "const generics". "Const generics" might be usable in a few months (I guess I'm not sure) when we raise the MSRV but it seems very unlikely to me we would use nightly features (from my understanding mainly to limit the maintenance burden).

ewoolsey commented 1 year ago

Hey thanks for the info @Philippe-Cholet . Const generics are obviously necessary for this to work, but this could still be implemented without the rest of the nightly features (although not quite as cleanly). Perhaps let's pick this conversation back up when const generics are usable.