Compact the trie/mtrie node tables where possible, to reduce memory usage.

Signed-off-by: Staffan Gimåker <staffan@spotify.com>

1 files changed, 87 insertions, 4 deletions

diff --git a/src/trie.cpp b/src/trie.cpp
index ea34787..7bb13e3 100644
--- a/src/trie.cpp
+++ b/src/trie.cpp
@@ -117,16 +117,18 @@ bool xs::trie_t::add (unsigned char *prefix_, size_t size_)
     if (count == 1) {
         if (!next.node) {
             next.node = new (std::nothrow) trie_t;
-            ++live_nodes;
             xs_assert (next.node);
+            ++live_nodes;
+            xs_assert (live_nodes == 1);
         }
         return next.node->add (prefix_ + 1, size_ - 1);
     }
     else {
         if (!next.table [c - min]) {
             next.table [c - min] = new (std::nothrow) trie_t;
-            ++live_nodes;
             xs_assert (next.table [c - min]);
+            ++live_nodes;
+            xs_assert (live_nodes > 1);
         }
         return next.table [c - min]->add (prefix_ + 1, size_ - 1);
     }
@@ -155,15 +157,96 @@ bool xs::trie_t::rm (unsigned char *prefix_, size_t size_)
 
      bool ret = next_node->rm (prefix_ + 1, size_ - 1);
 
+     //  Prune redundant nodes
      if (next_node->is_redundant ()) {
          delete next_node;
+         xs_assert (count > 0);
+
          if (count == 1) {
+             //  The just pruned node is was the only live node
              next.node = 0;
              count = 0;
+             --live_nodes;
+             xs_assert (live_nodes == 0);
          }
-         else
+         else {
              next.table [c - min] = 0;
-         --live_nodes;
+             xs_assert (live_nodes > 1);
+             --live_nodes;
+
+             //  Compact the table if possible
+             if (live_nodes == 1) {
+                 //  We can switch to using the more compact single-node
+                 //  representation since the table only contains one live node
+                 trie_t *node = 0;
+                 //  Since we always compact the table the pruned node must
+                 //  either be the left-most or right-most ptr in the node
+                 //  table
+                 if (c == min) {
+                     //  The pruned node is the left-most node ptr in the
+                     //  node table => keep the right-most node
+                     node = next.table [count - 1];
+                 }
+                 else if (c == min + count - 1) {
+                     //  The pruned node is the right-most node ptr in the
+                     //  node table => keep the left-most node
+                     node = next.table [0];
+                 }
+
+                 xs_assert (node);
+                 free (next.table);
+                 next.node = node;
+                 count = 1;
+             }
+             else if (c == min) {
+                 //  We can compact the table "from the left".
+                 //  Find the left-most non-null node ptr, which we'll use as
+                 //  our new min
+                 unsigned char new_min = min;
+                 for (unsigned short i = 1; i < count; ++i) {
+                     if (next.table [i]) {
+                         new_min = i + min;
+                         break;
+                     }
+                 }
+                 xs_assert (new_min != min);
+
+                 trie_t **old_table = next.table;
+                 xs_assert (new_min > min);
+                 xs_assert (count > new_min - min);
+
+                 count = count - (new_min - min);
+                 next.table = (trie_t**) malloc (sizeof (trie_t*) * count);
+                 xs_assert (next.table);
+
+                 memmove (next.table, old_table + (new_min - min),
+                          sizeof (trie_t*) * count);
+                 free (old_table);
+
+                 min = new_min;
+             }
+             else if (c == min + count - 1) {
+                 //  We can compact the table "from the right".
+                 //  Find the right-most non-null node ptr, which we'll use to
+                 //  determine the new table size
+                 unsigned short new_count = count;
+                 for (unsigned short i = 1; i < count; ++i) {
+                     if (next.table [count - 1 - i]) {
+                         new_count = count - i;
+                         break;
+                     }
+                 }
+                 xs_assert (new_count != count);
+                 count = new_count;
+
+                 trie_t **old_table = next.table;
+                 next.table = (trie_t**) malloc (sizeof (trie_t*) * count);
+                 xs_assert (next.table);
+
+                 memmove (next.table, old_table, sizeof (trie_t*) * count);
+                 free (old_table);
+             }
+         }
      }
 
      return ret;