Add method to get leaves of a NodeTrie (#8583)

From the oracles branch.

diff --git a/src/expr/node_trie.cpp b/src/expr/node_trie.cpp
index 80c06d4c55623300cf9e6fd89fc34584a1f3b8b0..18215c2b3eb50392cdc745200bbec80d228df4bf 100644 (file)
--- a/src/expr/node_trie.cpp
+++ b/src/expr/node_trie.cpp
@@ -91,4 +91,37 @@ template void NodeTemplateTrie<false>::debugPrint(const char* c,
 template void NodeTemplateTrie<true>::debugPrint(const char* c,
                                                  unsigned depth) const;
 
+
+template <bool ref_count>
+std::vector<Node> NodeTemplateTrie<ref_count>::getLeaves(size_t depth) const
+{
+  Assert(depth > 0);
+  std::vector<Node> vec;
+  std::vector<std::pair<const NodeTemplateTrie<ref_count>*, size_t>> visit;
+  visit.emplace_back(this, depth);
+  do
+  {
+    std::pair<const NodeTemplateTrie<ref_count>*, size_t> curr = visit.back();
+    visit.pop_back();
+    size_t currDepth = curr.second;
+    for (const std::pair<const NodeTemplate<ref_count>,
+                         NodeTemplateTrie<ref_count>>& p : curr.first->d_data)
+    {
+      if (currDepth == 0)
+      {
+        // we are at a leaf
+        vec.push_back(p.first);
+        break;
+      }
+      visit.emplace_back(&p.second, currDepth - 1);
+    }
+  } while (!visit.empty());
+  return vec;
+}
+
+template std::vector<Node> NodeTemplateTrie<false>::getLeaves(
+    size_t depth) const;
+template std::vector<Node> NodeTemplateTrie<true>::getLeaves(
+    size_t depth) const;
+
 }  // namespace cvc5::internal

diff --git a/src/expr/node_trie.h b/src/expr/node_trie.h
index 56db7d477b38f022c809e367b6b85d34cdd139fc..c61a0dfd4a7e3be5fa7b6b5f3093abbe729958f8 100644 (file)
--- a/src/expr/node_trie.h
+++ b/src/expr/node_trie.h
@@ -92,6 +92,20 @@ class NodeTemplateTrie
   void clear() { d_data.clear(); }
   /** Is this trie empty? */
   bool empty() const { return d_data.empty(); }
+  /**
+   * Get leaves at the given depth, where depth>0. This argument is necessary
+   * since we do not know apriori the depth of where data occurs.
+   *
+   * If this trie stores applications of a function f, then depth should be set
+   * to the arity of f.
+   *
+   * Notice this method will never throw an assertion error, even if the
+   * depth is not set to the proper value. In particular, it will return
+   * the empty vector if the provided depth is larger than the actual depth,
+   * and will return internal nodes if the provided depth is less than the
+   * actual depth of the trie.
+   */
+  std::vector<Node> getLeaves(size_t depth) const;
 }; /* class NodeTemplateTrie */
 
 template <bool ref_count>