Using `flat_map` to create iterator over fields for slice of struct

Question

Given a collection (vector/slice) of structs. How do I create a combined iterator over some fields in each struct?

Below is a concrete attempt using flat_map:

struct Game {
    home_team: u8,
    away_team: u8,
}

fn teams(games: &[Game]) -> impl Iterator<Item = u8> {
    games
        .iter()
        .flat_map(|game| [game.home_team, game.away_team].iter().map(|x| x.clone()))
}
fn main() {
    let data = &[
        Game {
            home_team: 1,
            away_team: 2,
        },
        Game {
            home_team: 1,
            away_team: 3,
        },
    ];
    let non_unique_teams: Vec<u8> = teams(data).collect();
}

My actual use-case is very similar. In particular, the fields that form the basis of the iterator implements Copy, making cloning perfectly fine. My intuition tells me that this should work since I'm cloning the only things I need to "take" from the incoming slice. Obviously, my understanding of the borrow checker is to poor for me to grasp this.

Sven Marnach · Accepted Answer · 2020-09-22T10:25:07.720

The iterator needs to own the memory containing the copies of the struct fields. In your code, you create a local array and call iter() on it, which results in an iterator over a slice reference that does not own the data.

The easiest way to make the iterator own the data is to allocate a vector for each struct:

fn teams(games: &[Game]) -> impl Iterator<Item = u8> + '_ {
    games
        .iter()
        .flat_map(|game| vec![game.home_team, game.away_team])
}

This will result in a heap allocation in each iteration. The performance penalty is probably minimal, since the allocator is probably able to reuse the allocation in each iteration. However, if you want to avoid the allocation for some reason, you could also use a combination of Iterator::chain() and std::iter::once():

use std::iter::once;

fn teams(games: &[Game]) -> impl Iterator<Item = u8> + '_ {
    games
        .iter()
        .flat_map(|game| once(game.home_team).chain(once(game.away_team)))
}

Further alternatives include implementing IntoIterator and Clone for Game, which would allow you to simply use games.iter().cloned().flatten(), the iter_vals crate or using generators, an unstable feature that makes implementing this kind of iterator more convenient.

Using `flat_map` to create iterator over fields for slice of struct

1 Answers1