function for Iterator should look like this
pub fn get_muts<K, F, const LEN: usize>(&mut self, f: F) -> Result<[&mut T; LEN]>
where
Self: Sized,
F: FnMut(Self::Item, [&mut K; LEN]) -> bool,
Is this posible? I think it would be hugely useful.
Could you write possible implementation on pseudo-code (without corresponding Rust's borrrow-checker rules), please?
Here is a horrible prototype of this function. Should get the general point across
pub fn get_muts<'a, T, K, F, I, const LEN: usize>(collection: I, keys: [&K; LEN], mut f: F) -> [&'a mut T; LEN]
where
T: Debug,
I: Sized + IntoIterator<Item = &'a mut T>,
F: FnMut(&I::Item, &K) -> bool,
{
let mut output = Vec::with_capacity(LEN);
unsafe { output.set_len(LEN) }
for item in collection.into_iter() {
for (i, key) in keys.into_iter().enumerate() {
if f(&item, key) {
output[i] = item;
}
}
}
output.try_into().unwrap()
}
There are a number of things that areproblematic here, But I hope this serves as a starting point for someone who is more comfortable with unsafe rust than I am.
The unstable HashMap::get_many_mut is a similar idea.
I didn't even know that existed, but yes the idea is exactly the same, this is just for iterators instead of hash maps.
There's a few options for slices, you can use a combination of split_at_mut/split_first_mut/split_last_mut or slice patterns (playground)
fn main() {
let mut v = vec![1, 2, 3, 4, 5];
let Some((a, rest)) = v.split_first_mut() else { unreachable!() };
let [_, b, c, ..] = rest else { unreachable!() };
*a = 7;
*b = 8;
*c = 9;
dbg!(v);
}
(it would be nice if there were ways to avoid the fallibility of the pattern here, e.g. with v[2..4] that could return an &mut [i32; 2] rather than an &mut [i32] so that the fallibility is all handled within the panicking indexing operation, but I don't know that current const-generics are powerful enough to implement something like that).
For iterators that give out mutable references like vec::IterMut you can simply collect the references (playground):
fn main() {
let mut v = vec![1, 2, 3, 4, 5];
let mut refs = v.iter_mut().collect::<Vec<&mut i32>>();
*refs[1] = 7;
*refs[2] = 8;
dbg!(v);
}
For by-value iterators, you can similarly just collect part of the iterator then operate on it, there's not really anything more efficient the iterator itself can do.
Along a similar line I've implemented a very useful function for Vec which allows for getting multiple indices. It could probably be improved with some unsafe code:
#[inline]
fn get_many_mut<'a, T, I, const LEN: usize>(vec: &'a mut Vec<T>, indexes: [usize; LEN]) -> Option<[&'a mut T; LEN]>
where
T: Debug,
&'a mut T: Copy,
I: Sized + IntoIterator<Item = &'a mut T>,
{
// Create uninitialized array
let mut data: [Option<&mut T>; LEN] = [None; LEN];
// Fill the data
for (vec_index, item) in vec.iter_mut().enumerate() {
for (i, index) in indexes.into_iter().enumerate() {
if index == vec_index {
data[i] = Some(item);
break;
}
}
}
// Assert initialized
let output = data.into_iter().collect::<Option<Vec<&mut T>>>();
output.map(|vec| vec.try_into().unwrap())
}
So I'm starting to hone in on this. My latest iteration looks like this and is actually pretty clean I think.
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: Sized + IntoIterator<Item = &'a mut T>,
T: Debug,
F: FnMut(&I::Item, &K) -> bool,
{
let mut refs = collection.into_iter();
keys.iter()
.map(|key| refs.find(|elem| f(elem, key)))
.collect::<Option<Vec<&mut T>>>()
.map(|x| x.try_into().unwrap())
}
I still think there are some improvements to be done with unsafe, but this is looking pretty good I think.
Here's some tests to see it in action.
#[test]
fn test_find_many() {
let mut v = vec![(0, 1), (0, 2), (0, 3), (0, 4), (0, 5)];
let keys = [&2, &3];
let res = find_many(&mut v, keys, |item, key| &item.1 == key);
let unwrapped = res.unwrap();
assert_eq!(*unwrapped[0], (0, 2));
assert_eq!(*unwrapped[1], (0, 3));
}
#[test]
fn test_duplicate_key() {
let mut v = vec![(0, 1), (0, 2), (0, 3), (0, 4), (0, 5)];
let keys = [&2, &2];
let res = find_many(&mut v, keys, |item, key| &item.1 == key);
assert!(res.is_none());
}
#[test]
fn test_missing_key() {
let mut v = vec![(0, 1), (0, 2), (0, 3), (0, 4), (0, 5)];
let keys = [&2, &7];
let res = find_many(&mut v, keys, |item, key| &item.1 == key);
assert!(res.is_none());
}
Note that your implementation will only find everything if collection and keys have them in the same order. For example, keys [&3, &2] will advance the refs iterator to the first 3, but then it won't find 2 because you already passed it. Iterators won't let you revisit prior items, especially because that would allow aliasing with &mut T items.
Damn... yup you're right. Any idea how to fix it?
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)
}
Yeah that's a fantastic approach. I didn't know about either of those features. Those are super useful for this.Great job!
This is very useful, what does it takes to add this to the stdlib as a method of Vec? Could you just send a PR?
I'm thinking a PR to Itertools could potentially go through.
Building on top. Another useful function to map the &mut could look like this:
/// finds multiple items in a collection and maps the elements to &muts.
///
/// ```
/// # use contracts_common::map::find_map_many;
/// # fn test_find_many() {
/// let mut v = vec![(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)];
/// let [left, right] =
/// find_map_many(&mut v, [&2, &3], |item, key| &item.0 == key, |item| &mut item.1).unwrap();
/// assert_eq!(*left, 3);
/// assert_eq!(*right, 4);
/// # }
/// ```
pub fn find_map_many<'a, I, T, U, F, M, K, const LEN: usize>(
collection: I, keys: [&K; LEN], mut find: F, mut map: M,
) -> Option<[&'a mut U; LEN]>
where
I: IntoIterator<Item = &'a mut T>,
T: Debug + 'a,
U: Debug,
F: FnMut(&T, &K) -> bool,
M: FnMut(&'a mut T) -> &'a mut U,
{
let mut remaining = LEN;
// When inline_const is stabilized this unsafe block can be removed
let mut output: [Option<&'a mut U>; LEN] = unsafe { std::mem::zeroed() };
'collection: for elem in collection {
for (key, out) in std::iter::zip(&keys, &mut output) {
if out.is_none() && find(elem, key) {
*out = Some(map(elem));
remaining -= 1;
if remaining == 0 {
break 'collection;
}
break;
}
}
}
// When array_try_map is stabilized this can be made better.
let vec = output.into_iter().collect::<Option<Vec<&mut U>>>();
vec.map(|x| x.try_into().unwrap())
}
That std::mem::zeroed should be fine, but you can avoid it either by using std::array::from_fn:
let mut output: [Option<&'a mut U>; LEN] = std::array::from_fn(|_| None);
Or using a const (needs some magic to make it generic though):
struct Helper<T>(std::marker::PhantomData<T>);
impl<T> Helper<T> {
const NONE: Option<T> = None;
}
let mut output: [Option<&'a mut U>; LEN] = [Helper::NONE; LEN];
You can also avoid collecting into a Vec if you first check the Options and then unwrap them inside an array::map:
let () = output.iter().map(|r| r.as_ref().map(|_| ())).collect::<Option<()>>()?;
Some(output.map(Option::unwrap))
Does it make sense to define an unsafe trait DisjointIterator: Iterator that promises to never yield the same value twice? Then these sort of multiple-mutable-fetch operations can use that as a bound instead of relying on their own internal checks.
Integer ranges, set iterators, and dictionary key iterators could all implement it, as could filters of these.
For that you also need a notion of equality that you can trust. User types can implement Eq in such a way that they appear different when being yielded from the iterator but equal afterwards. This is particularly important if you want to implement it for sets and dictionary keys iterators.
But even then, I'm not sure if this is really an improvement. It would forse users to statically prove that their keys are disjoint (What if you want to use [&1, &3, &4] for the keys?) and also wouldn't really improve the implementation (the break at the end of the if is required by the borrow checker, but is also an optimization which you likely want to keep).
It would force the disjointness check outside of the actual multiple-borrowing code, but it doesn't necessarily need to be statically proven: There would likely need to be a helper for dynamic checking, which could be omitted in those cases that are statically provable. For the case of slices like [&1, &3, &4], that helper could be implemented as a const fn so that the checking happens in the constant-folding optimization pass. I agree that it wouldn't improve matters much for this specific function because the Iterator interface forces considering each candidate in sequence— It would be more useful for inherent methods on the original collection types.
That may kill this idea as a practical matter. std couldn't have blanket TrustedIterator implementations, only ones for concrete types known to implement Eq sanely. The only way to get the blanket implementation back would be to add some kind of unsafe trait TrustedEq {}, which feels like the start of a potentially-very-deep rabbit hole.