Patches from sub-forward branch incorporated

Signed-off-by: Martin Sustrik <sustrik@250bpm.com>

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_);
+    }   
+}
+