Typing a closure that returns an anonymous type borrowing from one of its inputs without heap allocation or trait objects

Question

Let's say that I have the following working code:

use std::collections::VecDeque;

fn modify<S, F>(state: &mut S, func: F)
where
    F: for<'a> Fn(&'a mut S) -> Box<dyn Iterator<Item = &mut u64> + 'a>
{
    let mut prev = 1;
    for _ in 0..3 {
        for item in func(state) {
            let val = *item;
            *item = val * prev;
            prev = val;
        }
    }
}

fn main() {
    let mut state: VecDeque<u64> = vec![1,2,3,4].into();

    modify(&mut state, |s| Box::new(s.iter_mut()));
    assert_eq!(state, [48, 8, 24, 864]);

    modify(&mut state, |s| Box::new(s.iter_mut().take(2)));
    assert_eq!(state, [147456, 7077888, 24, 864]);

    modify(&mut state, |s| Box::new(s.iter_mut().skip(2)));
    assert_eq!(state, [147456, 7077888, 429981696, 10319560704]);
}

modify takes a function that generates an iterator multiple times from a given state variable (in this case, a VecDeque). This function works, but contains an unnecessary heap allocation as well as unneeded polymorphism in the form of a trait object. I'd very much like to do something like this:

use std::collections::VecDeque;

fn modify<'a, S, F, I>(state: &'a mut S, func: F)
where
    F: Fn(&'a mut S) -> I,
    I: Iterator<Item = &'a mut u64>
{
    let mut prev = 1;
    for _ in 0..3 {
        for item in func(state) {
            let val = *item;
            *item = val * prev;
            prev = val;
        }
    }
}

fn main() {
    let mut state: VecDeque<u64> = vec![1,2,3,4].into();

    modify(&mut state, |s| s.iter_mut());
    assert_eq!(state, [48, 8, 24, 864]);

    modify(&mut state, |s| s.iter_mut().take(2));
    assert_eq!(state, [147456, 7077888, 24, 864]);

    modify(&mut state, |s| s.iter_mut().skip(2));
    assert_eq!(state, [147456, 7077888, 429981696, 10319560704]);
}

But I now get this error:

error[E0499]: cannot borrow `*state` as mutable more than once at a time
  --> src/main.rs:10:26
   |
3  | fn modify<'a, S, F, I>(state: &'a mut S, func: F)
   |           -- lifetime `'a` defined here
...
10 |         for item in func(state) {
   |                     -----^^^^^-
   |                     |    |
   |                     |    `*state` was mutably borrowed here in the previous iteration of the loop
   |                     argument requires that `*state` is borrowed for `'a`

For more information about this error, try `rustc --explain E0499`.

And I can't figure out how to wrangle the borrow checker to do what I want. Any clues would be appreciated. Thanks!

Answer 1

TL;DR

No, not until closure HRTB inference is fixed. Current workarounds include using function pointers instead or implementing a helper trait on custom structs -- the helper trait is needed regardless of approach until higher-kinded types are introduced in Rust.

Playground

Details

To avoid returning a Box, you would need the type parameter I to be generic over the lifetime 'a, so that you can use it with any lifetime (in a for<'a> bound, for example). Unfortunately, as discussed in a similar question, Rust does not yet support higher-kinded types (type parameters that are themselves generic over other type parameters), so we must use a helper trait:

trait GenerateIterator<'a, S: 'a> {
    type Output: Iterator<Item = &'a mut u64>;
    fn generate(self, s: &'a mut S) -> Self::Output;
}

impl<'a, S: 'a, I, F> GenerateIterator<'a, S> for F
where
    F: Fn(&'a mut S) -> I,
    I: Iterator<Item = &'a mut u64>,
{
    type Output = I;
    fn generate(self, s: &'a mut S) -> Self::Output {
        self(s)
    }
}

fn modify<S>(state: &mut S, func: impl for<'a> GenerateIterator<'a, S>) {
    let mut prev = 1;
    for _ in 0..3 {
        for item in func.generate(state) {
            let val = *item;
            *item = val * prev;
            prev = val;
        }
    }
}

This is great in helping us work around the use of higher-kinded types in this situation, but due to another compiler limitation in closure HRTB (higher-ranked trait bound) inference, the compiler is unable to infer that the closures we pass to modify are general over any lifetime. The alternative is to use function pointers like VecDeque::iter_mut, or create your own structs that implement GenerateIterator to generate the Iterator as you want:

fn main() {
    fn take2(s: &mut VecDeque<u64>) -> impl Iterator<Item = &mut u64> {
        s.iter_mut().take(2)
    }

    fn skip2(s: &mut VecDeque<u64>) -> impl Iterator<Item = &mut u64> {
        s.iter_mut().skip(2)
    }

    let mut state: VecDeque<u64> = vec![1, 2, 3, 4].into();

    modify(&mut state, VecDeque::iter_mut);
    assert_eq!(state, [48, 8, 24, 864]);

    modify(&mut state, take2);
    assert_eq!(state, [147456, 7077888, 24, 864]);

    modify(&mut state, skip2);
    assert_eq!(state, [147456, 7077888, 429981696, 10319560704]);
}