From 697fbb5d4592230c82c838bf257d5000686b7de1 Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <isibboi@gmail.com>
Date: Sat, 13 Apr 2019 20:15:26 +0200
Subject: [PATCH] Implement iterators over variable and function identifiers

Relates to #37
---
 CHANGELOG.md        |  1 +
 src/operator/mod.rs | 18 ++++++++++++++++++
 src/tree/mod.rs     | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 62da407..c990fd2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@
 ### Added
 
  * Iterator over all identifiers within an expression, including duplicates
+ * Iterators over only variable or only function identifiers within an expression, including duplicates
  * Overload the `+` operator to concatenate strings
  * Overload `<`, `<=`, `>` and `>=` for strings using lexical ordering (Note: `==` and `!=` compare strings as expected)
  * Add `len`, `str::regex_matches`, `str::regex_replace`, `str::to_lowercase`, `str::to_uppercase`, `str::trim` functions for strings
diff --git a/src/operator/mod.rs b/src/operator/mod.rs
index e367546..04b1e45 100644
--- a/src/operator/mod.rs
+++ b/src/operator/mod.rs
@@ -42,6 +42,16 @@ pub trait Operator: Debug + Display {
     fn identifier(&self) -> Option<&str> {
         None
     }
+
+    /// Returns a variable identifier if this operator is a variable identifier, or `None` otherwise.
+    fn variable_identifier(&self) -> Option<&str> {
+        None
+    }
+
+    /// Returns a function identifier if this operator is a function identifier, or `None` otherwise.
+    fn function_identifier(&self) -> Option<&str> {
+        None
+    }
 }
 
 #[derive(Debug)]
@@ -719,6 +729,10 @@ impl Operator for VariableIdentifier {
     fn identifier(&self) -> Option<&str> {
         Some(&self.identifier)
     }
+
+    fn variable_identifier(&self) -> Option<&str> {
+        Some(&self.identifier)
+    }
 }
 
 impl Operator for FunctionIdentifier {
@@ -757,4 +771,8 @@ impl Operator for FunctionIdentifier {
     fn identifier(&self) -> Option<&str> {
         Some(&self.identifier)
     }
+
+    fn function_identifier(&self) -> Option<&str> {
+        Some(&self.identifier)
+    }
 }
diff --git a/src/tree/mod.rs b/src/tree/mod.rs
index 775fcb6..fb65b25 100644
--- a/src/tree/mod.rs
+++ b/src/tree/mod.rs
@@ -70,6 +70,42 @@ impl Node {
         self.iter().filter_map(|node| node.operator.identifier())
     }
 
+    /// Returns an iterator over all variable identifiers in this expression.
+    /// Each occurrence of a variable identifier is returned separately.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use evalexpr::*;
+    ///
+    /// let tree = build_operator_tree("a + f(b + c)").unwrap(); // Do proper error handling here
+    /// let mut iter = tree.iter_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_variable_identifiers(&self) -> impl Iterator<Item = &str> {
+        self.iter().filter_map(|node| node.operator.variable_identifier())
+    }
+
+    /// Returns an iterator over all function identifiers in this expression.
+    /// Each occurrence of a function identifier is returned separately.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use evalexpr::*;
+    ///
+    /// let tree = build_operator_tree("a + f(b + c)").unwrap(); // Do proper error handling here
+    /// let mut iter = tree.iter_function_identifiers();
+    /// assert_eq!(iter.next(), Some("f"));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn iter_function_identifiers(&self) -> impl Iterator<Item = &str> {
+        self.iter().filter_map(|node| node.operator.function_identifier())
+    }
+
     /// Evaluates the operator tree rooted at this node with the given context.
     ///
     /// Fails, if one of the operators in the expression tree fails.