Search code examples
dictionaryrustiteratormapping

Pass value out of iterator without stopping iteration


I have a function that parses an input string:

fn parse_input(s: &str) -> ((usize, usize), BTreeMap<(usize, usize), Tile>){
 let mut start: (usize, usize) = (0,0);
 let grid = s.split("\n").enumerate().flat_map(|(r,l)| {
     l.chars().enumerate().map(move |(col, c)| {
         let t = classify_tile(c);
         match t {
             Tile::Start => {
                  *start = (r, col);
                  ((r,col), t)
                 },
             _ => ((r,col), t) 
         }
     })
 }).collect::<BTreeMap<(usize, usize), Tile>>();
 (start, grid)
}

I basically want to capture the r and col value of the start tile (which is unique, only one occurrence). But currently if I try and modify the tuple inside the iterator I assume due to borrowing and scope reasons, the value doesn't get modified outside the iterator. It is important the iterator finishes though.

An alternative solution would be to search for the start tile in the btreemap afterwards but I'm hoping for a more efficient solution.

Should I just do this as nested for loops? Is the iteration actually meaningfully more efficient here?

edit: classify_tile function returns the enum type. Start, Soil or Pipe. The *start = (r,col) section is the bit that doesn't work. That was my attempt at solving it. Everything else works though.


Solution

  • You can do what you want, if you restructure your iterator a bit.


    But let's first address what the issue is.

    As a simpler example, let's count all uppercase letters (in a silly way):

    let text = "Hello World";
    
    let mut count = 0;
    
    text.split_whitespace()
       .flat_map(|s| {
            s.chars().map(|c| {
                if c.is_uppercase() {
                    count += 1;
                }
            })
        })
        .for_each(|_| {});
    
    assert_eq!(count, 2);
    

    If we try to compile this, we'll get the same issue you're getting:

    error: captured variable cannot escape `FnMut` closure body
      --> src\main.rs:85:13
       |
    80 |       let mut count = 0;
       |           --------- variable defined here
    ...
    84 |           .flat_map(|s| {
       |                       - inferred to be a `FnMut` closure
    85 | /             s.chars().map(|c| {
    86 | |                 if c.is_uppercase() {
    87 | |                     count += 1;
       | |                     ----- variable captured here
    88 | |                 }
    89 | |             })
       | |______________^ returns a closure that contains a reference to a captured variable, which then escapes the closure body
       |
       = note: `FnMut` closures only have access to their captured variables while they are executing...
       = note: ...therefore, they cannot allow references to captured variables to escape
    

    Why is that?

    The closure given to map() requires mutably borrowing count, this in itself is perfectly fine.

    The issue comes when we throw flat_map() into the mix. Since flat_map() potentially produces multiple iterators. Each of these produced iterators requires mutably borrowing count. This is of course not allowed, since we cannot mutably borrow anything more than once.


    Overall, this is straight forward to resolve, you simply restructure the iterator:

    text.split_whitespace()
        .flat_map(|s| s.chars())
        .map(|c| {
            if c.is_uppercase() {
                count += 1;
            }
        })
        .for_each(|_| {});
    

    or

    text.split_whitespace()
        .map(|s| s.chars())
        .flatten()
        .map(|c| {
            if c.is_uppercase() {
                count += 1;
            }
        })
        .for_each(|_| {});
    

    Resolving this for your iterator, could look something like this:

    let grid = s
        .split("\n")
        .enumerate()
        .flat_map(|(r, l)| {
            l.chars().enumerate().map(move |(col, c)| {
                let t = classify_tile(c);
                match t {
                    Tile::Start => ((r, col), t),
                    _ => ((r, col), t),
                }
            })
        })
        .inspect(|((r, col), t)| match t {
            Tile::Start => {
                start = (*r, *col);
            }
            _ => {}
        })
        .collect::<BTreeMap<(usize, usize), Tile>>();
    

    or

    let grid = s
        .split("\n")
        .enumerate()
        .flat_map(|(r, l)| l.chars().enumerate().map(move |(col, c)| ((r, col), c)))
        .map(|((r, col), c)| {
            let t = classify_tile(c);
            match t {
                Tile::Start => {
                    start = (r, col);
                    ((r, col), t)
                }
                _ => ((r, col), t),
            }
        })
        .collect::<BTreeMap<(usize, usize), Tile>>();
    

    It can be written in multiple ways.