Resumable continuation-passing style iterator reduce in Rust

user_915718 0
Viewed 888

I'm trying to write a continuation-passing-style "reduce" function that can be resumed at any point. I've got a version working, but in this I need to explicitly write a new version of the function if I want it to be able to make use of a borrow of some state. Rust Playground Link

fn reduce_async_with_store<'a, I, A, F, C>(
    store: &mut Store,
    mut iterator: I,
    accumulator: A,
    mut f: F,
    continuation: C,
) where
    I: Iterator + 'a,
    F: FnMut(&mut Store, I::Item, A, Box<dyn FnOnce(&mut Store, A) + 'a>) + Clone + 'a,
    C: FnOnce(&mut Store, A) + 'a,
{
    match iterator.next() {
        None => continuation(store, accumulator),
        Some(item) => {
            let next: Box<dyn FnOnce(&mut Store, A) + 'a> = {
                let f = f.clone();
                Box::new(move |store, accumulator| {
                    reduce_async_with_store(store, iterator, accumulator, f, continuation)
                })
            };
            f(store, item, accumulator, next);
        }
    }
}

fn some_operation(state: &mut Store, continuation: Box<dyn FnOnce(&mut Store) + 'static>) {
    let mut new_state = Store { foo: state.foo };
    continuation(&mut new_state);
}

#[derive(Debug)]
pub struct Store {
    foo: u8,
}

fn main() {
    let mut some_state = Store { foo: 0 };
    let arr = vec![1u8, 2u8, 3u8];
    reduce_async_with_store(
        &mut some_state,
        arr.into_iter(),
        Vec::new(),
        |store, item, mut acc, continuation| {
            println!("Item: {}", item);
            store.foo += item;
            acc.push(item);
            some_operation(
                store,
                Box::new(move |stor| {
                    continuation(stor, acc);
                }),
            );
        },
        |store, acc| {
            println!("Done!! {:?} {:?}", store, acc);
        },
    )
}

Here's the version of this function I'd like to write, where I can pass the Store in as part of the accumulator, and get it out - however, if I do this, I get cannot infer an appropriate lifetime due to conflicting requirements.

Rust Playground Link

fn reduce_async<'a, I, A, F, C>(mut iterator: I, accumulator: A, mut f: F, continuation: C)
where
    I: Iterator + 'a,
    F: FnMut(I::Item, A, Box<dyn FnOnce(A) + 'a>) + Clone + 'a,
    C: FnOnce(A) + 'a,
{
    match iterator.next() {
        None => continuation(accumulator),
        Some(item) => {
            let next: Box<dyn FnOnce(A) + 'a> = {
                let f = f.clone();
                Box::new(move |accumulator| reduce_async(iterator, accumulator, f, continuation))
            };
            f(item, accumulator, next);
        }
    }
}

fn some_operation(state: &mut Store, continuation: Box<dyn FnOnce(&mut Store) + 'static>) {
    let mut new_state = Store { foo: state.foo };
    continuation(&mut new_state);
}

#[derive(Debug)]
pub struct Store {
    foo: u8,
}

fn main() {
    let mut some_state = Store { foo: 0 };
    let arr = vec![1u8, 2u8, 3u8];
    reduce_async(
        arr.into_iter(),
        (&mut some_state, Vec::new()),
        |item, mut acc, continuation| {
            let (store, vec) = acc;
            println!("Item: {}", item);
            store.foo += item;
            vec.push(item);
            some_operation(
                store,
                Box::new(move |store| {
                    continuation((store, vec));
                }),
            );
        },
        |(store, vec)| {
            println!("Done!! {:?} {:?}", store, vec);
        },
    )
}

How can I write this non-specialized version of my function, and pass things like &mut Store through while respecting Rust's lifetimes?

How is it that my first example with reduce_async_with_store is permitted, even though I don't specify an explicit lifetime for &mut Store, and it could live until 'static?

some_operation takes a boxed closure because that's what the 3rd party API function I'm calling takes. I would like to eventually replace this code with async iterators, but the library I'm using doesn't have support for futures yet.

rust continuations lifetime
1 Answers

Let's start with some_operation; it's always easier to examine regular functions than closures because the compiler only checks their signatures.

Putting back the elided lifetimes, it looks like:

fn some_operation<'s>(state: &'s mut Store, continuation: Box<dyn for<'r> FnOnce(&'r mut Store) + 'static>) {
    let mut new_state = Store { foo: state.foo };
    continuation(&mut new_state);
}

There are two distinct lifetimes involved: 's and 'r — there is no connection between them.

Now let's look here:

Box::new(move |store| {
    continuation((store, vec));
}),

The continuation type should be Box<dyn FnOnce(A) + 'a> according to reduce_async's signature. What is the type of A after monomorphising? The argument passed to the function is a tuple:

(&mut some_state, Vec::new()),

The first element has type &'state mut State for some 'state and the second has Vec<u8>. Revisiting some_operation's signature: the first argument is &'s mut State, so we have chosen 'state = 's here. Then we call the closure using an argument with type &'r mut State.

Back in the main procedure, we are trying to build the accumulator from a value of type (&'r mut State, Vec<u8>) which is not the same as (&'state mut State, Vec<u8>).

That's what the compiler is trying to explain :) Let's check this explanation by changing some_operation's signature:

fn some_operation<'s>(state: &'s mut Store, continuation: Box<dyn FnOnce(&'s mut Store) + 's>) {
    continuation(state);
}

Here we explicitly mark that both lifetimes should be the same, and now the code compiles without any error.

Note that there were no problems in your first code snippet because the lifetime of store: &mut Store argument is different each time you invoke reduce_async_with_store! In the second snippet it is fixed to 'state.

In my opinion, the easiest fix would be to get rid of mutable references at all and passing Store by transferring ownership.

Rust playground link

fn reduce_async<'a, I, A, F, C>(mut iterator: I, accumulator: A, mut f: F, continuation: C)
where
    I: Iterator + 'a,
    F: FnMut(I::Item, A, Box<dyn FnOnce(A) + 'a>) + Clone + 'a,
    C: FnOnce(A) + 'a,
{
    match iterator.next() {
        None => continuation(accumulator),
        Some(item) => {
            let next: Box<dyn FnOnce(A) + 'a> = {
                let f = f.clone();
                Box::new(move |accumulator| reduce_async(iterator, accumulator, f, continuation))
            };
            f(item, accumulator, next);
        }
    }
}

fn some_operation(state: Store, continuation: Box<dyn FnOnce(Store) + 'static>) {
    let new_state = Store { foo: state.foo };
    continuation(new_state);
}

#[derive(Debug)]
pub struct Store {
    foo: u8,
}

fn main() {
    let some_state = Store { foo: 0 };
    let arr = vec![1u8, 2u8, 3u8];
    reduce_async(
        arr.into_iter(),
        (some_state, Vec::new()),
        |item, acc, continuation| {
            let (mut store, mut vec) = acc;
            println!("Item: {}", item);
            store.foo += item;
            vec.push(item);
            some_operation(
                store,
                Box::new(move |store| {
                    continuation((store, vec));
                }),
            );
        },
        |(store, vec)| {
            println!("Done!! {:?} {:?}", store, vec);
        },
    )
}

Keep in mind that the continuation invocations are not tail-recursive, so stack will grow on each iteration. You will probably need a trampoline here.

user_43473410
Related