From 532e2ed770bf7e254c8adb56fd4bf089c0e751ab Mon Sep 17 00:00:00 2001
From: Tobias Schmitt <ts@bobcat.fritz.box>
Date: Thu, 21 Jul 2022 21:12:26 +0200
Subject: [PATCH] adds iter_(read/write)_variable_identifiers

seperate iterator returns for read and write variable identifiers.
this is useful in our project and i suspect it might be of use for
others too.
---
 src/tree/mod.rs      | 42 ++++++++++++++++++++++++++++++++++++++++++
 tests/integration.rs | 12 +++++++++++-
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/src/tree/mod.rs b/src/tree/mod.rs
index 1956e56..4c5e39f 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<Item = &str> {
+        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<Item = &str> {
+        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);