From fad0df02d9d56c717c37269b50497ec5c58749f3 Mon Sep 17 00:00:00 2001
From: Andrew Reynolds <andrew.j.reynolds@gmail.com>
Date: Thu, 7 Apr 2022 16:04:31 -0500
Subject: [PATCH] Add method to get leaves of a NodeTrie (#8583)

From the oracles branch.
---
 src/expr/node_trie.cpp | 33 +++++++++++++++++++++++++++++++++
 src/expr/node_trie.h   | 14 ++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/src/expr/node_trie.cpp b/src/expr/node_trie.cpp
index 80c06d4c5..18215c2b3 100644
--- 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 56db7d477..c61a0dfd4 100644
--- 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>
-- 
2.30.2