expressive/src/tree/iter.rs
Sebastian Schmidt fd3063fc64 Reduce unwrap()s.
Relates to #14
2021-05-28 09:40:58 +03:00

49 lines
1.1 KiB
Rust

use crate::Node;
use std::slice::Iter;
/// An iterator that traverses an operator tree in pre-order.
pub struct NodeIter<'a> {
stack: Vec<Iter<'a, Node>>,
}
impl<'a> NodeIter<'a> {
fn new(node: &'a Node) -> Self {
NodeIter {
stack: vec![node.children.iter()],
}
}
}
impl<'a> Iterator for NodeIter<'a> {
type Item = &'a Node;
fn next(&mut self) -> Option<Self::Item> {
loop {
let mut result = None;
if let Some(last) = self.stack.last_mut() {
if let Some(next) = last.next() {
result = Some(next);
} else {
// Can not fail because we just borrowed last.
self.stack.pop().unwrap();
}
} else {
return None;
}
if let Some(result) = result {
self.stack.push(result.children.iter());
return Some(result);
}
}
}
}
impl Node {
/// Returns an iterator over all nodes in this tree.
pub fn iter(&self) -> impl Iterator<Item = &Node> {
NodeIter::new(self)
}
}