Implementing a find_map iterator for multiple searches

code review
1

the code is here

the issue i'm trying to solve

there's find_map: consume an iterator, evaluate a FnMut, first time a Some gets returned abort and return it.

  • ✓ no need to loop through the entire range
  • ✓ if you need to do some computation with the found result and the computation is somewhat related to the search criterion, you don't need to do work twice, you can do find and map in one call

now what if i have >1 independent things that need to be find_maped? I could do

let a = iter.find_map(fn1);
let b = iter.find_map(fn2);

but that would loop through the start of the range twice. I also don't know if both finds will actually find anything (a and/or b might be None) and i don't know if a or b will be found first.

side quest: let me know if i'm missing an existing solution.

what i did

I thought that could be reasonably be done with a try_fold

  • keep in the accumulator an optional for each of the searches
  • abort when all elements are Some

like this

    fn find_map_two<P1, P2, T1, T2>(mut self, map1: P1, map2: P2) -> (Option<T1>, Option<T2>)
    where
        Self: Sized,
        Self::Item: Clone,
        P1: FnMut(&Self::Item) -> Option<T1>,
        P2: FnMut(&Self::Item) -> Option<T2>,
    {
        let mut maps = (map1, map2);
        let init = (None, None);
        let rv = self.try_fold(init, |mut old, element| {
            use seq_macro::seq;
            seq!(i in 0..2 {
                old.i = old.i.or_else(||maps.i(&element));
            });
            let mut abort = true;
            seq!(i in 0..2 { abort = abort && old.i.is_some();});
            if abort {
                std::ops::ControlFlow::Break(old)
            } else {
                std::ops::ControlFlow::Continue(old)
            }
        });
        match rv {
            std::ops::ControlFlow::Break(rv) | std::ops::ControlFlow::Continue(rv) => rv,
        }
        //        .into_value();
    }

some odd things to remark

  • i mutate old.i = old.i.or_else because of what comes next
  • same for let mut abort = true; seq! … abort = abort && …, that's just old.0.is_some() && old.1.is_some()

side quest: anything to remark at this point? hints about how to do impl … MultiFindIterator correctly?

the real question

how do i generalize this over multiple FnMut?

  • I believe internally the FnMut and the Option should stay in a tuple and not a fixed size array because the functions might have different signatures (cf my poor little test).
  • afaik rust macros can't loop, and i might have to instantiate 1…10 tuples and functions.
  • i accept that for larger numbers of functions, something with boxing dyns might be necessary, that's out of scope

I switched to a macro for most of the code

macro_rules! find_map_many_impl {
    ($size:literal, $iter:expr, $init:expr, $maps:expr) => {
        {
            let rv = $iter.try_fold($init, |mut old, element| {
                use seq_macro::seq;
                seq!(i in 0..$size {
                    old.i = old.i.or_else(||$maps.i(&element));
                });
                let mut abort = true;
                seq!(i in 0..$size { abort = abort && old.i.is_some();});
                if abort {
                    std::ops::ControlFlow::Break(old)
                } else {
                    std::ops::ControlFlow::Continue(old)
                }
            });
            match rv {
                std::ops::ControlFlow::Break(rv) | std::ops::ControlFlow::Continue(rv) => rv,
            }
            //        .into_value();
        }
    };
}

And then the non-generic part that remains is


    fn find_map_3<P1, P2, P3, T1, T2, T3>(
        mut self,
        map1: P1,
        map2: P2,
        map3: P3,
    ) -> (Option<T1>, Option<T2>, Option<T3>)
    where
        Self: Sized,
        Self::Item: Clone,
        P1: FnMut(&Self::Item) -> Option<T1>,
        P2: FnMut(&Self::Item) -> Option<T2>,
        P3: FnMut(&Self::Item) -> Option<T3>,
    {
        let mut maps = (map1, map2, map3);
        let init = (None, None, None);

        find_map_many_impl!(3, self, init, maps)
    }
}

Can I do better than that? All the Ti and Pi look like they could do with some deduplication.

2

EDIT:

as pointed out by @chrefr, I mis-interpreted the problem in OP, this solution is wrong. please ignore it when discussing the original problem.

/EDIT

the code should be easier to understand if you implement it using find_map() than the current try_fold(), if you figure out how to compose multiple FnMut() predicates:

// can use tuples to emulate variadic function
fn at_least_one_of<
	T: Clone,
	U1,
	U2,
	F1: FnMut1<T, Output = Option<U1>>,
	F2: FnMut1<T, Output = Option<U2>>,
>(
	mut f1: F1,
	mut f2: F2,
) -> impl FnMut1<T, Output = Option<(Option<U1>, Option<U2>)>> {
	move |arg: T| {
		let u1 = f1.call_mut(arg.clone());
		let u2 = f2.call_mut(arg);
		if u1.is_some() || u2.is_some() {
			Some((u1, u2))
		} else {
			None
		}
	}
}
// again, can use tuple
fn find_map_multiple<
	T: Clone,
	I: Iterator<Item = T>,
	U1,
	U2,
	F1: FnMut1<T, Output = Option<U1>>,
	F2: FnMut1<T, Output = Option<U2>>,
>(
	mut i: I,
	f1: F1,
	f2: F2,
) -> (Option<U1>, Option<U2>) {
	let mut f12 = at_least_one_of(f1, f2);
	i.find_map(|x: T| f12.call_mut(x)).unwrap_or_default()
}
3

This does not work. It'll stop when either matches. OP needs to stop only when all match.

Also, in my opinion it's too general (and that can also cause problems with inference). I'll go for:

fn find_map_multiple<I: Iterator, U1, U2>(
    mut iter: I,
    mut f1: impl FnMut(&mut I::Item) -> Option<U1>,
    mut f2: impl FnMut(I::Item) -> Option<U2>,
) -> (Option<U1>, Option<U2>) {
    let (mut u1, mut u2) = (None, None);
    // `any()` short-circuits, which we need here.
    iter.any(|mut item| {
        if u1.is_none() {
            u1 = f1(&mut item);
        }
        if u2.is_none() {
            u2 = f2(item);
        }
        u1.is_some() && u2.is_some()
    });
    (u1, u2)
}
4

sorry, I misread the code.

I saw OP uses the return type (Option<T1>, Option<T2>), as opposed to Option<(T1, T2)>, and I made a wrong assumption without reading the code carefully.

you are right it can cause inference problem, especially for closures.

to clarify a bit, you may have noticed in the snippet I posted, the custom FnMut1<> trait is unnecessary, the combinator works just fine with the standard FnMut(). it is because the code is stripped down from a more complicated version, which used tuples and macros, but I feel it was distracting, so I scrubbed it.

in the original code, there's a single api accepting argument of tuples of functions, e.g . the callsite looks something like this:

let xs = vec![1, 2, 3, 4, 5];
// original intended usage:
let (u1, u2) = find_map_multiple(xs.iter(), (f1, f2));
let (u1, u2, u3) = find_map_mutiple(xs.iter(), (f1, f2, f3));
// instead of
//find_map2(xs.iter(), f1, f2);
//find_map3(xs.iter(), f1, f2, f3);

in your solution, the first callback takes a reference, which eliminates the Clone bound, which is good. on the other hand, this optimization might not always be desirable, since it makes the multiple predicates asymmetric.