1 files changed, 19 insertions, 33 deletions

diff --git a/src/expr/node_traversal.cpp b/src/expr/node_traversal.cpp
index ad1a9ec71..5968c4348 100644
--- a/src/expr/node_traversal.cpp
+++ b/src/expr/node_traversal.cpp
@@ -2,7 +2,7 @@
 /*! \file node_traversal.cpp
  ** \verbatim
  ** Top contributors (to current version):
- **   Alex Ozdemir
+ **   Alex Ozdemir, Andres Noetzli
  ** 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.
@@ -17,20 +17,20 @@
 namespace CVC4 {
 
 NodeDfsIterator::NodeDfsIterator(TNode n,
-                                 bool postorder,
+                                 VisitOrder order,
                                  std::function<bool(TNode)> skipIf)
     : d_stack{n},
       d_visited(),
-      d_postorder(postorder),
+      d_order(order),
       d_current(TNode()),
       d_skipIf(skipIf)
 {
 }
 
-NodeDfsIterator::NodeDfsIterator(bool postorder)
+NodeDfsIterator::NodeDfsIterator(VisitOrder order)
     : d_stack(),
       d_visited(),
-      d_postorder(postorder),
+      d_order(order),
       d_current(TNode()),
       d_skipIf([](TNode) { return false; })
 {
@@ -63,18 +63,21 @@ TNode& NodeDfsIterator::operator*()
   return d_current;
 }
 
-bool NodeDfsIterator::operator==(const NodeDfsIterator& other) const
+bool NodeDfsIterator::operator==(NodeDfsIterator& other)
 {
+  // Unitialize this node, and the other, before comparing.
+  initializeIfUninitialized();
+  other.initializeIfUninitialized();
   // 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, or
   // traversals with different skipIfs.
-  Assert(d_postorder == other.d_postorder);
+  Assert(d_order == other.d_order);
   return d_stack == other.d_stack && d_current == other.d_current;
 }
 
-bool NodeDfsIterator::operator!=(const NodeDfsIterator& other) const
+bool NodeDfsIterator::operator!=(NodeDfsIterator& other)
 {
   return !(*this == other);
 }
@@ -102,12 +105,12 @@ void NodeDfsIterator::advanceToNextVisit()
       {
         d_stack.push_back(back[i]);
       }
-      if (!d_postorder)
+      if (d_order == VisitOrder::PREORDER)
       {
         return;
       }
     }
-    else if (!d_postorder || visitEntry->second)
+    else if (d_order == VisitOrder::PREORDER || visitEntry->second)
     {
       // if we're previsiting or we've already post-visited this node: skip it
       d_stack.pop_back();
@@ -134,38 +137,21 @@ void NodeDfsIterator::initializeIfUninitialized()
   }
 }
 
-NodeDfsIterable::NodeDfsIterable(TNode n)
-    : d_node(n), d_postorder(true), d_skipIf([](TNode) { return false; })
-{
-}
-
-NodeDfsIterable& NodeDfsIterable::inPostorder()
-{
-  d_postorder = true;
-  return *this;
-}
-
-NodeDfsIterable& NodeDfsIterable::inPreorder()
-{
-  d_postorder = false;
-  return *this;
-}
-
-NodeDfsIterable& NodeDfsIterable::skipIf(
-    std::function<bool(TNode)> skipCondition)
+NodeDfsIterable::NodeDfsIterable(TNode n,
+                                 VisitOrder order,
+                                 std::function<bool(TNode)> skipIf)
+    : d_node(n), d_order(order), d_skipIf(skipIf)
 {
-  d_skipIf = skipCondition;
-  return *this;
 }
 
 NodeDfsIterator NodeDfsIterable::begin() const
 {
-  return NodeDfsIterator(d_node, d_postorder, d_skipIf);
+  return NodeDfsIterator(d_node, d_order, d_skipIf);
 }
 
 NodeDfsIterator NodeDfsIterable::end() const
 {
-  return NodeDfsIterator(d_postorder);
+  return NodeDfsIterator(d_order);
 }
 
 }  // namespace CVC4