evalexpr/tree/
iter.rs

1use crate::{operator::Operator, value::numeric_types::EvalexprNumericTypes, Node};
2use std::slice::{Iter, IterMut};
3
4/// An iterator that traverses an operator tree in pre-order.
5pub struct NodeIter<'a, NumericTypes: EvalexprNumericTypes> {
6    stack: Vec<Iter<'a, Node<NumericTypes>>>,
7}
8
9impl<'a, NumericTypes: EvalexprNumericTypes> NodeIter<'a, NumericTypes> {
10    fn new(node: &'a Node<NumericTypes>) -> Self {
11        NodeIter {
12            stack: vec![node.children.iter()],
13        }
14    }
15}
16
17impl<'a, NumericTypes: EvalexprNumericTypes> Iterator for NodeIter<'a, NumericTypes> {
18    type Item = &'a Node<NumericTypes>;
19
20    fn next(&mut self) -> Option<Self::Item> {
21        loop {
22            let mut result = None;
23
24            if let Some(last) = self.stack.last_mut() {
25                if let Some(next) = last.next() {
26                    result = Some(next);
27                } else {
28                    // Can not fail because we just borrowed last.
29                    // We just checked that the iterator is empty, so we can safely discard it.
30                    let _ = self.stack.pop().unwrap();
31                }
32            } else {
33                return None;
34            }
35
36            if let Some(result) = result {
37                self.stack.push(result.children.iter());
38                return Some(result);
39            }
40        }
41    }
42}
43
44/// An iterator that mutably traverses an operator tree in pre-order.
45pub struct OperatorIterMut<'a, NumericTypes: EvalexprNumericTypes> {
46    stack: Vec<IterMut<'a, Node<NumericTypes>>>,
47}
48
49impl<'a, NumericTypes: EvalexprNumericTypes> OperatorIterMut<'a, NumericTypes> {
50    fn new(node: &'a mut Node<NumericTypes>) -> Self {
51        OperatorIterMut {
52            stack: vec![node.children.iter_mut()],
53        }
54    }
55}
56
57impl<'a, NumericTypes: EvalexprNumericTypes> Iterator for OperatorIterMut<'a, NumericTypes> {
58    type Item = &'a mut Operator<NumericTypes>;
59
60    fn next(&mut self) -> Option<Self::Item> {
61        loop {
62            let mut result = None;
63
64            if let Some(last) = self.stack.last_mut() {
65                if let Some(next) = last.next() {
66                    result = Some(next);
67                } else {
68                    // Can not fail because we just borrowed last.
69                    // We just checked that the iterator is empty, so we can safely discard it.
70                    let _ = self.stack.pop().unwrap();
71                }
72            } else {
73                return None;
74            }
75
76            if let Some(result) = result {
77                self.stack.push(result.children.iter_mut());
78                return Some(&mut result.operator);
79            }
80        }
81    }
82}
83
84impl<NumericTypes: EvalexprNumericTypes> Node<NumericTypes> {
85    /// Returns an iterator over all nodes in this tree.
86    pub fn iter(&self) -> impl Iterator<Item = &Node<NumericTypes>> {
87        NodeIter::new(self)
88    }
89
90    /// Returns a mutable iterator over all operators in this tree.
91    pub fn iter_operators_mut(&mut self) -> impl Iterator<Item = &mut Operator<NumericTypes>> {
92        OperatorIterMut::new(self)
93    }
94}