diff options
Diffstat (limited to 'src/trie.cpp')
-rw-r--r-- | src/trie.cpp | 63 |
1 files changed, 55 insertions, 8 deletions
diff --git a/src/trie.cpp b/src/trie.cpp index 4198ff3..cd6cb7b 100644 --- a/src/trie.cpp +++ b/src/trie.cpp @@ -50,12 +50,12 @@ zmq::trie_t::~trie_t () } } -void zmq::trie_t::add (unsigned char *prefix_, size_t size_) +bool zmq::trie_t::add (unsigned char *prefix_, size_t size_) { // We are at the node corresponding to the prefix. We are done. if (!size_) { ++refcnt; - return; + return refcnt == 1; } unsigned char c = *prefix_; @@ -74,7 +74,7 @@ void zmq::trie_t::add (unsigned char *prefix_, size_t size_) count = (min < c ? c - min : min - c) + 1; next.table = (trie_t**) malloc (sizeof (trie_t*) * count); - alloc_assert (next.table); + zmq_assert (next.table); for (unsigned short i = 0; i != count; ++i) next.table [i] = 0; min = std::min (min, c); @@ -111,26 +111,28 @@ void zmq::trie_t::add (unsigned char *prefix_, size_t size_) if (count == 1) { if (!next.node) { next.node = new (std::nothrow) trie_t; - alloc_assert (next.node); + zmq_assert (next.node); } - next.node->add (prefix_ + 1, size_ - 1); + return next.node->add (prefix_ + 1, size_ - 1); } else { if (!next.table [c - min]) { next.table [c - min] = new (std::nothrow) trie_t; - alloc_assert (next.table [c - min]); + zmq_assert (next.table [c - min]); } - next.table [c - min]->add (prefix_ + 1, size_ - 1); + return next.table [c - min]->add (prefix_ + 1, size_ - 1); } } bool zmq::trie_t::rm (unsigned char *prefix_, size_t size_) { + // TODO: Shouldn't an error be reported if the key does not exist? + if (!size_) { if (!refcnt) return false; refcnt--; - return true; + return refcnt == 0; } unsigned char c = *prefix_; @@ -179,3 +181,48 @@ bool zmq::trie_t::check (unsigned char *data_, size_t size_) size_--; } } + +void zmq::trie_t::apply (void (*func_) (unsigned char *data_, size_t size_, + void *arg_), void *arg_) +{ + unsigned char *buff = NULL; + apply_helper (&buff, 0, 0, func_, arg_); + free (buff); +} + +void zmq::trie_t::apply_helper ( + unsigned char **buff_, size_t buffsize_, size_t maxbuffsize_, + void (*func_) (unsigned char *data_, size_t size_, void *arg_), void *arg_) +{ + // If this node is a subscription, apply the function. + if (refcnt) + func_ (*buff_, buffsize_, arg_); + + // Adjust the buffer. + if (buffsize_ >= maxbuffsize_) { + maxbuffsize_ = buffsize_ + 256; + *buff_ = (unsigned char*) realloc (*buff_, maxbuffsize_); + zmq_assert (*buff_); + } + + // If there are no subnodes in the trie, return. + if (count == 0) + return; + + // If there's one subnode (optimisation). + if (count == 1) { + (*buff_) [buffsize_] = min; + buffsize_++; + next.node->apply_helper (buff_, buffsize_, maxbuffsize_, func_, arg_); + return; + } + + // If there are multiple subnodes. + for (unsigned char c = 0; c != count; c++) { + (*buff_) [buffsize_] = min + c; + if (next.table [c]) + next.table [c]->apply_helper (buff_, buffsize_ + 1, maxbuffsize_, + func_, arg_); + } +} + |