1 files changed, 131 insertions, 0 deletions
diff --git a/src/expr/node_traversal.h b/src/expr/node_traversal.h
new file mode 100644
index 000000000..fffc1d746
--- /dev/null
+++ b/src/expr/node_traversal.h
@@ -0,0 +1,131 @@
+/*********************                                                        */
+/*! \file node_traversal.h
+ ** \verbatim
+ ** Top contributors (to current version):
+ **   Alex Ozdemir
+ ** This file is part of the CVC4 project.
+ ** Copyright (c) 2009-2020 by the authors listed in the file AUTHORS
+ ** in the top-level source directory) and their institutional affiliations.
+ ** All rights reserved.  See the file COPYING in the top-level source
+ ** directory for licensing information.\endverbatim
+ **
+ ** \brief Iterators for traversing nodes.
+ **/
+
+#include "cvc4_private.h"
+
+#ifndef CVC4__EXPR__NODE_TRAVERSAL_H
+#define CVC4__EXPR__NODE_TRAVERSAL_H
+
+#include <cstddef>
+#include <iterator>
+#include <unordered_map>
+#include <vector>
+
+#include "expr/node.h"
+
+namespace CVC4 {
+
+// Iterator for traversing a node in post-order
+// It does DAG-traversal, so indentical sub-nodes will be visited once only.
+class NodeDfsIterator
+{
+ public:
+  // STL type definitions for an iterator
+  using value_type = TNode;
+  using pointer = TNode*;
+  using reference = TNode&;
+  using iterator_category = std::forward_iterator_tag;
+  using difference_type = std::ptrdiff_t;
+
+  // Construct a traversal iterator beginning at `n`
+  NodeDfsIterator(TNode n, bool postorder);
+  // Construct an end-of-traversal iterator
+  NodeDfsIterator(bool postorder);
+
+  // Move/copy construction and assignment. Destructor.
+  NodeDfsIterator(NodeDfsIterator&&) = default;
+  NodeDfsIterator& operator=(NodeDfsIterator&&) = default;
+  NodeDfsIterator(NodeDfsIterator&) = default;
+  NodeDfsIterator& operator=(NodeDfsIterator&) = default;
+  ~NodeDfsIterator() = default;
+
+  // Preincrement
+  NodeDfsIterator& operator++();
+  // Postincrement
+  NodeDfsIterator operator++(int);
+  // Dereference
+  reference operator*();
+  // Equals
+  bool operator==(const NodeDfsIterator&) const;
+  // Not equals
+  bool operator!=(const NodeDfsIterator&) const;
+
+ private:
+  // While we're not at an appropriate visit (see d_postorder), advance.
+  // In each step:
+  //   * enqueue children of a not-yet-pre-visited node (and mark it
+  //     previsited)
+  //   * pop a not-yet-post-visited node (and mark it post-visited)
+  //   * pop an already post-visited node.
+  // After calling this, `d_current` will be changed to the next node, if there
+  // is another node in the traversal.
+  void advanceToNextVisit();
+
+  // If this iterator hasn't been dereferenced or incremented yet, advance to
+  // first visit.
+  // Necessary because we are lazy and don't find our first visit node at
+  // construction time.
+  void initializeIfUninitialized();
+
+  // Stack of nodes to visit.
+  std::vector<TNode> d_stack;
+
+  // Whether (and how) we've visited a node.
+  // Absent if we haven't visited it.
+  // Set to `false` if we've already pre-visited it (enqueued its children).
+  // Set to `true` if we've also already post-visited it.
+  std::unordered_map<TNode, bool, TNodeHashFunction> d_visited;
+
+  // Whether this is a post-order iterator (the alternative is pre-order)
+  bool d_postorder;
+
+  // Whether this iterator has been initialized (advanced to its first
+  // visit)
+  bool d_initialized;
+
+  // Current referent node. A valid node to visit if non-null.
+  // Null after construction (but before first access) and at the end.
+  TNode d_current;
+};
+
+// Node wrapper that is iterable in DAG post-order
+class NodeDfsIterable
+{
+ public:
+  NodeDfsIterable(TNode n);
+
+  // Modifying the traversal order
+  // Modify this iterable to be in post-order (default)
+  NodeDfsIterable& inPostorder();
+  // Modify this iterable to be in pre-order
+  NodeDfsIterable& inPreorder();
+
+  // Move/copy construction and assignment. Destructor.
+  NodeDfsIterable(NodeDfsIterable&&) = default;
+  NodeDfsIterable& operator=(NodeDfsIterable&&) = default;
+  NodeDfsIterable(NodeDfsIterable&) = default;
+  NodeDfsIterable& operator=(NodeDfsIterable&) = default;
+  ~NodeDfsIterable() = default;
+
+  NodeDfsIterator begin() const;
+  NodeDfsIterator end() const;
+
+ private:
+  TNode d_node;
+  bool d_postorder;
+};
+
+}  // namespace CVC4
+
+#endif  // CVC4__EXPR__NODE_TRAVERSAL_H