Merge pull request #222 from gimaker/trie-pruning

Fix for LIBZMQ-305

src/mtrie.cpp

 /*
     Copyright (c) 2011 250bpm s.r.o.
+    Copyright (c) 2011-2012 Spotify AB
     Copyright (c) 2011 Other contributors as noted in the AUTHORS file
 
     This file is part of 0MQ.
@@ -34,14 +35,18 @@
 
 zmq::mtrie_t::mtrie_t () :
     min (0),
-    count (0)
+    count (0),
+    live_nodes (0)
 {
 }
 
 zmq::mtrie_t::~mtrie_t ()
 {
-    if (count == 1)
+    if (count == 1) {
+        zmq_assert (next.node);
         delete next.node;
+        next.node = 0;
+    }
     else if (count > 1) {
         for (unsigned short i = 0; i != count; ++i)
             if (next.table [i])
@@ -118,6 +123,7 @@ bool zmq::mtrie_t::add_helper (unsigned char *prefix_, size_t size_,
     if (count == 1) {
         if (!next.node) {
             next.node = new (std::nothrow) mtrie_t;
+            ++live_nodes;
             zmq_assert (next.node);
         }
         return next.node->add_helper (prefix_ + 1, size_ - 1, pipe_);
@@ -125,6 +131,7 @@ bool zmq::mtrie_t::add_helper (unsigned char *prefix_, size_t size_,
     else {
         if (!next.table [c - min]) {
             next.table [c - min] = new (std::nothrow) mtrie_t;
+            ++live_nodes;
             zmq_assert (next.table [c - min]);
         }
         return next.table [c - min]->add_helper (prefix_ + 1, size_ - 1, pipe_);
@@ -167,16 +174,28 @@ void zmq::mtrie_t::rm_helper (pipe_t *pipe_, unsigned char **buff_,
         buffsize_++;
         next.node->rm_helper (pipe_, buff_, buffsize_, maxbuffsize_,
             func_, arg_);
+        if (next.node->is_redundant ()) {
+            delete next.node;
+            next.node = 0;
+            count = 0;
+            --live_nodes;
+        }
         return;
     }
 
     //  If there are multiple subnodes.
     for (unsigned short c = 0; c != count; c++) {
         (*buff_) [buffsize_] = min + c;
-        if (next.table [c])
+        if (next.table [c]) {
             next.table [c]->rm_helper (pipe_, buff_, buffsize_ + 1,
                 maxbuffsize_, func_, arg_);
-    }  
+            if (next.table [c]->is_redundant ()) {
+                delete next.table [c];
+                next.table [c] = 0;
+                --live_nodes;
+            }
+        }
+    }
 }
 
 bool zmq::mtrie_t::rm (unsigned char *prefix_, size_t size_, pipe_t *pipe_)
@@ -203,7 +222,20 @@ bool zmq::mtrie_t::rm_helper (unsigned char *prefix_, size_t size_,
     if (!next_node)
         return false;
 
-    return next_node->rm_helper (prefix_ + 1, size_ - 1, pipe_);
+    bool ret = next_node->rm_helper (prefix_ + 1, size_ - 1, pipe_);
+
+    if (next_node->is_redundant ()) {
+        delete next_node;
+        if (count == 1) {
+            next.node = 0;
+            count = 0;
+        }
+        else
+            next.table [c - min] = 0;
+        --live_nodes;
+    }
+
+    return ret;
 }
 
 void zmq::mtrie_t::match (unsigned char *data_, size_t size_,
@@ -247,3 +279,7 @@ void zmq::mtrie_t::match (unsigned char *data_, size_t size_,
     }
 }
 
+bool zmq::mtrie_t::is_redundant () const
+{
+    return pipes.empty () && live_nodes == 0;
+}

src/mtrie.hpp

 /*
     Copyright (c) 2011 250bpm s.r.o.
+    Copyright (c) 2011-2012 Spotify AB
     Copyright (c) 2011 Other contributors as noted in the AUTHORS file
 
     This file is part of 0MQ.
@@ -69,12 +70,14 @@ namespace zmq
             void *arg_);
         bool rm_helper (unsigned char *prefix_, size_t size_,
             zmq::pipe_t *pipe_);
+        bool is_redundant () const;
 
         typedef std::set <zmq::pipe_t*> pipes_t;
         pipes_t pipes;
 
         unsigned char min;
         unsigned short count;
+        unsigned short live_nodes;
         union {
             class mtrie_t *node;
             class mtrie_t **table;

src/trie.cpp

 /*
     Copyright (c) 2009-2011 250bpm s.r.o.
     Copyright (c) 2007-2009 iMatix Corporation
+    Copyright (c) 2011-2012 Spotify AB
     Copyright (c) 2007-2011 Other contributors as noted in the AUTHORS file
 
     This file is part of 0MQ.
@@ -35,14 +36,18 @@
 zmq::trie_t::trie_t () :
     refcnt (0),
     min (0),
-    count (0)
+    count (0),
+    live_nodes (0)
 {
 }
 
 zmq::trie_t::~trie_t ()
 {
-    if (count == 1)
+    if (count == 1) {
+        zmq_assert (next.node);
         delete next.node;
+        next.node = 0;
+    }
     else if (count > 1) {
         for (unsigned short i = 0; i != count; ++i)
             if (next.table [i])
@@ -113,6 +118,7 @@ bool zmq::trie_t::add (unsigned char *prefix_, size_t size_)
         if (!next.node) {
             next.node = new (std::nothrow) trie_t;
             zmq_assert (next.node);
+            ++live_nodes;
         }
         return next.node->add (prefix_ + 1, size_ - 1);
     }
@@ -120,6 +126,7 @@ bool zmq::trie_t::add (unsigned char *prefix_, size_t size_)
         if (!next.table [c - min]) {
             next.table [c - min] = new (std::nothrow) trie_t;
             zmq_assert (next.table [c - min]);
+            ++live_nodes;
         }
         return next.table [c - min]->add (prefix_ + 1, size_ - 1);
     }
@@ -146,7 +153,20 @@ bool zmq::trie_t::rm (unsigned char *prefix_, size_t size_)
      if (!next_node)
          return false;
 
-     return next_node->rm (prefix_ + 1, size_ - 1);
+     bool ret = next_node->rm (prefix_ + 1, size_ - 1);
+
+     if (next_node->is_redundant ()) {
+         delete next_node;
+         if (count == 1) {
+             next.node = 0;
+             count = 0;
+         }
+         else
+             next.table [c - min] = 0;
+         --live_nodes;
+     }
+
+     return ret;
 }
 
 bool zmq::trie_t::check (unsigned char *data_, size_t size_)
@@ -224,6 +244,10 @@ void zmq::trie_t::apply_helper (
         if (next.table [c])
             next.table [c]->apply_helper (buff_, buffsize_ + 1, maxbuffsize_,
                 func_, arg_);
-    }   
+    }
 }
 
+bool zmq::trie_t::is_redundant () const
+{
+    return refcnt == 0 && live_nodes == 0;
+}

src/trie.hpp

 /*
     Copyright (c) 2009-2011 250bpm s.r.o.
     Copyright (c) 2007-2009 iMatix Corporation
+    Copyright (c) 2011-2012 Spotify AB
     Copyright (c) 2007-2011 Other contributors as noted in the AUTHORS file
 
     This file is part of 0MQ.
@@ -57,10 +58,12 @@ namespace zmq
             unsigned char **buff_, size_t buffsize_, size_t maxbuffsize_,
             void (*func_) (unsigned char *data_, size_t size_, void *arg_),
             void *arg_);
+        bool is_redundant () const;
 
         uint32_t refcnt;
         unsigned char min;
         unsigned short count;
+        unsigned short live_nodes;
         union {
             class trie_t *node;
             class trie_t **table;