From 0d549b1c36797e3f42343c38aeb4b8fb6a874bc9 Mon Sep 17 00:00:00 2001 From: Ben Weinstein-Raun Date: Fri, 2 Jun 2023 18:56:04 -0700 Subject: [PATCH] add mutable iteration over identifiers --- src/tree/iter.rs | 49 +++++++++++++++- src/tree/mod.rs | 150 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 197 insertions(+), 2 deletions(-) diff --git a/src/tree/iter.rs b/src/tree/iter.rs index acf64f7..931f42e 100644 --- a/src/tree/iter.rs +++ b/src/tree/iter.rs @@ -1,5 +1,5 @@ -use crate::Node; -use std::slice::Iter; +use crate::{Node, operator::Operator}; +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>, +} + +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 { + 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 { NodeIter::new(self) } + + /// Returns a mutable iterator over all operators in this tree. + pub fn iter_operator_mut(&mut self) -> impl Iterator { + NodeIterMut::new(self) + } } + diff --git a/src/tree/mod.rs b/src/tree/mod.rs index 347f787..35a3f96 100644 --- a/src/tree/mod.rs +++ b/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 { + 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 { + 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 { + 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 { + 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 { + 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.