Iterator stacks

Let's say I have a vector of enums. One of the enum types can be a subvector, and it can contain another subvector, etc.

Now I want to perform a recursive iteration over the entire thing.

I have a working implementation that uses a stack of "iterator nodes", each of which contains a reference to the vector and an index indicating how far along the iteration it is. Each time it encounters a new subvector is just pushes the node on to the stack and reiterates.

Is there a way to do this with a stack of actual iterators?

Can you use flat_map?

If I recall correctly flat_map() is for two level structures only, but I have N layers.

Well that's called a tree then. You'll probably need to implement one of the many tree traversal algorithms, depending on what order you want your nodes to be visited in.

A simple approach would be to use naïve recursion, like this (playground):

enum Tree {
    Leaf(u32),
    Node(Vec<Tree>),
}

fn traverse_naive<F: FnMut(u32)>(tree: &Tree, callback: &mut F) {
    match *tree {
        Tree::Leaf(x) => callback(x),
        Tree::Node(ref subtrees) => {
            subtrees.iter().for_each(|st| traverse_naive(st, callback))   
        }
    }
}
struct S;

enum E {
    A(Vec<E>),
    B(Vec<E>),
    C(S),
}

impl E {
    fn iter<'a>(&'a self) -> Box<dyn Iterator<Item=&'a S> + 'a> {
        match *self {
            E::A(ref v) => Box::new(v.iter().flat_map(|e| e.iter())),
            E::B(ref v) => Box::new(v.iter().flat_map(|e| e.iter())),
            E::C(ref s) => Box::new(std::iter::once(s))
        }
    }
}
2 Likes

Fundamentally, Iterator is just a wrapper around a next() function that adds a bunch of utilities. If you want a solution that’s similar to the one you have, but uses and appears as an iterator, you could do something like this:

(Playground)

pub struct DFS<'a,Node:'a,NodeIter:Iterator<Item=&'a Node>, Visit:Fn(&'a Node)->NodeIter> {
    visit: Visit,
    stack: Vec<(NodeIter, &'a Node)>,
}

impl<'a,Node:'a,NodeIter:Iterator<Item=&'a Node>, Visit:Fn(&'a Node)->NodeIter>
DFS<'a,Node,NodeIter,Visit> {
    pub fn new(root: &'a Node, visit: Visit)->Self {
        DFS {
            stack: vec![ (visit(root), root) ],
            visit
        }
    }
}

impl<'a,Node:'a,NodeIter:Iterator<Item=&'a Node>, Visit:Fn(&'a Node)->NodeIter>
Iterator for DFS<'a,Node,NodeIter,Visit> {
    type Item = &'a Node;
    fn next(&mut self)->Option<&'a Node> {
        if self.stack.len() == 0 { return None; }
        while let Some(n) = self.stack.last_mut().unwrap().0.next() {
            self.stack.push( ( (self.visit)(n), n ) );
        }
        Some(self.stack.pop().unwrap().1)
    }
}
2 Likes

Doing it with a stack as you have done is I think the current, standard way to do it. I have just been programming a similar iterator this last week ( I found it quite tricky ).

However, it says here: "generators will likely be extended to ergonomic implementations of iterators and other primitives in the future".

You might consider using an array rather than Vec for the stack, that's what I did. It avoids any heap allocations.

One thing to distinguish is:

  • external iteration

    (for x in ... { ... } syntax, whereby the iterator yields items back to the caller),

    This is what the standard Iterator trait is for.

    Workflow
    • Caller: give me an item, please.

    • Iterator: here it is (or not, if exhausted).

    • Caller: Thanks (politeness may compile down to a no-op :grin:)—proceeds to do stuff with that item, such as print!-ing it or performing an early return from its own function.

    • Caller: Give me an(other) item, please.

    • Iterator: Here it is (or not etc.),

    • etc.

  • internal iteration

    (... .for_each(|x| { ... }) syntax or other adaptors, whereby the iterator is actually an object taking, as a query, a processing request over its internal items

    Workflow
    • Caller (thinking out loud): hmm, if I were to have each item of that container, I would like for them to be printed, one at a time

    • Container: I can do that for you, if you so wish.

    • Caller: Really? Great. Yes, please, for each item, I'd like you to do that.

    • Container: Roger roger, will do that. I'll be back once I'm done—proceeds to go to the back room, at which point some metallic clanging noises can be heard.

Now, depending on the Rust data structures, sometimes an item can only be obtained out of a drawer that is about to dangle, and that needs that it be held in place while operating on the item. In that case, it can be extremely painful, or incur in a seemingly unnecessary runtime cost, or even downright impossible to be able to offer external iteration, i.e., to implement Iterator.

Example [RefCell<Item>]

It is not possible (within sane code), to implement the following API:

fn iter_refcells<'iter, Item> (
    refcells: &'iter [RefCell<Item>],
) -> impl 'iter + Iterator<Item = &'iter Item>

It is, however, quite trivial to implement internal iteration on it:

impl<Item> InternallyIterable for [RefCell<Item>] {
    type Item = Item;
    
    fn try_fold<Acc, Err, F> (
        self: &'_ Self,
        mut acc: Acc,
        mut f: F,
    ) -> Result<Acc, Err>
    where
        F : FnMut(Acc, &'_ Self::Item) -> Result<Acc, Err>,
    {
        for refcell in self {
            acc = f(acc, &*refcell.borrow())?;
        }
        Ok(acc)
    }
}

fn main ()
{
    let refcells = [RefCell::new(42), RefCell::new(27)];
    let mut sum = 0;
    refcells.for_each(|x| {
        eprintln!("Got {}", x);
        sum += x;
    });
    dbg!(sum);
}

At this point, we can see that your recursive data-structure can easily implement InternallyIterable or something similar, like @H2CO3 did.

External iteration can be possible too here, so as long as we avoid creating a recursive type, which is possible by using type-erased iterators, by using virtual method for the dynamic dispatch at runtime:
that is @jethrogb solution.

1 Like