trie.cpp 11 KB
Newer Older
1
/*
2
    Copyright (c) 2007-2016 Contributors as noted in the AUTHORS file
3

4
    This file is part of libzmq, the ZeroMQ core engine in C++.
5

6 7 8
    libzmq is free software; you can redistribute it and/or modify it under
    the terms of the GNU Lesser General Public License (LGPL) as published
    by the Free Software Foundation; either version 3 of the License, or
9 10
    (at your option) any later version.

11 12 13 14 15 16 17 18 19 20 21 22 23 24
    As a special exception, the Contributors give you permission to link
    this library with independent modules to produce an executable,
    regardless of the license terms of these independent modules, and to
    copy and distribute the resulting executable under terms of your choice,
    provided that you also meet, for each linked independent module, the
    terms and conditions of the license of that module. An independent
    module is a module which is not derived from or based on this library.
    If you modify this library, you must extend this exception to your
    version of the library.

    libzmq is distributed in the hope that it will be useful, but WITHOUT
    ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
    FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
    License for more details.
25

26
    You should have received a copy of the GNU Lesser General Public License
27 28 29
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

30
#include "precompiled.hpp"
31 32 33 34
#include "macros.hpp"
#include "err.hpp"
#include "trie.hpp"

35 36 37 38 39
#include <stdlib.h>

#include <new>
#include <algorithm>

40
zmq::trie_t::trie_t () : refcnt (0), min (0), count (0), live_nodes (0)
41 42 43
{
}

44
zmq::trie_t::~trie_t ()
45
{
46 47
    if (count == 1) {
        zmq_assert (next.node);
48 49
        LIBZMQ_DELETE (next.node);
    } else if (count > 1) {
50
        for (unsigned short i = 0; i != count; ++i) {
51
            LIBZMQ_DELETE (next.table[i]);
52
        }
53 54 55 56
        free (next.table);
    }
}

57
bool zmq::trie_t::add (unsigned char *prefix_, size_t size_)
58 59 60 61
{
    //  We are at the node corresponding to the prefix. We are done.
    if (!size_) {
        ++refcnt;
62
        return refcnt == 1;
63 64 65 66 67
    }

    unsigned char c = *prefix_;
    if (c < min || c >= min + count) {
        //  The character is out of range of currently handled
68
        //  characters. We have to extend the table.
69 70 71 72
        if (!count) {
            min = c;
            count = 1;
            next.node = NULL;
73
        } else if (count == 1) {
74
            unsigned char oldc = min;
75
            trie_t *oldp = next.node;
76
            count = (min < c ? c - min : min - c) + 1;
77 78
            next.table =
              static_cast<trie_t **> (malloc (sizeof (trie_t *) * count));
79
            alloc_assert (next.table);
80
            for (unsigned short i = 0; i != count; ++i)
81
                next.table[i] = 0;
82
            min = std::min (min, c);
83 84
            next.table[oldc - min] = oldp;
        } else if (min < c) {
85
            //  The new character is above the current character range.
86
            unsigned short old_count = count;
87
            count = c - min + 1;
88 89
            next.table = static_cast<trie_t **> (
              realloc ((void *) next.table, sizeof (trie_t *) * count));
90
            zmq_assert (next.table);
91
            for (unsigned short i = old_count; i != count; i++)
92 93
                next.table[i] = NULL;
        } else {
94
            //  The new character is below the current character range.
95
            unsigned short old_count = count;
96
            count = (min + old_count) - c;
97 98
            next.table = static_cast<trie_t **> (
              realloc ((void *) next.table, sizeof (trie_t *) * count));
99 100
            zmq_assert (next.table);
            memmove (next.table + min - c, next.table,
101
                     old_count * sizeof (trie_t *));
102
            for (unsigned short i = 0; i != min - c; i++)
103
                next.table[i] = NULL;
104 105 106 107 108 109 110
            min = c;
        }
    }

    //  If next node does not exist, create one.
    if (count == 1) {
        if (!next.node) {
111
            next.node = new (std::nothrow) trie_t;
112
            alloc_assert (next.node);
113
            ++live_nodes;
114
            zmq_assert (live_nodes == 1);
115
        }
116
        return next.node->add (prefix_ + 1, size_ - 1);
117 118 119 120
    } else {
        if (!next.table[c - min]) {
            next.table[c - min] = new (std::nothrow) trie_t;
            alloc_assert (next.table[c - min]);
121
            ++live_nodes;
122
            zmq_assert (live_nodes > 1);
123
        }
124
        return next.table[c - min]->add (prefix_ + 1, size_ - 1);
125 126 127
    }
}

128
bool zmq::trie_t::rm (unsigned char *prefix_, size_t size_)
129
{
130 131 132 133 134 135 136 137 138 139
    //  TODO: Shouldn't an error be reported if the key does not exist?
    if (!size_) {
        if (!refcnt)
            return false;
        refcnt--;
        return refcnt == 0;
    }
    unsigned char c = *prefix_;
    if (!count || c < min || c >= min + count)
        return false;
140

141
    trie_t *next_node = count == 1 ? next.node : next.table[c - min];
142

143 144
    if (!next_node)
        return false;
145

146
    bool ret = next_node->rm (prefix_ + 1, size_ - 1);
147

148 149
    //  Prune redundant nodes
    if (next_node->is_redundant ()) {
150
        LIBZMQ_DELETE (next_node);
151
        zmq_assert (count > 0);
152

153 154 155 156 157 158
        if (count == 1) {
            //  The just pruned node is was the only live node
            next.node = 0;
            count = 0;
            --live_nodes;
            zmq_assert (live_nodes == 0);
159 160
        } else {
            next.table[c - min] = 0;
161 162 163 164 165 166 167 168 169 170 171 172 173 174
            zmq_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
175
                    node = next.table[count - 1];
176
                    min += count - 1;
177
                } else if (c == min + count - 1) {
178 179
                    //  The pruned node is the right-most node ptr in the
                    //  node table => keep the left-most node
180
                    node = next.table[0];
181 182 183 184 185
                }
                zmq_assert (node);
                free (next.table);
                next.node = node;
                count = 1;
186
            } else if (c == min) {
187 188 189 190 191
                //  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) {
192
                    if (next.table[i]) {
193 194 195 196 197 198 199 200 201 202 203
                        new_min = i + min;
                        break;
                    }
                }
                zmq_assert (new_min != min);

                trie_t **old_table = next.table;
                zmq_assert (new_min > min);
                zmq_assert (count > new_min - min);

                count = count - (new_min - min);
204 205
                next.table =
                  static_cast<trie_t **> (malloc (sizeof (trie_t *) * count));
206 207 208
                alloc_assert (next.table);

                memmove (next.table, old_table + (new_min - min),
209
                         sizeof (trie_t *) * count);
210 211 212
                free (old_table);

                min = new_min;
213
            } else if (c == min + count - 1) {
214 215 216 217 218
                //  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) {
219
                    if (next.table[count - 1 - i]) {
220 221 222 223 224 225 226 227
                        new_count = count - i;
                        break;
                    }
                }
                zmq_assert (new_count != count);
                count = new_count;

                trie_t **old_table = next.table;
228 229
                next.table =
                  static_cast<trie_t **> (malloc (sizeof (trie_t *) * count));
230 231
                alloc_assert (next.table);

232
                memmove (next.table, old_table, sizeof (trie_t *) * count);
233 234 235 236 237
                free (old_table);
            }
        }
    }
    return ret;
238 239
}

240
bool zmq::trie_t::check (unsigned char *data_, size_t size_)
241 242 243
{
    //  This function is on critical path. It deliberately doesn't use
    //  recursion to get a bit better performance.
244
    trie_t *current = this;
245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263
    while (true) {
        //  We've found a corresponding subscription!
        if (current->refcnt)
            return true;

        //  We've checked all the data and haven't found matching subscription.
        if (!size_)
            return false;

        //  If there's no corresponding slot for the first character
        //  of the prefix, the message does not match.
        unsigned char c = *data_;
        if (c < current->min || c >= current->min + current->count)
            return false;

        //  Move to the next character.
        if (current->count == 1)
            current = current->next.node;
        else {
264
            current = current->next.table[c - current->min];
265 266 267 268 269 270 271
            if (!current)
                return false;
        }
        data_++;
        size_--;
    }
}
272

273 274
void zmq::trie_t::apply (
  void (*func_) (unsigned char *data_, size_t size_, void *arg_), void *arg_)
275 276 277 278 279 280
{
    unsigned char *buff = NULL;
    apply_helper (&buff, 0, 0, func_, arg_);
    free (buff);
}

281 282 283 284 285 286 287
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_)
288 289 290 291 292 293 294 295
{
    //  If this node is a subscription, apply the function.
    if (refcnt)
        func_ (*buff_, buffsize_, arg_);

    //  Adjust the buffer.
    if (buffsize_ >= maxbuffsize_) {
        maxbuffsize_ = buffsize_ + 256;
296
        *buff_ = static_cast<unsigned char *> (realloc (*buff_, maxbuffsize_));
297 298 299 300 301 302 303 304 305
        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) {
306
        (*buff_)[buffsize_] = min;
307 308 309 310 311 312
        buffsize_++;
        next.node->apply_helper (buff_, buffsize_, maxbuffsize_, func_, arg_);
        return;
    }

    //  If there are multiple subnodes.
313
    for (unsigned short c = 0; c != count; c++) {
314 315 316 317
        (*buff_)[buffsize_] = min + c;
        if (next.table[c])
            next.table[c]->apply_helper (buff_, buffsize_ + 1, maxbuffsize_,
                                         func_, arg_);
318
    }
319 320
}

321
bool zmq::trie_t::is_redundant () const
322 323 324
{
    return refcnt == 0 && live_nodes == 0;
}