API design for navigating a tree

I'm designing a tree structure around the idea that higher nodes "approximate" their children in some sense, and when iterating, the calling code will want to make decisions on a per-node basis about whether the iteration should continue descending the tree or to skip the rest of the subtree and go to the current node's siblings.

However, I don't know a good way to express this as a code interface. If I wanted to use straightforward iteration, I'd need to separate the decision-making logic from actually acting on the nodes, like this:

enum Control {
	Descend,
	Continue
}

fn flow(n: &Node) -> Control { /* Decide where to go */ }

fn do_something(t: &Tree) {
	for node in t.iter(flow)
	{
		/* Do something with the node */
	}
}

I'm not fond of this design because the action I want to take with the node might depend on intermediate values that I'm also using to make the control flow decision, and it'd have to be calculated twice.

The other idea I had was to use a generator-like structure, something like this:

fn update_particle(t: &Tree) {
    let mut cursor = t.query();
    let mut choice = Control::Descend;

    while let Yield(node) = cursor.next(choice) {
        /* Do something with the node, but also update choice */
    }
}

I don't like this because as you can see, it's not actually an iterator anymore and the while let stands out as unusual. I also don't like mutable state if it can be avoided, and this pattern pretty much requires it.

Does anyone have any suggestions about how this could be improved, or what would be more idiomatic?

This looks like a good use case for the newish std::ops::ControlFlow:

pub enum ControlFlow<B, C = ()> {
    Continue(C),
    Break(B),
}

The docs even contain an example of using it for a tree traversal :blush:

1 Like