Rust error "cannot infer an appropriate lifetime for borrow expression" when attempting to mutate state inside a closure returning an Iterator

Question

I'm trying to learn Rust and have encountered a lifetime-related problem while trying to emulate nested Python generators. The problem is with the lifetime of a value mutated by a closure, as reported by the compiler. The crux of the code is flat_mapping a closure that calls a function that mutates a value supplied from the outer-scope in its returned Iterator. See line 39 in the Rust playground example.

The code here is a simplified, trivial version of the original program. Since my ultimate goal is to learn more about Rust, I'd appreciate some insight even more than a fix for my code!

For example, one "solution" is the commented-out code on line 44. It "works" but it misses the point by always allocating a Vec that contains all points on the trace even if the user only wants to check the first Point on a trace.

I think the problem has something to do with how the mutable borrow to point lives on in the Iterator that trace_steps returns. I have tried far too many variations to list here, from passing in the point that is mutated from main (more similar to how trace_step works) to attempts at blindly using Rc<RefCell<Point>> when I started to get desperate.

Below is the code copied from the Rust Playground is:

#[derive(Debug, Eq, PartialEq)]
struct Point {
    x: i32,
    y: i32,
}

impl Point {
    fn new(x: i32, y: i32) -> Point {
        Point { x, y }
    }
}

// Intention is that this is like a Python generator.  Normally the "step" would
// be a struct with a direction and a length but this is a simplified version.
fn trace_step<'a>(point: &'a mut Point, step: u8) -> impl Iterator<Item = Point> + 'a {
    let mut len = step;
    std::iter::from_fn(move || {
        if len == 0 {
            None
        } else {
            len -= 1;
            point.x += 1;
            Some(Point { ..*point })
        }
    })
}

// FIXME: See compiler error!!!
// Compiler cannot infer an appropriate lifetime for the borrow &mut point.
// Can't the borrow just live as long as the closure?
//
// Intention is that this produces points along a path defined by multiple
// steps.  Simplified.
fn trace_steps(steps: Vec<u8>) -> impl Iterator<Item = Point> {
    let mut point: Point = Point::new(0, 0);

    // FIXME: This doesn't work.
    let f = |x: &u8| trace_step(&mut point, *x);
    steps.iter().flat_map(f)

    // This works, but we don't want to commit to allocating the space for all
    // points if the user only needs to, for example, count the number of points.
    /*
    let mut ret: Vec<Point> = Vec::new();
    for step in steps {
        ret.extend(trace_step(&mut point, step));
    }
    ret.into_iter()
    */
}

fn main() {
    let mut point: Point = Point::new(0, 0);
    let points: Vec<Point> = trace_step(&mut point, 3).collect();

    // Outputs: [Point { x: 1, y: 0 }, Point { x: 2, y: 0 }, Point { x: 3, y: 0 }]
    println!("{:?}", points);

    // Should trace the first from (0, 0) to (1, 0) and then trace the second step
    // from (1, 0) to (2, 0) to (3, 0).
    let points: Vec<Point> = trace_steps(vec![1, 2]).collect();
    println!("{:?}", points);
}

And the error when this is run in the Rust Playground is:

   Compiling playground v0.0.1 (/playground)
error[E0495]: cannot infer an appropriate lifetime for borrow expression due to conflicting requirements
  --> src/main.rs:38:33
   |
38 |     let f = |x: &u8| trace_step(&mut point, *x);
   |                                 ^^^^^^^^^^
   |
note: first, the lifetime cannot outlive the lifetime '_ as defined on the body at 38:13...
  --> src/main.rs:38:13
   |
38 |     let f = |x: &u8| trace_step(&mut point, *x);
   |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
note: ...so that closure can access `point`
  --> src/main.rs:38:33
   |
38 |     let f = |x: &u8| trace_step(&mut point, *x);
   |                                 ^^^^^^^^^^
note: but, the lifetime must be valid for the destruction scope surrounding expression at 34:63...
  --> src/main.rs:34:63
   |
34 |   fn trace_steps(steps: Vec<u8>) -> impl Iterator<Item = Point> {
   |  _______________________________________________________________^
35 | |     let mut point: Point = Point::new(0, 0);
36 | |     
37 | |     // FIXME: This doesn't work.
...  |
49 | |     */
50 | | }
   | |_^
note: ...so that references are valid when the destructor runs
  --> src/main.rs:34:63
   |
34 |   fn trace_steps(steps: Vec<u8>) -> impl Iterator<Item = Point> {
   |  _______________________________________________________________^
35 | |     let mut point: Point = Point::new(0, 0);
36 | |     
37 | |     // FIXME: This doesn't work.
...  |
49 | |     */
50 | | }
   | |_^

error: aborting due to previous error

error: could not compile `playground`.

Answer 1

The problem is that Rust is very strict about copying mutable references. This is a problem because when you return the iterator inside flat_map, that iterator has to have a mutable (sole) reference to the point, but flat_map isn't robust enough to give the iterator back to you, and therefore Rust can't prove that the last iterator doesn't still reference the point by the time the closure is called again. Once generators are stabilized, this will be trivial to do properly. In the meantime, it is still possible, but MUCH harder than I expected, you need to manually implement the Iterator trait. Here you go:

Playground link

use std::iter::{ExactSizeIterator, FusedIterator};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Point {
    x: i32,
    y: i32,
}

impl Point {
    fn new(x: i32, y: i32) -> Point {
        Self { x, y }
    }
}

#[derive(Debug)]
struct StepTracer<'a> {
    point: &'a mut Point,
    len: u8,
}

impl<'a> StepTracer<'a> {
    fn new(point: &'a mut Point, len: u8) -> Self {
        Self { point, len }
    }

    fn into_inner(self) -> &'a mut Point {
        self.point
    }
}

impl<'a> Iterator for StepTracer<'a> {
    type Item = Point;

    fn next(&mut self) -> Option<Self::Item> {
        if self.len == 0 {
            None
        } else {
            self.len -= 1;
            self.point.x += 1;
            Some(*self.point)
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.len as usize, Some(self.len as usize))
    }
}

impl FusedIterator for StepTracer<'_> {}
impl ExactSizeIterator for StepTracer<'_> {}

// You may also want to consider implementing DoubleEndedIterator
// Additional traits: https://doc.rust-lang.org/std/iter/index.html#traits

enum MultiStepTracerState<'a> {
    First(&'a mut Point),
    Second(&'a mut Point),
    Tracer(StepTracer<'a>),
    Done,
}

/// Intention is that this produces points along a path defined by multiple
/// steps. Simplified.
struct MultiStepTracer<'a, I: Iterator<Item = u8>> {
    steps: I,
    state: MultiStepTracerState<'a>,
}

impl<'a, I: Iterator<Item = u8>> MultiStepTracer<'a, I> {
    fn new(point: &'a mut Point, steps: I) -> Self {
        Self {
            steps,
            state: MultiStepTracerState::First(point),
        }
    }
}

impl<I: Iterator<Item = u8>> Iterator for MultiStepTracer<'_, I> {
    type Item = Point;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let mut temp_state = MultiStepTracerState::Done;
            std::mem::swap(&mut self.state, &mut temp_state);
            let point_ref = match temp_state {
                MultiStepTracerState::First(point) => {
                    let result = *point;
                    self.state = MultiStepTracerState::Second(point);
                    return Some(result);
                }
                MultiStepTracerState::Second(point) => point,
                MultiStepTracerState::Tracer(mut tracer) => {
                    if let Some(result) = tracer.next() {
                        self.state = MultiStepTracerState::Tracer(tracer);
                        return Some(result);
                    } else {
                        tracer.into_inner()
                    }
                }
                MultiStepTracerState::Done => {
                    return None;
                }
            };

            if let Some(len) = self.steps.next() {
                self.state = MultiStepTracerState::Tracer(StepTracer::new(point_ref, len));
            } else {
                self.state = MultiStepTracerState::Done;
                return None;
            }
        }
    }
}

impl<I: Iterator<Item = u8>> FusedIterator for MultiStepTracer<'_, I> {}

fn main() {
    let mut point: Point = Point::new(0, 0);
    let points: Vec<Point> = StepTracer::new(&mut point, 3).collect();

    // Outputs: [Point { x: 1, y: 0 }, Point { x: 2, y: 0 }, Point { x: 3, y: 0 }]
    println!("{:?}", points);

    // Should trace the first from (0, 0) to (1, 0) and then trace the second step
    // from (1, 0) to (2, 0) to (3, 0).
    let points: Vec<Point> =
        MultiStepTracer::new(&mut Point::new(0, 0), [1, 2].iter().copied()).collect();
    println!("{:?}", points);
}

Answer 2

The original question asked for an Iterator of points along some path defined by run lengths and the answer below does not provide an Iterator. The accepted answer above still deserves all credit for being the best answer to the original question.

The code below achieves essentially the same result by ditching the mutable state and fully embracing the functional approach struggling to break through the via flat_map in the original question's mess of code.

Run on the Rust playground.

Code:

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Point {
    x: i32,
    y: i32,
}

impl Point {
    fn new(x: i32, y: i32) -> Point {
        Self { x, y }
    }
}

fn main() {
    let origin: Point = Point::new(0, 0);
    let lengths: Vec<u16> = vec![1, 2];

    // Function that returns the next point after "taking a step"
    fn step(p: Point) -> Point {
         Point {x: p.x + 1, y: p.y }
    };

    /*****************************************
     * ORIGINAL EXAMPLE: Collect all points along the path
     *****************************************/

    // The crux of this version of the answer is to create all of the steps we 
    // intend to take for each length.  Steps will be an iterator that is 
    // something like: [|x| step(x), |x| step(x), |x| step(x)]
    let steps = lengths.iter().flat_map(|num_steps: &u16| (0..*num_steps).map(|_| |x| step(x)) );

    // `fold` lets us chain steps one after the other.  Unfortunately, this
    // doesn't give us an iterator, so it's not a good answer to the original 
    // question.
    let path_points: Vec<Point> = steps.fold(vec![origin], |mut acc, f| {
        acc.push(f(*acc.last().unwrap()));
        acc
    }).split_off(1);  // split_off gets rid of the initial "origin" point at (0, 0)
    println!("Path for original example: {:?}", path_points);

    /*****************************************
     * BONUS EXAMPLE: Get just the endpoint
     *****************************************/

    // Same as above
    let steps = lengths.iter().flat_map(|num_steps: &u16| (0..*num_steps).map(|_| |x| step(x)) );

    // Note that this has the same space-saving benefits of the iterator 
    // solution, but it requires the user to do more work in general having to
    // think about how to write the folding function
    let end_point: Point = steps.fold(origin, |acc, f| {
        f(acc)
    });
    println!("End point for bonus example: {:?}", end_point);
}

Output:

Path for original example: [Point { x: 1, y: 0 }, Point { x: 2, y: 0 }, Point { x: 3, y: 0 }]
End point for bonus example: Point { x: 3, y: 0 }