Techniques for turning recursive functions into iterators in Rust?

user_318206 0
Viewed 951

I'm struggling to turn a simple recursive function into a simple iterator. The problem is that the recursive function maintains state in its local variables and call stack -- and to turn this into a rust iterator means basically externalizing all the function state into mutable properties on some custom iterator struct. It's quite a messy endeavor.

In a language like javascript or python, yield comes to the rescue. Are there any techniques in Rust to help manage this complexity?

Simple example using yield (pseudocode):

function one_level(state, depth, max_depth) {
  if depth == max_depth {
    return
  }
  for s in next_states_from(state) {
    yield state_to_value(s);
    yield one_level(s, depth+1, max_depth);
  }
}

To make something similar work in Rust, I'm basically creating a Vec<Vec<State>> on my iterator struct, to reflect the data returned by next_states_from at each level of the call stack. Then for each next() invocation, carefully popping pieces off of this to restore state. I feel like I may be missing something.

rust recursion iterator
1 Answers

You are performing a (depth-limited) depth-first search on your state graph. You can do it iteratively by using a single stack of unprocessed subtrees(depending on your state graph structure).

struct Iter {
    stack: Vec<(State, u32)>,
    max_depth: u32,
}

impl Iter {
    fn new(root: State, max_depth: u32) -> Self {
        Self {
            stack: vec![(root, 0)],
            max_depth
        }
    }
}

impl Iterator for Iter {
    type Item = u32; // return type of state_to_value
    fn next(&mut self) -> Option<Self::Item> {
        let (state, depth) = self.stack.pop()?;
        if depth < self.max_depth {
            for s in next_states_from(state) {
                self.stack.push((s, depth+1));
            }
        }
        return Some(state_to_value(state));
    }
}

There are some slight differences to your code:

  • The iterator yields the value of the root element, while your version does not. This can be easily fixed using .skip(1)
  • Children are processed in right-to-left order (reversed from the result of next_states_from). Otherwise, you will need to reverse the order of pushing the next states (depending on the result type of next_states_from you can just use .rev(), otherwise you will need a temporary)
user_138068680
Related