radix_tree.cpp

/* SPDX-License-Identifier: MPL-2.0 */

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

#include <stdlib.h>
#include <string.h>
#include <iterator>
#include <vector>

node_t::node_t (unsigned char *data_) : _data (data_)
{
}

uint32_t node_t::refcount ()
{
    uint32_t u32;
    memcpy (&u32, _data, sizeof (u32));
    return u32;
}

void node_t::set_refcount (uint32_t value_)
{
    memcpy (_data, &value_, sizeof (value_));
}

uint32_t node_t::prefix_length ()
{
    uint32_t u32;
    memcpy (&u32, _data + sizeof (uint32_t), sizeof (u32));
    return u32;
}

void node_t::set_prefix_length (uint32_t value_)
{
    memcpy (_data + sizeof (value_), &value_, sizeof (value_));
}

uint32_t node_t::edgecount ()
{
    uint32_t u32;
    memcpy (&u32, _data + 2 * sizeof (uint32_t), sizeof (u32));
    return u32;
}

void node_t::set_edgecount (uint32_t value_)
{
    memcpy (_data + 2 * sizeof (value_), &value_, sizeof (value_));
}

unsigned char *node_t::prefix ()
{
    return _data + 3 * sizeof (uint32_t);
}

void node_t::set_prefix (const unsigned char *bytes_)
{
    memcpy (prefix (), bytes_, prefix_length ());
}

unsigned char *node_t::first_bytes ()
{
    return prefix () + prefix_length ();
}

void node_t::set_first_bytes (const unsigned char *bytes_)
{
    memcpy (first_bytes (), bytes_, edgecount ());
}

unsigned char node_t::first_byte_at (size_t index_)
{
    zmq_assert (index_ < edgecount ());
    return first_bytes ()[index_];
}

void node_t::set_first_byte_at (size_t index_, unsigned char byte_)
{
    zmq_assert (index_ < edgecount ());
    first_bytes ()[index_] = byte_;
}

unsigned char *node_t::node_pointers ()
{
    return prefix () + prefix_length () + edgecount ();
}

void node_t::set_node_pointers (const unsigned char *pointers_)
{
    memcpy (node_pointers (), pointers_, edgecount () * sizeof (void *));
}

node_t node_t::node_at (size_t index_)
{
    zmq_assert (index_ < edgecount ());

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

void node_t::set_node_at (size_t index_, node_t node_)
{
    zmq_assert (index_ < edgecount ());
    memcpy (node_pointers () + index_ * sizeof (void *), &node_._data,
            sizeof (node_._data));
}

void node_t::set_edge_at (size_t index_,
                          unsigned char first_byte_,
                          node_t node_)
{
    set_first_byte_at (index_, first_byte_);
    set_node_at (index_, node_);
}

bool node_t::operator== (node_t other_) const
{
    return _data == other_._data;
}

bool node_t::operator!= (node_t other_) const
{
    return !(*this == other_);
}

void node_t::resize (size_t prefix_length_, size_t edgecount_)
{
    const size_t node_size = 3 * sizeof (uint32_t) + prefix_length_
                             + edgecount_ * (1 + sizeof (void *));
    unsigned char *new_data =
      static_cast<unsigned char *> (realloc (_data, node_size));
    zmq_assert (new_data);
    _data = new_data;
    set_prefix_length (static_cast<uint32_t> (prefix_length_));
    set_edgecount (static_cast<uint32_t> (edgecount_));
}

node_t make_node (size_t refcount_, size_t prefix_length_, size_t edgecount_)
{
    const size_t node_size = 3 * sizeof (uint32_t) + prefix_length_
                             + edgecount_ * (1 + sizeof (void *));

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

    node_t node (data);
    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_));
    return node;
}

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

zmq::radix_tree_t::radix_tree_t () : _root (make_node (0, 0, 0)), _size (0)
{
}

static void free_nodes (node_t node_)
{
    for (size_t i = 0, count = node_.edgecount (); i < count; ++i)
        free_nodes (node_.node_at (i));
    free (node_._data);
}

zmq::radix_tree_t::~radix_tree_t ()
{
    free_nodes (_root);
}

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_)
{
}

match_result_t zmq::radix_tree_t::match (const unsigned char *key_,
                                         size_t key_size_,
                                         bool is_lookup_ = false) const
{
    zmq_assert (key_);

    // Node we're currently at in the traversal and its predecessors.
    node_t current_node = _root;
    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;

    while (current_node.prefix_length () > 0 || current_node.edgecount () > 0) {
        const unsigned char *const prefix = current_node.prefix ();
        const size_t prefix_length = current_node.prefix_length ();

        for (prefix_byte_index = 0;
             prefix_byte_index < prefix_length && key_byte_index < key_size_;
             ++prefix_byte_index, ++key_byte_index) {
            if (prefix[prefix_byte_index] != key_[key_byte_index])
                break;
        }

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

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

        // We need to match the rest of the key. Check if there's an
        // outgoing edge from this node.
        node_t next_node = current_node;
        for (size_t i = 0, edgecount = current_node.edgecount (); i < edgecount;
             ++i) {
            if (current_node.first_byte_at (i) == key_[key_byte_index]) {
                parent_edge_index = edge_index;
                edge_index = i;
                next_node = current_node.node_at (i);
                break;
            }
        }

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

    return match_result_t (key_byte_index, prefix_byte_index, edge_index,
                           parent_edge_index, current_node, parent_node,
                           grandparent_node);
}

bool zmq::radix_tree_t::add (const unsigned char *key_, size_t key_size_)
{
    const match_result_t match_result = match (key_, key_size_);
    const size_t key_bytes_matched = match_result._key_bytes_matched;
    const size_t prefix_bytes_matched = match_result._prefix_bytes_matched;
    const size_t edge_index = match_result._edge_index;
    node_t current_node = match_result._current_node;
    node_t parent_node = match_result._parent_node;

    if (key_bytes_matched != key_size_) {
        // Not all characters match, we might have to split the node.
        if (prefix_bytes_matched == current_node.prefix_length ()) {
            // 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.
            node_t key_node = make_node (1, key_size_ - key_bytes_matched, 0);
            key_node.set_prefix (key_ + key_bytes_matched);

            // 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()
            // increments the edgecount by 1, node_pointers() tells us the
            // 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.
            memmove (current_node.node_pointers (),
                     current_node.node_pointers () - 1,
                     (current_node.edgecount () - 1) * sizeof (void *));

            // Add an edge to the new node.
            current_node.set_edge_at (current_node.edgecount () - 1,
                                      key_[key_bytes_matched], key_node);

            // We need to update all pointers to the current node
            // after the call to resize().
            if (current_node.prefix_length () == 0)
                _root._data = current_node._data;
            else
                parent_node.set_node_at (edge_index, current_node);
            _size.add (1);
            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.
        node_t key_node = make_node (1, key_size_ - key_bytes_matched, 0);
        node_t split_node =
          make_node (current_node.refcount (),
                     current_node.prefix_length () - prefix_bytes_matched,
                     current_node.edgecount ());

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

        // Copy the current node's edges to the new node.
        split_node.set_first_bytes (current_node.first_bytes ());
        split_node.set_node_pointers (current_node.node_pointers ());

        // 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.
        current_node.resize (prefix_bytes_matched, 2);
        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);

        _size.add (1);
        parent_node.set_node_at (edge_index, current_node);
        return true;
    }

    // All characters in the key match, but we still might need to split.
    if (prefix_bytes_matched != current_node.prefix_length ()) {
        // 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.
        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);
        split_node.set_first_bytes (current_node.first_bytes ());
        split_node.set_node_pointers (current_node.node_pointers ());

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

        // 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
        // set the prefix because the first `prefix_bytes_matched` bytes
        // in the prefix are preserved by resize().
        current_node.set_edge_at (0, split_node.prefix ()[0], split_node);
        current_node.set_refcount (1);

        _size.add (1);
        parent_node.set_node_at (edge_index, current_node);
        return true;
    }

    zmq_assert (key_bytes_matched == key_size_);
    zmq_assert (prefix_bytes_matched == current_node.prefix_length ());

    _size.add (1);
    current_node.set_refcount (current_node.refcount () + 1);
    return current_node.refcount () == 1;
}

bool zmq::radix_tree_t::rm (const unsigned char *key_, size_t key_size_)
{
    const match_result_t match_result = match (key_, key_size_);
    const size_t key_bytes_matched = match_result._key_bytes_matched;
    const size_t prefix_bytes_matched = match_result._prefix_bytes_matched;
    const size_t edge_index = match_result._edge_index;
    const 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;

    if (key_bytes_matched != key_size_
        || prefix_bytes_matched != current_node.prefix_length ()
        || current_node.refcount () == 0)
        return false;

    current_node.set_refcount (current_node.refcount () - 1);
    _size.sub (1);
    if (current_node.refcount () > 0)
        return false;

    // Don't delete the root node.
    if (current_node == _root)
        return true;

    const 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.
        node_t child = current_node.node_at (0);

        // Make room for the child node's prefix and edges. We need to
        // keep the old prefix length since resize() will overwrite
        // it.
        const 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 ());
        current_node.set_node_pointers (child.node_pointers ());
        current_node.set_refcount (child.refcount ());

        free (child._data);
        parent_node.set_node_at (edge_index, current_node);
        return true;
    }

    if (parent_node.edgecount () == 2 && parent_node.refcount () == 0
        && parent_node != _root) {
        // 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.
        zmq_assert (edge_index < 2);
        node_t other_child = parent_node.node_at (!edge_index);

        // Make room for the child node's prefix and edges. We need to
        // keep the old prefix length since resize() will overwrite
        // it.
        const 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 ());
        parent_node.set_node_pointers (other_child.node_pointers ());
        parent_node.set_refcount (other_child.refcount ());

        free (current_node._data);
        free (other_child._data);
        grandparent_node.set_node_at (parent_edge_index, parent_node);
        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);

    // 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.
    const size_t last_index = parent_node.edgecount () - 1;
    const unsigned char last_byte = parent_node.first_byte_at (last_index);
    const node_t last_node = parent_node.node_at (last_index);
    parent_node.set_edge_at (edge_index, last_byte, last_node);

    // Move the chunk of pointers one byte to the left, effectively
    // deleting the last byte in the region of first bytes by
    // overwriting it.
    memmove (parent_node.node_pointers () - 1, parent_node.node_pointers (),
             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.
    free (current_node._data);

    if (parent_node.prefix_length () == 0)
        _root._data = parent_node._data;
    else
        grandparent_node.set_node_at (parent_edge_index, parent_node);
    return true;
}

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

    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;
}

static void
visit_keys (node_t node_,
            std::vector<unsigned char> &buffer_,
            void (*func_) (unsigned char *data_, size_t size_, void *arg_),
            void *arg_)
{
    const size_t prefix_length = node_.prefix_length ();
    buffer_.reserve (buffer_.size () + prefix_length);
    std::copy (node_.prefix (), node_.prefix () + prefix_length,
               std::back_inserter (buffer_));

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

    for (size_t i = 0, edgecount = node_.edgecount (); i < edgecount; ++i) {
        visit_keys (node_.node_at (i), buffer_, func_, arg_);
    }
    buffer_.resize (static_cast<uint32_t> (buffer_.size () - prefix_length));
}

void zmq::radix_tree_t::apply (
  void (*func_) (unsigned char *data_, size_t size_, void *arg_), void *arg_)
{
    if (_root.refcount () > 0)
        func_ (NULL, 0, arg_); // Root node is always empty.

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

size_t zmq::radix_tree_t::size () const
{
    return _size.get ();
}