trie.cpp 11.2 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 48 49 50 51
    if (_count == 1) {
        zmq_assert (_next.node);
        LIBZMQ_DELETE (_next.node);
    } else if (_count > 1) {
        for (unsigned short i = 0; i != _count; ++i) {
            LIBZMQ_DELETE (_next.table[i]);
52
        }
53
        free (_next.table);
54 55 56
    }
}

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

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

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

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

140
    trie_t *next_node = _count == 1 ? _next.node : _next.table[c - _min];
141

142 143
    if (!next_node)
        return false;
144

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

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

152
        if (_count == 1) {
153
            //  The just pruned node is was the only live node
154 155 156 157
            _next.node = 0;
            _count = 0;
            --_live_nodes;
            zmq_assert (_live_nodes == 0);
158
        } else {
159 160 161
            _next.table[c - _min] = 0;
            zmq_assert (_live_nodes > 1);
            --_live_nodes;
162 163

            //  Compact the table if possible
164
            if (_live_nodes == 1) {
165 166 167 168 169 170
                //  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
171
                if (c == _min) {
172 173
                    //  The pruned node is the left-most node ptr in the
                    //  node table => keep the right-most node
174 175 176
                    node = _next.table[_count - 1];
                    _min += _count - 1;
                } else if (c == _min + _count - 1) {
177 178
                    //  The pruned node is the right-most node ptr in the
                    //  node table => keep the left-most node
179
                    node = _next.table[0];
180 181
                }
                zmq_assert (node);
182 183 184 185
                free (_next.table);
                _next.node = node;
                _count = 1;
            } else if (c == _min) {
186 187 188
                //  We can compact the table "from the left".
                //  Find the left-most non-null node ptr, which we'll use as
                //  our new min
189 190 191 192
                unsigned char new_min = _min;
                for (unsigned short i = 1; i < _count; ++i) {
                    if (_next.table[i]) {
                        new_min = i + _min;
193 194 195
                        break;
                    }
                }
196
                zmq_assert (new_min != _min);
197

198 199 200
                trie_t **old_table = _next.table;
                zmq_assert (new_min > _min);
                zmq_assert (_count > new_min - _min);
201

202 203 204 205
                _count = _count - (new_min - _min);
                _next.table =
                  static_cast<trie_t **> (malloc (sizeof (trie_t *) * _count));
                alloc_assert (_next.table);
206

207 208
                memmove (_next.table, old_table + (new_min - _min),
                         sizeof (trie_t *) * _count);
209 210
                free (old_table);

211 212
                _min = new_min;
            } else if (c == _min + _count - 1) {
213 214 215
                //  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
216 217 218 219
                unsigned short new_count = _count;
                for (unsigned short i = 1; i < _count; ++i) {
                    if (_next.table[_count - 1 - i]) {
                        new_count = _count - i;
220 221 222
                        break;
                    }
                }
223 224
                zmq_assert (new_count != _count);
                _count = new_count;
225

226 227 228 229
                trie_t **old_table = _next.table;
                _next.table =
                  static_cast<trie_t **> (malloc (sizeof (trie_t *) * _count));
                alloc_assert (_next.table);
230

231
                memmove (_next.table, old_table, sizeof (trie_t *) * _count);
232 233 234 235 236
                free (old_table);
            }
        }
    }
    return ret;
237 238
}

239
bool zmq::trie_t::check (unsigned char *data_, size_t size_)
240 241 242
{
    //  This function is on critical path. It deliberately doesn't use
    //  recursion to get a bit better performance.
243
    trie_t *current = this;
244 245
    while (true) {
        //  We've found a corresponding subscription!
246
        if (current->_refcnt)
247 248 249 250 251 252 253 254
            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.
255
        const unsigned char c = *data_;
256
        if (c < current->_min || c >= current->_min + current->_count)
257 258 259
            return false;

        //  Move to the next character.
260 261
        if (current->_count == 1)
            current = current->_next.node;
262
        else {
263
            current = current->_next.table[c - current->_min];
264 265 266 267 268 269 270
            if (!current)
                return false;
        }
        data_++;
        size_--;
    }
}
271

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

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

    //  Adjust the buffer.
    if (buffsize_ >= maxbuffsize_) {
        maxbuffsize_ = buffsize_ + 256;
295
        *buff_ = static_cast<unsigned char *> (realloc (*buff_, maxbuffsize_));
296 297 298 299
        zmq_assert (*buff_);
    }

    //  If there are no subnodes in the trie, return.
300
    if (_count == 0)
301 302 303
        return;

    //  If there's one subnode (optimisation).
304 305
    if (_count == 1) {
        (*buff_)[buffsize_] = _min;
306
        buffsize_++;
307
        _next.node->apply_helper (buff_, buffsize_, maxbuffsize_, func_, arg_);
308 309 310 311
        return;
    }

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

320
bool zmq::trie_t::is_redundant () const
321
{
322
    return _refcnt == 0 && _live_nodes == 0;
323
}