diff --git a/src/tests/util/trie.cpp b/src/tests/util/trie.cpp index 671ea03290..6ee1b4c24e 100644 --- a/src/tests/util/trie.cpp +++ b/src/tests/util/trie.cpp @@ -27,6 +27,7 @@ static void tst1() { t2 = insert(t2, "help", 12); lean_assert(*find(t2, "abd") == 11); lean_assert(!find(t, "abd")); + t2.display(std::cout); ctrie t3 = *t2.find('a'); lean_assert(*find(t3, "bc") == 10); lean_assert(*find(t3, "bd") == 11); @@ -54,6 +55,8 @@ static void tst2() { lean_assert(*find(t1, "abd") == 12); lean_assert(!find(t2, "abc")); lean_assert(*find(t2, "abd") == 12); + std::cout << "---------\n"; + t1.display(std::cout); } int main() { diff --git a/src/util/trie.h b/src/util/trie.h index e6397e3217..9f843d6e74 100644 --- a/src/util/trie.h +++ b/src/util/trie.h @@ -9,6 +9,7 @@ Author: Leonardo de Moura #include #include "util/rb_map.h" #include "util/optional.h" +#include "util/buffer.h" namespace lean { template @@ -83,6 +84,18 @@ class trie : public KeyCMP { return new_t1; } + template + static void for_each(F && f, node const & n, buffer & prefix) { + if (n->m_value) { + f(prefix.size(), prefix.data(), *(n->m_value)); + } + n->m_children.for_each([&](Key const & k, node const & c) { + prefix.push_back(k); + for_each(f, c, prefix); + prefix.pop_back(); + }); + } + node m_root; trie(node const & n):m_root(n) {} public: @@ -120,6 +133,23 @@ public: void merge(trie const & t) { m_root = merge(m_root.steal(), t.m_root); } + + template + void for_each(F && f) const { + buffer prefix; + for_each(f, m_root, prefix); + } + + // for debugging purposes + void display(std::ostream & out) const { + for_each([&](unsigned num_keys, Key const * keys, Val const & v) { + for (unsigned i = 0; i < num_keys; i++) { + if (i > 0) out << ", "; + out << keys[i]; + } + out << " -> " << v << "\n"; + }); + } }; struct char_cmp { int operator()(char c1, char c2) const { return c1 < c2 ? -1 : (c1 == c2 ? 0 : 1); } };