Node traversal iterator (#3845)

Implement an iterator for pre- and post-order traversals.

I believe that this will be useful in pre-processing passes, many of
which do postorder traversals that they implement by hand.

Right now this iterator does not support modification of the traversal
pattern, but we could add this later on, if we want it.

Co-authored-by: Andres Noetzli <andres.noetzli@gmail.com>
Co-authored-by: Mathias Preiner <mathias.preiner@gmail.com>

3 files changed, 283 insertions, 0 deletions

diff --git a/src/expr/CMakeLists.txt b/src/expr/CMakeLists.txt
index d1faa8ffb..00bd121cb 100644
--- a/src/expr/CMakeLists.txt
+++ b/src/expr/CMakeLists.txt
@@ -29,6 +29,8 @@ libcvc4_add_sources(
   node_self_iterator.h
   node_trie.cpp
   node_trie.h
+  node_traversal.cpp
+  node_traversal.h
   node_value.cpp
   node_value.h
   node_visitor.h
diff --git a/src/expr/node_traversal.cpp b/src/expr/node_traversal.cpp
new file mode 100644
index 000000000..9e7a82c24
--- /dev/null
+++ b/src/expr/node_traversal.cpp
@@ -0,0 +1,150 @@
+/*********************                                                        */
+/*! \file node_traversal.cpp
+ ** \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 "node_traversal.h"
+
+namespace CVC4 {
+
+NodeDfsIterator::NodeDfsIterator(TNode n, bool postorder)
+    : d_stack{n},
+      d_visited(),
+      d_postorder(postorder),
+      d_current(TNode())
+{
+}
+
+NodeDfsIterator::NodeDfsIterator(bool postorder)
+    : d_stack(),
+      d_visited(),
+      d_postorder(postorder),
+      d_current(TNode())
+{
+}
+
+NodeDfsIterator& NodeDfsIterator::operator++()
+{
+  // If we were just constructed, advance to first visit, **before**
+  // advancing past it to the next visit (below).
+  initializeIfUninitialized();
+
+  // Advance to the next visit
+  advanceToNextVisit();
+  return *this;
+}
+
+NodeDfsIterator NodeDfsIterator::operator++(int)
+{
+  NodeDfsIterator copyOfOld(*this);
+  ++*this;
+  return copyOfOld;
+}
+
+TNode& NodeDfsIterator::operator*()
+{
+  // If we were just constructed, advance to first visit
+  initializeIfUninitialized();
+  Assert(!d_current.isNull());
+
+  return d_current;
+}
+
+bool NodeDfsIterator::operator==(const NodeDfsIterator& other) const
+{
+  // The stack and current node uniquely represent traversal state. We need not
+  // use the scheduled node set.
+  //
+  // Users should not compare iterators for traversals of different nodes.
+  Assert(d_postorder == other.d_postorder);
+  return d_stack == other.d_stack && d_current == other.d_current;
+}
+
+bool NodeDfsIterator::operator!=(const NodeDfsIterator& other) const
+{
+  return !(*this == other);
+}
+
+void NodeDfsIterator::advanceToNextVisit()
+{
+  // While a node is enqueued and we're not at the right visit type
+  while (!d_stack.empty())
+  {
+    TNode back = d_stack.back();
+    auto visitEntry = d_visited.find(back);
+    if (visitEntry == d_visited.end())
+    {
+      // if we haven't pre-visited this node, pre-visit it
+      d_visited[back] = false;
+      d_current = back;
+      // Use integer underflow to reverse-iterate
+      for (size_t n = back.getNumChildren(), i = n - 1; i < n; --i)
+      {
+        d_stack.push_back(back[i]);
+      }
+      if (!d_postorder)
+      {
+        return;
+      }
+    }
+    else if (!d_postorder || visitEntry->second)
+    {
+      // if we're previsiting or we've already post-visited this node: skip it
+      d_stack.pop_back();
+    }
+    else
+    {
+      // otherwise, this is a post-visit
+      visitEntry->second = true;
+      d_current = back;
+      d_stack.pop_back();
+      return;
+    }
+  }
+  // We're at the end of the traversal: nullify the current node to agree
+  // with the "end" iterator.
+  d_current = TNode();
+}
+
+void NodeDfsIterator::initializeIfUninitialized()
+{
+  if (d_current.isNull())
+  {
+    advanceToNextVisit();
+  }
+}
+
+NodeDfsIterable::NodeDfsIterable(TNode n) : d_node(n), d_postorder(true) {}
+
+NodeDfsIterable& NodeDfsIterable::inPostorder()
+{
+  d_postorder = true;
+  return *this;
+}
+
+NodeDfsIterable& NodeDfsIterable::inPreorder()
+{
+  d_postorder = false;
+  return *this;
+}
+
+NodeDfsIterator NodeDfsIterable::begin() const
+{
+  return NodeDfsIterator(d_node, d_postorder);
+}
+
+NodeDfsIterator NodeDfsIterable::end() const
+{
+  return NodeDfsIterator(d_postorder);
+}
+
+}  // namespace CVC4
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