diff --git a/src/tree/mod.rs b/src/tree/mod.rs index 1956e56..c225994 100644 --- a/src/tree/mod.rs +++ b/src/tree/mod.rs @@ -99,6 +99,48 @@ impl Node { }) } + /// Returns an iterator over all read variable identifiers in this expression. + /// Each occurrence of a variable identifier is returned separately. + /// + /// # Examples + /// + /// ```rust + /// use evalexpr::*; + /// + /// let tree = build_operator_tree("d = a + f(b + c)").unwrap(); // Do proper error handling here + /// let mut iter = tree.iter_read_variable_identifiers(); + /// assert_eq!(iter.next(), Some("a")); + /// assert_eq!(iter.next(), Some("b")); + /// assert_eq!(iter.next(), Some("c")); + /// assert_eq!(iter.next(), None); + /// ``` + pub fn iter_read_variable_identifiers(&self) -> impl Iterator { + self.iter().filter_map(|node| match node.operator() { + Operator::VariableIdentifierRead { identifier } => Some(identifier.as_str()), + _ => None, + }) + } + + /// Returns an iterator over all write variable identifiers in this expression. + /// Each occurrence of a variable identifier is returned separately. + /// + /// # Examples + /// + /// ```rust + /// use evalexpr::*; + /// + /// let tree = build_operator_tree("d = a + f(b + c)").unwrap(); // Do proper error handling here + /// let mut iter = tree.iter_write_variable_identifiers(); + /// assert_eq!(iter.next(), Some("d")); + /// assert_eq!(iter.next(), None); + /// ``` + pub fn iter_write_variable_identifiers(&self) -> impl Iterator { + self.iter().filter_map(|node| match node.operator() { + Operator::VariableIdentifierWrite { identifier } => Some(identifier.as_str()), + _ => None, + }) + } + /// Returns an iterator over all function identifiers in this expression. /// Each occurrence of a function identifier is returned separately. /// diff --git a/tests/integration.rs b/tests/integration.rs index 66468ab..f5e8354 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -1616,16 +1616,26 @@ fn test_error_constructors() { #[test] fn test_iterators() { - let tree = build_operator_tree("5 + 3 + fun(4) + var").unwrap(); + let tree = build_operator_tree("writevar = 5 + 3 + fun(4) + var").unwrap(); let mut iter = tree.iter_identifiers(); + assert_eq!(iter.next(), Some("writevar")); assert_eq!(iter.next(), Some("fun")); assert_eq!(iter.next(), Some("var")); assert_eq!(iter.next(), None); let mut iter = tree.iter_variable_identifiers(); + assert_eq!(iter.next(), Some("writevar")); assert_eq!(iter.next(), Some("var")); assert_eq!(iter.next(), None); + let mut iter = tree.iter_read_variable_identifiers(); + assert_eq!(iter.next(), Some("var")); + assert_eq!(iter.next(), None); + + let mut iter = tree.iter_write_variable_identifiers(); + assert_eq!(iter.next(), Some("writevar")); + assert_eq!(iter.next(), None); + let mut iter = tree.iter_function_identifiers(); assert_eq!(iter.next(), Some("fun")); assert_eq!(iter.next(), None);