Add skip predicate to node traversal. (#4222)

Sometime you want to skip specific sub-DAGs when traversing a node. For example, you might be doing a transformation with a cache, and want to skip sub-DAGs that you've already processed.

This PR would add a skipIf builder method to NodeDfsIterable, which allows the user to provide a predicate according to which nodes will be omitted from the subsequent traversal.

1 files changed, 32 insertions, 0 deletions

diff --git a/test/unit/expr/node_traversal_black.h b/test/unit/expr/node_traversal_black.h
index b4a7c449c..b751a0999 100644
--- a/test/unit/expr/node_traversal_black.h
+++ b/test/unit/expr/node_traversal_black.h
@@ -161,6 +161,22 @@ class NodePostorderTraversalBlack : public CxxTest::TestSuite
     std::copy(traversal.begin(), traversal.end(), std::back_inserter(actual));
     TS_ASSERT_EQUALS(actual, expected);
   }
+
+  void testSkipIf()
+  {
+    Node tb = d_nodeManager->mkConst(true);
+    Node eb = d_nodeManager->mkConst(false);
+    Node cnd = d_nodeManager->mkNode(XOR, tb, eb);
+    Node top = d_nodeManager->mkNode(XOR, cnd, cnd);
+    std::vector<TNode> expected = {top};
+
+    auto traversal = NodeDfsIterable(top).inPostorder().skipIf(
+        [&cnd](TNode n) { return n == cnd; });
+
+    std::vector<TNode> actual;
+    std::copy(traversal.begin(), traversal.end(), std::back_inserter(actual));
+    TS_ASSERT_EQUALS(actual, expected);
+  }
 };
 
 class NodePreorderTraversalBlack : public CxxTest::TestSuite
@@ -278,4 +294,20 @@ class NodePreorderTraversalBlack : public CxxTest::TestSuite
     std::copy(traversal.begin(), traversal.end(), std::back_inserter(actual));
     TS_ASSERT_EQUALS(actual, expected);
   }
+
+  void testSkipIf()
+  {
+    Node tb = d_nodeManager->mkConst(true);
+    Node eb = d_nodeManager->mkConst(false);
+    Node cnd = d_nodeManager->mkNode(XOR, tb, eb);
+    Node top = d_nodeManager->mkNode(XOR, cnd, cnd);
+    std::vector<TNode> expected = {top, cnd, eb};
+
+    auto traversal = NodeDfsIterable(top).inPreorder().skipIf(
+        [&tb](TNode n) { return n == tb; });
+
+    std::vector<TNode> actual;
+    std::copy(traversal.begin(), traversal.end(), std::back_inserter(actual));
+    TS_ASSERT_EQUALS(actual, expected);
+  }
 };