radix_tree.cpp 20.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36
/*
    Copyright (c) 2018 Contributors as noted in the AUTHORS file

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

    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
    (at your option) any later version.

    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.

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

#include "precompiled.hpp"
#include "macros.hpp"
#include "err.hpp"
#include "radix_tree.hpp"

#include <stdlib.h>
#include <string.h>
37
#include <vector>
38

39
node_t::node_t (unsigned char *data_) : _data (data_)
40 41 42
{
}

43
uint32_t node_t::refcount ()
44 45
{
    uint32_t u32;
46
    memcpy (&u32, _data, sizeof (u32));
47 48 49
    return u32;
}

50
void node_t::set_refcount (uint32_t value_)
51
{
52
    memcpy (_data, &value_, sizeof (value_));
53 54
}

55
uint32_t node_t::prefix_length ()
56 57
{
    uint32_t u32;
58
    memcpy (&u32, _data + sizeof (uint32_t), sizeof (u32));
59 60 61
    return u32;
}

62
void node_t::set_prefix_length (uint32_t value_)
63
{
64
    memcpy (_data + sizeof (value_), &value_, sizeof (value_));
65 66
}

67
uint32_t node_t::edgecount ()
68 69
{
    uint32_t u32;
70
    memcpy (&u32, _data + 2 * sizeof (uint32_t), sizeof (u32));
71 72 73
    return u32;
}

74
void node_t::set_edgecount (uint32_t value_)
75
{
76
    memcpy (_data + 2 * sizeof (value_), &value_, sizeof (value_));
77 78
}

79
unsigned char *node_t::prefix ()
80
{
81
    return _data + 3 * sizeof (uint32_t);
82 83
}

84
void node_t::set_prefix (const unsigned char *bytes_)
85
{
86
    memcpy (prefix (), bytes_, prefix_length ());
87 88
}

89
unsigned char *node_t::first_bytes ()
90 91 92 93
{
    return prefix () + prefix_length ();
}

94
void node_t::set_first_bytes (const unsigned char *bytes_)
95
{
96
    memcpy (first_bytes (), bytes_, edgecount ());
97 98
}

99
unsigned char node_t::first_byte_at (size_t index_)
100
{
101 102
    zmq_assert (index_ < edgecount ());
    return first_bytes ()[index_];
103 104
}

105
void node_t::set_first_byte_at (size_t index_, unsigned char byte_)
106
{
107 108
    zmq_assert (index_ < edgecount ());
    first_bytes ()[index_] = byte_;
109 110
}

111
unsigned char *node_t::node_pointers ()
112 113 114 115
{
    return prefix () + prefix_length () + edgecount ();
}

116
void node_t::set_node_pointers (const unsigned char *pointers_)
117
{
118
    memcpy (node_pointers (), pointers_, edgecount () * sizeof (void *));
119 120
}

121
node_t node_t::node_at (size_t index_)
122
{
123
    zmq_assert (index_ < edgecount ());
124 125

    unsigned char *data;
126
    memcpy (&data, node_pointers () + index_ * sizeof (void *), sizeof (data));
127
    return node_t (data);
128 129
}

130
void node_t::set_node_at (size_t index_, node_t node_)
131
{
132 133 134
    zmq_assert (index_ < edgecount ());
    memcpy (node_pointers () + index_ * sizeof (void *), &node_._data,
            sizeof (node_._data));
135 136
}

137 138 139
void node_t::set_edge_at (size_t index_,
                          unsigned char first_byte_,
                          node_t node_)
140
{
141 142
    set_first_byte_at (index_, first_byte_);
    set_node_at (index_, node_);
143 144
}

145
bool node_t::operator== (node_t other_) const
146
{
147
    return _data == other_._data;
148 149
}

150
bool node_t::operator!= (node_t other_) const
151
{
152
    return !(*this == other_);
153 154
}

155
void node_t::resize (size_t prefix_length_, size_t edgecount_)
156
{
157 158
    size_t node_size = 3 * sizeof (uint32_t) + prefix_length_
                       + edgecount_ * (1 + sizeof (void *));
159
    unsigned char *new_data =
160
      static_cast<unsigned char *> (realloc (_data, node_size));
161
    zmq_assert (new_data);
162 163 164
    _data = new_data;
    set_prefix_length (static_cast<uint32_t> (prefix_length_));
    set_edgecount (static_cast<uint32_t> (edgecount_));
165 166
}

167
node_t make_node (size_t refcount_, size_t prefix_length_, size_t edgecount_)
168
{
169 170
    size_t node_size = 3 * sizeof (uint32_t) + prefix_length_
                       + edgecount_ * (1 + sizeof (void *));
171

172
    unsigned char *data = static_cast<unsigned char *> (malloc (node_size));
173 174
    zmq_assert (data);

175
    node_t node (data);
176 177 178
    node.set_refcount (static_cast<uint32_t> (refcount_));
    node.set_prefix_length (static_cast<uint32_t> (prefix_length_));
    node.set_edgecount (static_cast<uint32_t> (edgecount_));
179
    return node;
180 181 182 183
}

// ----------------------------------------------------------------------

184
zmq::radix_tree_t::radix_tree_t () : _root (make_node (0, 0, 0)), _size (0)
185 186 187
{
}

188
static void free_nodes (node_t node_)
189
{
190 191 192
    for (size_t i = 0; i < node_.edgecount (); ++i)
        free_nodes (node_.node_at (i));
    free (node_._data);
193 194
}

195
zmq::radix_tree_t::~radix_tree_t ()
196
{
197
    free_nodes (_root);
198 199
}

200 201 202 203 204 205 206 207 208 209 210 211 212 213
match_result_t::match_result_t (size_t key_bytes_matched_,
                                size_t prefix_bytes_matched_,
                                size_t edge_index_,
                                size_t parent_edge_index_,
                                node_t current_,
                                node_t parent_,
                                node_t grandparent_) :
    _key_bytes_matched (key_bytes_matched_),
    _prefix_bytes_matched (prefix_bytes_matched_),
    _edge_index (edge_index_),
    _parent_edge_index (parent_edge_index_),
    _current_node (current_),
    _parent_node (parent_),
    _grandparent_node (grandparent_)
214 215 216
{
}

217 218 219
match_result_t zmq::radix_tree_t::match (const unsigned char *key_,
                                         size_t key_size_,
                                         bool is_lookup_ = false) const
220
{
221
    zmq_assert (key_);
222

223
    // Node we're currently at in the traversal and its predecessors.
224
    node_t current_node = _root;
225 226 227 228 229 230 231 232 233 234
    node_t parent_node = current_node;
    node_t grandparent_node = current_node;
    // Index of the next byte to match in the key.
    size_t key_byte_index = 0;
    // Index of the next byte to match in the current node's prefix.
    size_t prefix_byte_index = 0;
    // Index of the edge from parent to current node.
    size_t edge_index = 0;
    // Index of the edge from grandparent to parent.
    size_t parent_edge_index = 0;
235 236

    while (current_node.prefix_length () > 0 || current_node.edgecount () > 0) {
237 238
        for (prefix_byte_index = 0;
             prefix_byte_index < current_node.prefix_length ()
239
             && key_byte_index < key_size_;
240 241
             ++prefix_byte_index, ++key_byte_index) {
            if (current_node.prefix ()[prefix_byte_index]
242
                != key_[key_byte_index])
243 244 245 246 247
                break;
        }

        // Even if a prefix of the key matches and we're doing a
        // lookup, this means we've found a matching subscription.
248
        if (is_lookup_ && prefix_byte_index == current_node.prefix_length ()
249
            && current_node.refcount () > 0) {
250
            key_byte_index = key_size_;
251 252 253 254 255
            break;
        }

        // There was a mismatch or we've matched the whole key, so
        // there's nothing more to do.
256
        if (prefix_byte_index != current_node.prefix_length ()
257
            || key_byte_index == key_size_)
258 259 260 261
            break;

        // We need to match the rest of the key. Check if there's an
        // outgoing edge from this node.
262 263
        node_t next_node = current_node;
        for (size_t i = 0; i < current_node.edgecount (); ++i) {
264
            if (current_node.first_byte_at (i) == key_[key_byte_index]) {
265 266 267
                parent_edge_index = edge_index;
                edge_index = i;
                next_node = current_node.node_at (i);
268 269 270 271 272 273 274 275 276 277 278
                break;
            }
        }

        if (next_node == current_node)
            break; // No outgoing edge.
        grandparent_node = parent_node;
        parent_node = current_node;
        current_node = next_node;
    }

279 280 281
    return match_result_t (key_byte_index, prefix_byte_index, edge_index,
                           parent_edge_index, current_node, parent_node,
                           grandparent_node);
282 283
}

284
bool zmq::radix_tree_t::add (const unsigned char *key_, size_t key_size_)
285
{
286 287 288 289 290 291
    match_result_t match_result = match (key_, key_size_);
    size_t key_bytes_matched = match_result._key_bytes_matched;
    size_t prefix_bytes_matched = match_result._prefix_bytes_matched;
    size_t edge_index = match_result._edge_index;
    node_t current_node = match_result._current_node;
    node_t parent_node = match_result._parent_node;
292

293
    if (key_bytes_matched != key_size_) {
294
        // Not all characters match, we might have to split the node.
295
        if (prefix_bytes_matched == current_node.prefix_length ()) {
296 297 298
            // The mismatch is at one of the outgoing edges, so we
            // create an edge from the current node to a new leaf node
            // that has the rest of the key as the prefix.
299 300
            node_t key_node = make_node (1, key_size_ - key_bytes_matched, 0);
            key_node.set_prefix (key_ + key_bytes_matched);
301 302 303 304 305 306 307

            // Reallocate for one more edge.
            current_node.resize (current_node.prefix_length (),
                                 current_node.edgecount () + 1);

            // Make room for the new edge. We need to shift the chunk
            // of node pointers one byte to the right. Since resize()
308
            // increments the edgecount by 1, node_pointers() tells us the
309 310 311 312
            // destination address. The chunk of node pointers starts
            // at one byte to the left of this destination.
            //
            // Since the regions can overlap, we use memmove.
313 314
            memmove (current_node.node_pointers (),
                     current_node.node_pointers () - 1,
315 316 317
                     (current_node.edgecount () - 1) * sizeof (void *));

            // Add an edge to the new node.
318
            current_node.set_edge_at (current_node.edgecount () - 1,
319
                                      key_[key_bytes_matched], key_node);
320 321 322 323

            // We need to update all pointers to the current node
            // after the call to resize().
            if (current_node.prefix_length () == 0)
324
                _root._data = current_node._data;
325
            else
326
                parent_node.set_node_at (edge_index, current_node);
327
            ++_size;
328 329 330 331 332 333 334 335 336
            return true;
        }

        // There was a mismatch, so we need to split this node.
        //
        // Create two nodes that will be reachable from the parent.
        // One node will have the rest of the characters from the key,
        // and the other node will have the rest of the characters
        // from the current node's prefix.
337
        node_t key_node = make_node (1, key_size_ - key_bytes_matched, 0);
338 339 340 341
        node_t split_node =
          make_node (current_node.refcount (),
                     current_node.prefix_length () - prefix_bytes_matched,
                     current_node.edgecount ());
342 343

        // Copy the prefix chunks to the new nodes.
344
        key_node.set_prefix (key_ + key_bytes_matched);
345
        split_node.set_prefix (current_node.prefix () + prefix_bytes_matched);
346 347 348

        // Copy the current node's edges to the new node.
        split_node.set_first_bytes (current_node.first_bytes ());
349
        split_node.set_node_pointers (current_node.node_pointers ());
350 351 352 353 354

        // Resize the current node to accommodate a prefix comprising
        // the matched characters and 2 outgoing edges to the above
        // nodes. Set the refcount to 0 since this node doesn't hold a
        // key.
355
        current_node.resize (prefix_bytes_matched, 2);
356 357 358 359 360 361 362
        current_node.set_refcount (0);

        // Add links to the new nodes. We don't need to copy the
        // prefix since resize() retains it in the current node.
        current_node.set_edge_at (0, key_node.prefix ()[0], key_node);
        current_node.set_edge_at (1, split_node.prefix ()[0], split_node);

363
        ++_size;
364
        parent_node.set_node_at (edge_index, current_node);
365 366 367 368
        return true;
    }

    // All characters in the key match, but we still might need to split.
369
    if (prefix_bytes_matched != current_node.prefix_length ()) {
370 371 372 373 374 375
        // All characters in the key match, but not all characters
        // from the current node's prefix match.

        // Create a node that contains the rest of the characters from
        // the current node's prefix and the outgoing edges from the
        // current node.
376 377 378 379 380
        node_t split_node =
          make_node (current_node.refcount (),
                     current_node.prefix_length () - prefix_bytes_matched,
                     current_node.edgecount ());
        split_node.set_prefix (current_node.prefix () + prefix_bytes_matched);
381
        split_node.set_first_bytes (current_node.first_bytes ());
382
        split_node.set_node_pointers (current_node.node_pointers ());
383 384 385

        // Resize the current node to hold only the matched characters
        // from its prefix and one edge to the new node.
386
        current_node.resize (prefix_bytes_matched, 1);
387 388 389

        // Add an edge to the split node and set the refcount to 1
        // since this key wasn't inserted earlier. We don't need to
390 391
        // set the prefix because the first `prefix_bytes_matched` bytes
        // in the prefix are preserved by resize().
392 393 394
        current_node.set_edge_at (0, split_node.prefix ()[0], split_node);
        current_node.set_refcount (1);

395
        ++_size;
396
        parent_node.set_node_at (edge_index, current_node);
397 398 399
        return true;
    }

400
    zmq_assert (key_bytes_matched == key_size_);
401
    zmq_assert (prefix_bytes_matched == current_node.prefix_length ());
402

403
    ++_size;
404 405 406 407
    current_node.set_refcount (current_node.refcount () + 1);
    return current_node.refcount () == 1;
}

408
bool zmq::radix_tree_t::rm (const unsigned char *key_, size_t key_size_)
409
{
410 411 412 413 414 415 416 417
    match_result_t match_result = match (key_, key_size_);
    size_t key_bytes_matched = match_result._key_bytes_matched;
    size_t prefix_bytes_matched = match_result._prefix_bytes_matched;
    size_t edge_index = match_result._edge_index;
    size_t parent_edge_index = match_result._parent_edge_index;
    node_t current_node = match_result._current_node;
    node_t parent_node = match_result._parent_node;
    node_t grandparent_node = match_result._grandparent_node;
418

419
    if (key_bytes_matched != key_size_
420
        || prefix_bytes_matched != current_node.prefix_length ()
421 422 423 424
        || current_node.refcount () == 0)
        return false;

    current_node.set_refcount (current_node.refcount () - 1);
425
    --_size;
426 427 428 429
    if (current_node.refcount () > 0)
        return false;

    // Don't delete the root node.
430
    if (current_node == _root)
431 432 433 434 435 436 437 438 439 440
        return true;

    size_t outgoing_edges = current_node.edgecount ();
    if (outgoing_edges > 1)
        // This node can't be merged with any other node, so there's
        // nothing more to do.
        return true;

    if (outgoing_edges == 1) {
        // Merge this node with the single child node.
441
        node_t child = current_node.node_at (0);
442 443 444 445 446 447 448 449 450 451 452 453 454 455

        // Make room for the child node's prefix and edges. We need to
        // keep the old prefix length since resize() will overwrite
        // it.
        uint32_t old_prefix_length = current_node.prefix_length ();
        current_node.resize (old_prefix_length + child.prefix_length (),
                             child.edgecount ());

        // Append the child node's prefix to the current node.
        memcpy (current_node.prefix () + old_prefix_length, child.prefix (),
                child.prefix_length ());

        // Copy the rest of child node's data to the current node.
        current_node.set_first_bytes (child.first_bytes ());
456
        current_node.set_node_pointers (child.node_pointers ());
457 458
        current_node.set_refcount (child.refcount ());

459
        free (child._data);
460
        parent_node.set_node_at (edge_index, current_node);
461 462 463 464
        return true;
    }

    if (parent_node.edgecount () == 2 && parent_node.refcount () == 0
465
        && parent_node != _root) {
466 467 468
        // Removing this node leaves the parent with one child.
        // If the parent doesn't hold a key or if it isn't the root,
        // we can merge it with its single child node.
469 470
        zmq_assert (edge_index < 2);
        node_t other_child = parent_node.node_at (!edge_index);
471 472 473 474 475 476 477 478 479 480 481 482 483 484

        // Make room for the child node's prefix and edges. We need to
        // keep the old prefix length since resize() will overwrite
        // it.
        uint32_t old_prefix_length = parent_node.prefix_length ();
        parent_node.resize (old_prefix_length + other_child.prefix_length (),
                            other_child.edgecount ());

        // Append the child node's prefix to the current node.
        memcpy (parent_node.prefix () + old_prefix_length,
                other_child.prefix (), other_child.prefix_length ());

        // Copy the rest of child node's data to the current node.
        parent_node.set_first_bytes (other_child.first_bytes ());
485
        parent_node.set_node_pointers (other_child.node_pointers ());
486 487
        parent_node.set_refcount (other_child.refcount ());

488 489
        free (current_node._data);
        free (other_child._data);
490
        grandparent_node.set_node_at (parent_edge_index, parent_node);
491 492 493 494 495 496 497 498
        return true;
    }

    // This is a leaf node that doesn't leave its parent with one
    // outgoing edge. Remove the outgoing edge to this node from the
    // parent.
    zmq_assert (outgoing_edges == 0);

499 500 501 502 503 504 505
    // Replace the edge to the current node with the last edge. An
    // edge consists of a byte and a pointer to the next node. First
    // replace the byte.
    size_t last_index = parent_node.edgecount () - 1;
    unsigned char last_byte = parent_node.first_byte_at (last_index);
    node_t last_node = parent_node.node_at (last_index);
    parent_node.set_edge_at (edge_index, last_byte, last_node);
506 507 508 509

    // Move the chunk of pointers one byte to the left, effectively
    // deleting the last byte in the region of first bytes by
    // overwriting it.
510
    memmove (parent_node.node_pointers () - 1, parent_node.node_pointers (),
511 512 513 514 515 516 517 518
             parent_node.edgecount () * sizeof (void *));

    // Shrink the parent node to the new size, which "deletes" the
    // last pointer in the chunk of node pointers.
    parent_node.resize (parent_node.prefix_length (),
                        parent_node.edgecount () - 1);

    // Nothing points to this node now, so we can reclaim it.
519
    free (current_node._data);
520 521

    if (parent_node.prefix_length () == 0)
522
        _root._data = parent_node._data;
523
    else
524
        grandparent_node.set_node_at (parent_edge_index, parent_node);
525 526 527
    return true;
}

528
bool zmq::radix_tree_t::check (const unsigned char *key_, size_t key_size_)
529
{
530
    if (_root.refcount () > 0)
531 532
        return true;

533 534 535 536 537
    match_result_t match_result = match (key_, key_size_, true);
    return match_result._key_bytes_matched == key_size_
           && match_result._prefix_bytes_matched
                == match_result._current_node.prefix_length ()
           && match_result._current_node.refcount () > 0;
538 539 540
}

static void
541 542 543 544
visit_keys (node_t node_,
            std::vector<unsigned char> &buffer_,
            void (*func_) (unsigned char *data, size_t size, void *arg),
            void *arg_)
545
{
546 547
    for (size_t i = 0; i < node_.prefix_length (); ++i)
        buffer_.push_back (node_.prefix ()[i]);
548

549 550 551
    if (node_.refcount () > 0) {
        zmq_assert (!buffer_.empty ());
        func_ (&buffer_[0], buffer_.size (), arg_);
552 553
    }

554 555 556 557
    for (size_t i = 0; i < node_.edgecount (); ++i)
        visit_keys (node_.node_at (i), buffer_, func_, arg_);
    for (size_t i = 0; i < node_.prefix_length (); ++i)
        buffer_.pop_back ();
558 559
}

560 561
void zmq::radix_tree_t::apply (
  void (*func_) (unsigned char *data, size_t size, void *arg), void *arg_)
562
{
563 564
    if (_root.refcount () > 0)
        func_ (NULL, 0, arg_); // Root node is always empty.
565 566

    std::vector<unsigned char> buffer;
567 568
    for (size_t i = 0; i < _root.edgecount (); ++i)
        visit_keys (_root.node_at (i), buffer, func_, arg_);
569 570
}

571
size_t zmq::radix_tree_t::size () const
572
{
573
    return _size;
574
}