add mutable iteration over identifiers

src/tree/iter.rs

@@ -1,5 +1,5 @@
-use crate::Node;
+use crate::{Node, operator::Operator};
-use std::slice::Iter;
+use std::slice::{Iter, IterMut};
 
 /// An iterator that traverses an operator tree in pre-order.
 pub struct NodeIter<'a> {
@@ -41,9 +41,54 @@ impl<'a> Iterator for NodeIter<'a> {
     }
 }
 
+pub struct NodeIterMut<'a> {
+    stack: Vec<IterMut<'a, Node>>,
+}
+
+impl<'a> NodeIterMut<'a> {
+    fn new(node: &'a mut Node) -> Self {
+        NodeIterMut {
+            stack: vec![node.children.iter_mut()]
+        }
+    }
+}
+
+impl<'a> Iterator for NodeIterMut<'a> {
+    type Item = &'a mut Operator;
+
+    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.
+                    // We just checked that the iterator is empty, so we can safely discard it.
+                    let _ = self.stack.pop().unwrap();
+                }
+            } else {
+                return None;
+            }
+
+            if let Some(result) = result {
+                self.stack.push(result.children.iter_mut());
+                return Some(&mut result.operator);
+            }
+        }
+    }
+}
+
 impl Node {
     /// Returns an iterator over all nodes in this tree.
     pub fn iter(&self) -> impl Iterator<Item = &Node> {
         NodeIter::new(self)
     }
+
+    /// Returns a mutable iterator over all operators in this tree.
+    pub fn iter_operator_mut(&mut self) -> impl Iterator<Item = &mut Operator> {
+        NodeIterMut::new(self)
+    }
 }
+

src/tree/mod.rs

@@ -76,6 +76,36 @@ impl Node {
         })
     }
 
+    /// Returns an iterator over all identifiers in this expression, allowing mutation.
+    /// Each occurrence of an identifier is returned separately.
+    /// # Examples
+    ///
+    /// ```rust
+    /// use evalexpr::*;
+    ///
+    /// let mut tree = build_operator_tree("a + b + c * f()").unwrap(); // Do proper error handling here
+    ///
+    /// for identifier in tree.iter_identifiers_mut() {
+    ///     *identifier = String::from("x");
+    /// }
+    ///
+    /// let mut iter = tree.iter_identifiers();
+    ///
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn iter_identifiers_mut(&mut self) -> impl Iterator<Item = &mut String> {
+        self.iter_operator_mut().filter_map(|operator| match operator {
+            Operator::VariableIdentifierWrite { identifier }
+            | Operator::VariableIdentifierRead { identifier }
+            | Operator::FunctionIdentifier { identifier } => Some(identifier),
+            _ => None
+        })
+    }
+
     /// Returns an iterator over all variable identifiers in this expression.
     /// Each occurrence of a variable identifier is returned separately.
     ///
@@ -99,6 +129,36 @@ impl Node {
         })
     }
 
+    /// Returns an iterator over all variable identifiers in this expression, allowing mutation.
+    /// Each occurrence of a variable identifier is returned separately.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use evalexpr::*;
+    ///
+    /// let mut tree = build_operator_tree("a + b + c * f()").unwrap(); // Do proper error handling here
+    ///
+    /// for identifier in tree.iter_variable_identifiers_mut() {
+    ///     *identifier = String::from("x");
+    /// }
+    ///
+    /// let mut iter = tree.iter_identifiers();
+    ///
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("f"));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn iter_variable_identifiers_mut(&mut self) -> impl Iterator<Item = &mut String> {
+        self.iter_operator_mut().filter_map(|operator| match operator {
+            Operator::VariableIdentifierWrite { identifier }
+            | Operator::VariableIdentifierRead { identifier } => Some(identifier),
+            _ => None,
+        })
+    }
+
     /// Returns an iterator over all read variable identifiers in this expression.
     /// Each occurrence of a variable identifier is returned separately.
     ///
@@ -121,6 +181,36 @@ impl Node {
         })
     }
 
+    /// Returns an iterator over all read variable identifiers in this expression, allowing mutation.
+    /// Each occurrence of a variable identifier is returned separately.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use evalexpr::*;
+    ///
+    /// let mut tree = build_operator_tree("d = a + f(b + c)").unwrap(); // Do proper error handling here
+    ///
+    /// for identifier in tree.iter_read_variable_identifiers_mut() {
+    ///     *identifier = String::from("x");
+    /// }
+    ///
+    /// let mut iter = tree.iter_identifiers();
+    ///
+    /// assert_eq!(iter.next(), Some("d"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("f"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn iter_read_variable_identifiers_mut(&mut self) -> impl Iterator<Item = &mut String> {
+        self.iter_operator_mut().filter_map(|operator| match operator {
+            Operator::VariableIdentifierRead { identifier } => Some(identifier),
+            _ => None,
+        })
+    }
+
     /// Returns an iterator over all write variable identifiers in this expression.
     /// Each occurrence of a variable identifier is returned separately.
     ///
@@ -141,6 +231,36 @@ impl Node {
         })
     }
 
+    /// Returns an iterator over all write variable identifiers in this expression, allowing mutation.
+    /// Each occurrence of a variable identifier is returned separately.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use evalexpr::*;
+    ///
+    /// let mut tree = build_operator_tree("d = a + f(b + c)").unwrap(); // Do proper error handling here
+    ///
+    /// for identifier in tree.iter_write_variable_identifiers_mut() {
+    ///     *identifier = String::from("x");
+    /// }
+    ///
+    /// let mut iter = tree.iter_identifiers();
+    ///
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("a"));
+    /// assert_eq!(iter.next(), Some("f"));
+    /// assert_eq!(iter.next(), Some("b"));
+    /// assert_eq!(iter.next(), Some("c"));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn iter_write_variable_identifiers_mut(&mut self) -> impl Iterator<Item = &mut String> {
+        self.iter_operator_mut().filter_map(|operator| match operator {
+            Operator::VariableIdentifierWrite { identifier } => Some(identifier),
+            _ => None,
+        })
+    }
+
     /// Returns an iterator over all function identifiers in this expression.
     /// Each occurrence of a function identifier is returned separately.
     ///
@@ -161,6 +281,36 @@ impl Node {
         })
     }
 
+    /// Returns an iterator over all function identifiers in this expression, allowing mutation.
+    /// Each occurrence of a variable identifier is returned separately.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use evalexpr::*;
+    ///
+    /// let mut tree = build_operator_tree("d = a + f(b + c)").unwrap(); // Do proper error handling here
+    ///
+    /// for identifier in tree.iter_function_variable_identifiers_mut() {
+    ///     *identifier = String::from("x");
+    /// }
+    ///
+    /// let mut iter = tree.iter_identifiers();
+    ///
+    /// assert_eq!(iter.next(), Some("d"));
+    /// assert_eq!(iter.next(), Some("a"));
+    /// assert_eq!(iter.next(), Some("x"));
+    /// assert_eq!(iter.next(), Some("b"));
+    /// assert_eq!(iter.next(), Some("c"));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn iter_function_variable_identifiers_mut(&mut self) -> impl Iterator<Item = &mut String> {
+        self.iter_operator_mut().filter_map(|operator| match operator {
+            Operator::FunctionIdentifier { identifier } => Some(identifier),
+            _ => None,
+        })
+    }
+
     /// Evaluates the operator tree rooted at this node with the given context.
     ///
     /// Fails, if one of the operators in the expression tree fails.