From e328adb97b04f5699f524f0de8d04a1134b807ae Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <isibboi@gmail.com>
Date: Fri, 12 Apr 2019 19:27:34 +0200
Subject: [PATCH] Implement `Node::iter_identifiers()`

This commit implements an iterator over the identifiers in an expression.
The identifiers are iterated pre-order, meaning that higher nodes in the tree are emitted first.
The identifier iterator is based on a general iterator over all nodes in an operator tree, which can be used for further featues.

Relates to #37
---
 src/operator/mod.rs | 13 +++++++++++
 src/tree/iter.rs    | 53 +++++++++++++++++++++++++++++++++++++++++++++
 src/tree/mod.rs     | 21 ++++++++++++++++++
 3 files changed, 87 insertions(+)
 create mode 100644 src/tree/iter.rs

diff --git a/src/operator/mod.rs b/src/operator/mod.rs
index 4252854..1f6e7a0 100644
--- a/src/operator/mod.rs
+++ b/src/operator/mod.rs
@@ -37,6 +37,11 @@ pub trait Operator: Debug + Display {
     fn eval_mut(&self, arguments: &[Value], context: &mut dyn Context) -> EvalexprResult<Value> {
         self.eval(arguments, context)
     }
+
+    /// Returns an identifier if this operator is a function or variable identifier, or `None` otherwise.
+    fn identifier(&self) -> Option<&str> {
+        None
+    }
 }
 
 #[derive(Debug)]
@@ -681,6 +686,10 @@ impl Operator for VariableIdentifier {
             ))
         }
     }
+
+    fn identifier(&self) -> Option<&str> {
+        Some(&self.identifier)
+    }
 }
 
 impl Operator for FunctionIdentifier {
@@ -715,4 +724,8 @@ impl Operator for FunctionIdentifier {
             ))
         }
     }
+
+    fn identifier(&self) -> Option<&str> {
+        Some(&self.identifier)
+    }
 }
diff --git a/src/tree/iter.rs b/src/tree/iter.rs
new file mode 100644
index 0000000..672261f
--- /dev/null
+++ b/src/tree/iter.rs
@@ -0,0 +1,53 @@
+use Node;
+use std::slice::Iter;
+
+/// An iterator that traverses an operator tree in pre-order.
+pub struct NodeIter<'a> {
+    stack: Vec<Iter<'a, Node>>,
+}
+
+impl<'a> NodeIter<'a> {
+    fn new(node: &'a Node) -> Self {
+        NodeIter {
+            stack: vec![node.children.iter()],
+        }
+    }
+}
+
+impl<'a> Iterator for NodeIter<'a> {
+    type Item = &'a Node;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let mut pop_stack = false;
+            let mut result = None;
+
+            if let Some(last) = self.stack.last_mut() {
+                if let Some(next) = last.next() {
+                    result = Some(next);
+                } else {
+                    pop_stack = true;
+                }
+            } else {
+                return None;
+            }
+
+            if pop_stack {
+                // Can not fail because we borrowed last before.
+                self.stack.pop().unwrap();
+            }
+
+            if let Some(result) = result {
+                self.stack.push(result.children.iter());
+                return Some(result);
+            }
+        }
+    }
+}
+
+impl Node {
+    /// Returns an iterator over all nodes in this tree.
+    pub fn iter(&self) -> impl Iterator<Item=&Node> {
+        NodeIter::new(self)
+    }
+}
\ No newline at end of file
diff --git a/src/tree/mod.rs b/src/tree/mod.rs
index 12b3de7..775fcb6 100644
--- a/src/tree/mod.rs
+++ b/src/tree/mod.rs
@@ -13,6 +13,7 @@ use crate::{
 };
 
 mod display;
+mod iter;
 
 /// A node in the operator tree.
 /// The operator tree is created by the crate-level `build_operator_tree` method.
@@ -49,6 +50,26 @@ impl Node {
         Self::new(RootNode)
     }
 
+    /// Returns an iterator over all identifiers in this expression.
+    /// Each occurrence of an identifier is returned separately.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use evalexpr::*;
+    ///
+    /// let tree = build_operator_tree("a + b + c * f()").unwrap(); // Do proper error handling here
+    /// let mut iter = tree.iter_identifiers();
+    /// assert_eq!(iter.next(), Some("a"));
+    /// assert_eq!(iter.next(), Some("b"));
+    /// assert_eq!(iter.next(), Some("c"));
+    /// assert_eq!(iter.next(), Some("f"));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn iter_identifiers(&self) -> impl Iterator<Item = &str> {
+        self.iter().filter_map(|node| node.operator.identifier())
+    }
+
     /// Evaluates the operator tree rooted at this node with the given context.
     ///
     /// Fails, if one of the operators in the expression tree fails.