Struggling with lifetimes & returning iterator from closure

I have a pattern in my code that occurs often enough that I'd like to create a reusable abstraction for it. Basically my GameState struct can contain a random number generator:

struct GameState {
  pub rng: Option<StdRng>
}

and functions for building iterators of various kinds:

fn legal_actions::evaluate<'a>(game: &'a GameState) -> Box<dyn Iterator<Item = UserAction> + 'a>;

If this generator is present, we use it for random operations, otherwise we use the standard thread_rng(). The way this usually works is as follows:

fn random_action(game: &mut GameState) -> Option<UserAction> {
    if game.rng.is_some() {
        let actions = legal_actions::evaluate(&game).collect::<Vec<_>>();
        actions.into_iter().choose(game.rng.as_mut()?)
    } else {
        legal_actions::evaluate(&game).choose(&mut rand::thread_rng())
    }
}

In the default case we don't need to .collect() anything because there's no conflict between mutable and immutable borrowing, which is a nice performance optimization.

I'm trying to abstract this out in a reusable pattern, but I've had a really hard time getting the lifetimes right. I think the obvious starting point would be something like:

fn choose_randomly<It>(
    game: &mut GameState,
    generator: impl FnOnce(&GameState) -> It,
) -> Option<It::Item>
where
    It: Iterator,
{
    if game.rng.is_some() {
        let actions = generator(game).collect::<Vec<_>>();
        actions.into_iter().choose(game.rng.as_mut()?)
    } else {
        generator(game).choose(&mut rand::thread_rng())
    }
}

but I haven't been able to get something like this to work because of various lifetime issues:

error: lifetime may not live long enough
choose_randomly(&mut game, |g| legal_actions::evaluate(&g))
   | -- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ returning this value requires that `'1` must outlive `'2`
   | return type of closure is Box<(dyn Iterator<Item = UserAction> + '2)>
   | has type `&'1 GameState`

Is there a way to create a generalized version of this function?

I'm not sure exactly what your lifetime problem is because I can't reproduce it with the amount of code you're given, and I suspect it depends on the signature of legal_actions::evaluate() in particular — but I think you can simplify this problem by, instead of accepting a closure, returning a Rng that can be used together with just one call to choose():

use rand::{prelude::IteratorRandom, rngs::StdRng, Rng, RngCore};

struct GameState {
    pub rng: Option<StdRng>,
}

fn game_rng(game: &mut GameState) -> impl Rng + '_ {
    if let Some(rng) = &mut game.rng {
        Box::new(rng) as Box<dyn RngCore + '_>
    } else {
        Box::new(rand::thread_rng())
    }
}

fn evaluate(_game: &GameState) -> impl Iterator<Item = i32> {
    [1, 2, 3].into_iter()
}

fn main() {
    let mut game = GameState { rng: None };
    evaluate(&game).choose(&mut game_rng(&mut game));
}

This does allocate a Box for the rng, but if this pattern works out for your needs, you can then optimize it by writing an enum of the two possible RNGs, then implementing RngCore for that enum.

The trouble I have is that my iterator functions are usually tied to the lifetime of the GameState itself. In your example that would be equivalent to

fn evaluate<'a>(game: &'a GameState) -> impl Iterator<Item = i32> + 'a {
    [1, 2, 3].into_iter()
}

which gives an error.

In the error output: if the closure takes a reference, why are you passing &g instead of g to evaluate?

The problem is that here:

fn choose_randomly<It>(
    game: &mut GameState,
    generator: impl FnOnce(&GameState) -> It,
) -> Option<It::Item>
where
    It: Iterator,

It must be a single type, and generator must produce that type for any input lifetime on &'lifetime GameState. If It captures the lifetime, it's a different type for every lifetime, and thus cannot meet this bound.

This is a case where the sugar of Fn traits becomes salt, because you can't not name the output type when using them. You're going to have to introduce your own redirection to be able to have the result vary in lifetime but still be arbitrarily short.

Here's one way. I rushed it, so maybe it could be a little cleaner. I also put a 'static bound on the eventual item, as otherwise you can't borrow &mut game for both the iterator and the RNG at the same time.

The hrtb is because Rust is bad at inferring higher-lifetime bounds of closures when you need them, so it needs a little assistance in this case.


Edit: Cleaned up a bit.

Other details I rushed over:

  • Type parameters like It must monomorphize to a single type.
  • &'lifetime Thing is a single type, but there is no single &Thing type that covers references of any lifetimes (i.e. there is no higher-ranked reference type). When you see &Thing in type position, it's either part of a higher-ranked type that operates over references of different lifetimes, or it's a &'lifetime Thing where 'lifetime is singular and inferred.
  • impl FnOnce(&GameState) -> It is short for impl for<'any> FnOnce(&'any GameState) -> It; this one is a higher-ranked type as a whole. But It here must be a singular type (i.e. the same for all lifetimes of 'any).
  • If you know the output type except for some lifetime parameter and single-type generics -- like if the output is always a Box<dyn Iterator<Item = Item> + '_> -- then you can use lifetime elision on stable.
  • Similarly, if you could use -> (impl Iterator<Item = Item> + '_), that would also work, but you can't yet
4 Likes

thanks! that does work :slight_smile: