Lazily mapping over iterator with `FnMut`

Hello,

I am trying to lazily map over an iterator with a closure that uses a &mut reference.
This is not permitted by the compiler with the error message "captured variable cannot escape FnMut closure body".
I have created an example at the Playground that demonstrates roughly what I would like to do. My non-working attempt is the (commented) function countdown_repeat_mut in lines 26--34.
I know that the functions in the playground example look quite artificial and that you would never write such a countdown_mut in real life, but it is a reasonable simplification of the actual problem I'm trying to solve.

Is this actually possible what I am aiming for? If so, how?
Please keep in mind that I am aiming to create a lazy iterator, so I am not allowed to call collect() when producing the iterator in countdown_repeat_mut.

Thank you!

Would something like this work (playground)?

/// Produce a countdown from `n-1` to `0`.
fn countdown(n: usize) -> Box<dyn Iterator<Item = usize>> {
    if n == 0 {
        Box::new(empty())
    } else {
        Box::new(once(n - 1).chain(countdown(n - 1)))
    }
}

/// Produce `n` countdowns from `n-1` to `0`.
fn countdown_repeat(n: usize) -> Box<dyn Iterator<Item = usize>> {
    Box::new(repeat(n).take(n).flat_map(countdown))
}

Or if the &mut arguments are important, you can make them work like this (playground):

fn countdown_mut(n: &mut usize) -> Box<dyn Iterator<Item = usize>> {
    if *n == 0 {
        Box::new(empty())
    } else {
        *n -= 1;
        Box::new(once(*n).chain(countdown_mut(n)))
    }
}

fn countdown_repeat_mut(i: &mut usize) -> Box<dyn Iterator<Item = usize> + '_> {
    Box::new(repeat(*i).take(*i).flat_map(move |m| {
        *i = m;
        countdown_mut(i)
    }))
}

The important thing to note is that the return value of countdown_mut does not actually borrow its argument, so you can make your original code compile just by removing the lifetime from the return type. It's important that it doesn't borrow its argument, because this allows the Map iterator to call the closure multiple times without worrying about overlapping borrows.

If this doesn't work in your real use-case, then you might not be able to compose your iterator out of generic adapters like Map, and might instead need to create a type with a custom Iterator impl. When generators become stable, they will provide an easier way to do this in many cases.

Thank you so much for your reply!

You are right, I can omit the lifetime from the return type in countdown_mut.
Unfortunately, I am not sure if I can do something like this in my actual use case ... which actually looks a bit more like this (Playground), but does not compile due to lifetime errors:

fn multi<'i, 's: 'i, C>(
    mut cl: impl Iterator<Item = C> + Clone + 'i,
    st: &'s mut C,
) -> Box<dyn Iterator<Item = usize> + 'i> {
    match cl.next() {
        // `once(0)` is just a placeholder for a more complicated function here
        Some(lit) => Box::new(core::iter::once(0).flat_map(move |_| multi(cl, st))),
        None => Box::new(core::iter::once(0)),
    }
}

The important thing is that I would like to call the multi function from within the multi function as part of a flat_map (thanks, by the way, for pointing this function out to me!), passing the mutable argument st. I have already tried quite a few lifetime combinations, but did not get things to work.

If I understand correctly what you write, it seems to me that the way I want to write this code will not work because I think that I borrow st, but I'm actually not sure about that ...
Can you help me confirm or disprove my hypothesis? :slight_smile:

Yeah, I think it's likely true that there's no simple way to write your algorithm in this way (where "this way" means something like "recursively calling a function from within a closure that captures the function's arguments and gets returned from the function").

It's a tricky case, and I'm not sure I could even explain exactly where the problems are. But my first step toward solving it would probably be to express it in an iterative loop (without recursion), and then write a lazy iterator based on that version, maybe by manually implementing Iterator for a custom type.

I see. Thank you for taking the time to discuss my problem with me. :slight_smile: It has been very helpful!

The problem is that if the closure returns a mutable reference to something (or something containing a mutable reference), then calling it multiple times like this would be invalid:

let val1 = my_closure();
let val2 = my_closure();
use_variable(val1); // oops, val2 should have exclusive access

This is invalid because mutable references have exclusive access, so the above code should not be valid as it violates the exclusivity of mutable references.

The problem is then that the FnMut trait defines the following signature:

// with irrelevant details removed
trait FnMut<Args>: FnOnce<Args> {
    type Output;
    fn call_mut(&mut self, args: Args) -> Self::Output;
}

With this signature there is nothing to tie the (elided) lifetime on &mut self together with the type Self::Output. Thus the snippet from before must be valid, as calling the closure does not mutably borrow it until the returned value is no longer valid.

3 Likes