SkipTrie.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cmath> // for log2
#include <cstddef> // for size_t
#include <memory> // for std::shared_ptr
#include <random> // for random height generation
#include <vector> // for std::vector
#include <unordered_map> // for std::unordered_map
#include <fstream> // for std::ostream
#include <string> // for std::string
#include <sstream> // for std::ostringstream in to_string
#include <compare> // for std::strong_ordering
 
#include "BitString.cuh"
 
enum class Direction
{
    FORWARD, 
    BACKWARD, 
    INPLACE 
};
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS = sizeof(CHAR_T) * 8>
class SkipTrie
{
public:
    using KEY_T = BitString<CHAR_T, CHAR_SIZE_BITS>;
 
    SkipTrie();
 
    static size_t get_random_height();
 
    bool insert(const KEY_T* key) noexcept;
 
    bool insert(const KEY_T* key, size_t height) noexcept;
 
    bool contains(const KEY_T* key) const noexcept;
 
    bool remove(const KEY_T* key) noexcept;
 
    std::ostream& print(std::ostream& os) const noexcept;
 
    std::string to_string() const noexcept;
 
    size_t size() const noexcept { return m_size; }
 
    size_t height() const noexcept { return m_height; }
 
    size_t lcp(const KEY_T* key) const noexcept;
 
    size_t lcp_with_others(const KEY_T* key) const noexcept;
 
    std::vector<const KEY_T*> suffix_search(const KEY_T* key) const noexcept;
 
    std::vector<const KEY_T*> range_search(const KEY_T* key1, const KEY_T* key2) const noexcept;
 
    std::unordered_map<size_t, size_t> get_lcp_group_sizes() const noexcept;
 
protected:
    size_t m_size;
    size_t m_height;
 
    struct Node
    {
        const KEY_T* key = nullptr;
        std::shared_ptr<Node> next{nullptr};
        Node* prev = nullptr;
        std::shared_ptr<Node> down{nullptr};
        size_t lcp_next{0};
 
        size_t lcp(Direction direction) const noexcept;
 
        Node* next_node(Direction direction) const noexcept;
    };
 
    std::shared_ptr<Node> m_head = nullptr;
    std::shared_ptr<Node> m_lower_head = nullptr;
    std::shared_ptr<Node> m_lower_tail = nullptr;
 
    void add_layer() noexcept;
 
    void remove_layer() noexcept;
 
    std::shared_ptr<Node> insert_recursive(const KEY_T* key, Node* curr, size_t height, Direction direction, size_t lcp, size_t current_height) noexcept;
 
    Direction iter_layer(const KEY_T* key, Node*& curr, Direction direction, size_t& lcp) const noexcept;
 
    struct NodeLCP
    {
        Node* node;
        size_t lcp;
    };
 
    virtual NodeLCP find_first(const KEY_T* key, const bool require_level0 = false) const noexcept;
 
    struct EqualOrSuccessor
    {
        Node* node;
        bool is_equal;
        size_t lcp;
    };
 
    virtual EqualOrSuccessor find_equal_or_successor(const KEY_T* key, const bool require_level0) const noexcept;
 
    struct ResultLCP
    {
        std::strong_ordering result;
        size_t lcp;
    };
 
    virtual ResultLCP compare(const KEY_T* key1, const KEY_T* key2, size_t lcp) const noexcept;
 
public:
    class Iterator
    {
    public:
        Iterator(Node* node, Direction direction) noexcept;
 
        Iterator& operator++() noexcept;
 
        Iterator operator++(int) noexcept;
 
        const KEY_T& operator*() const noexcept;
 
        const KEY_T* operator->() const noexcept;
 
        bool operator==(const Iterator& other) const noexcept;
 
        bool operator!=(const Iterator& other) const noexcept;
 
    private:
        Node* m_node;
        Direction m_direction;
    };
 
    Iterator begin() const noexcept;
 
    Iterator end() const noexcept;
};
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
std::ostream& operator<<(std::ostream& os, const SkipTrie<CHAR_T, CHAR_SIZE_BITS>& ssl);
 
 
// ========================================================================== //
//                           Implementation Details                           //
// ========================================================================== //
 
// --- Iterator Implementation ---
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator::Iterator(Node* node, Direction direction) noexcept
    : m_node(node), m_direction(direction)
{
    // Constructor body intentionally empty
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
typename SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator& SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator::operator++() noexcept
{
    m_node = m_node->next_node(m_direction);
    return *this;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
typename SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator::operator++(int) noexcept
{
    Iterator temp = *this;
    ++(*this); // Use pre-increment
    return temp;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
const BitString<CHAR_T, CHAR_SIZE_BITS>& SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator::operator*() const noexcept
{
    return *m_node->key;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
const BitString<CHAR_T, CHAR_SIZE_BITS>* SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator::operator->() const noexcept
{
    return m_node->key;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
bool SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator::operator==(const Iterator& other) const noexcept
{
    return m_node == other.m_node;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
bool SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator::operator!=(const Iterator& other) const noexcept
{
    return !(*this == other); // Delegate to operator==
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
typename SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator SkipTrie<CHAR_T, CHAR_SIZE_BITS>::begin() const noexcept
{
    // Returns iterator starting from the first actual node after the head sentinel at the base level.
    return Iterator(m_lower_head->next.get(), Direction::FORWARD);
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
typename SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Iterator SkipTrie<CHAR_T, CHAR_SIZE_BITS>::end() const noexcept
{
    // Returns iterator pointing to the tail sentinel at the base level.
    return Iterator(m_lower_tail.get(), Direction::FORWARD);
}
 
// --- Node Implementation ---
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
size_t SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Node::lcp(Direction direction) const noexcept
{
    if (direction == Direction::FORWARD)
    {
        // LCP with the next node is stored directly.
        return lcp_next;
    }
 
    // LCP with the previous node is stored in the previous node's lcp_next.
    return prev ? prev->lcp_next : 0;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
typename SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Node* SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Node::next_node(Direction direction) const noexcept
{
    if (direction == Direction::FORWARD)
    {
        return next.get();
    }
 
    return prev;
}
 
// --- SkipTrie Implementation ---
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
SkipTrie<CHAR_T, CHAR_SIZE_BITS>::SkipTrie()
    : m_size(0), m_height(0)
{
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
size_t SkipTrie<CHAR_T, CHAR_SIZE_BITS>::get_random_height()
{
    // Use static variables for random number generation.
    // Note: Not thread-safe if multiple threads call this concurrently.
    static std::random_device rd; // Obtain a random seed from the hardware device
    static std::mt19937 gen(rd()); // Seed the default random engine
    // static std::mt19937 gen(1); // Use for deterministic testing (fixed seed)
    static std::geometric_distribution<size_t> dist(0.5); // p=0.5 gives 50% chance of 0, 25% of 1, etc.
 
    return dist(gen);
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
void SkipTrie<CHAR_T, CHAR_SIZE_BITS>::add_layer() noexcept
{
    // Create new sentinel nodes for the layer.
    std::shared_ptr<Node> new_head = std::make_shared<Node>();
    std::shared_ptr<Node> new_tail = std::make_shared<Node>();
 
    // Link the new head down to the old head (or null if no layers yet).
    new_head->down = m_head;
    // Link the new tail down to the old tail (or null if no layers yet).
    // Access old tail via old head's next pointer.
    if (m_head)
    {
        new_tail->down = m_head->next;
    }
    // Link new head and tail horizontally.
    new_head->next = new_tail;
    new_tail->prev = new_head.get();
 
    // Update the main head pointer to the new head.
    m_head = new_head;
    // Increment the height counter.
    ++m_height;
 
    // If this is the very first layer added, set the lower head/tail pointers.
    if (!m_lower_head)
    {
        m_lower_head = m_head;
        m_lower_tail = new_tail;
    }
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
void SkipTrie<CHAR_T, CHAR_SIZE_BITS>::remove_layer() noexcept
{
    // Move the head pointer down to the next layer.
    m_head = m_head->down;
    // Decrement the height.
    --m_height;
 
    // If height becomes 0, reset lower head/tail pointers.
    if (m_height == 0)
    {
        m_lower_head = nullptr;
        m_lower_tail = nullptr;
    }
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
auto SkipTrie<CHAR_T, CHAR_SIZE_BITS>::compare(const KEY_T* key1, const KEY_T* key2, size_t lcp) const noexcept -> ResultLCP
{
    // Delegate comparison to the BitString's sequential compare method.
    auto [comparison, next_lcp] = key1->seq_k_compare(*key2, lcp);
    return { comparison, next_lcp };
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
Direction SkipTrie<CHAR_T, CHAR_SIZE_BITS>::iter_layer(const KEY_T* key, Node*& curr, Direction direction, size_t& lcp) const noexcept
{
    // If already found inplace, no need to iterate.
    if (direction == Direction::INPLACE)
    {
        return Direction::INPLACE;
    }
 
    // Get the next node in the current direction.
    Node* next = curr->next_node(direction);
 
    // Iterate while the next node exists and the LCP allows potential match/crossing.
    // Check `next->key` to ensure we don't compare against the tail sentinel.
    // `curr->lcp(direction)` gives LCP between `curr` and `next`.
    while (next->key && curr->lcp(direction) >= lcp)
    {
        // Optimization: If LCP between curr and next is strictly greater than current path LCP,
        // we know the target key cannot be between curr and next based on prefixes. Skip over next.
        if (curr->lcp(direction) > lcp)
        {
            curr = next;
            next = curr->next_node(direction);
            continue;
        }
 
        // If LCPs are equal, we must perform a full comparison starting from `lcp`.
        auto [comparison, next_lcp] = compare(key, next->key, lcp);
        // Update the path LCP with the result of the comparison.
        lcp = next_lcp;
 
        // Move to the next node.
        curr = next;
 
        // Check comparison result:
        if (comparison == std::strong_ordering::equal)
        {
            // Exact match found.
            return Direction::INPLACE;
        }
        else if (comparison == std::strong_ordering::less)
        {
            // Key is less than curr->key.
            if (direction == Direction::FORWARD)
            {
                // If moving forward, we've gone past the insertion point. Need to go backward on level below.
                return Direction::BACKWARD;
            }
            // If already moving backward, continue moving backward.
        }
        else // comparison == std::strong_ordering::greater
        {
            // Key is greater than curr->key.
            if (direction == Direction::BACKWARD)
            {
                // If moving backward, we've gone past the insertion point. Need to go forward on level below.
                return Direction::FORWARD;
            }
            // If already moving forward, continue moving forward.
        }
 
        // Get the next node for the next iteration.
        next = curr->next_node(direction);
    }
 
    // Loop finished: `next` is either the tail sentinel or `curr->lcp(direction) < lcp`.
    // The correct direction to proceed on the level below is the current `direction`.
    return direction;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
typename SkipTrie<CHAR_T, CHAR_SIZE_BITS>::NodeLCP SkipTrie<CHAR_T, CHAR_SIZE_BITS>::find_first(const KEY_T* key, const bool require_level0) const noexcept
{
    // Use find_equal_or_successor to perform the search.
    auto [node, is_equal, lcp] = find_equal_or_successor(key, require_level0);
 
    // Return the node only if an exact match was found.
    return { is_equal ? node : nullptr, lcp };
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
bool SkipTrie<CHAR_T, CHAR_SIZE_BITS>::contains(const KEY_T* key) const noexcept
{
    // Check if find_first returns a non-null node.
    return find_first(key).node != nullptr;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
bool SkipTrie<CHAR_T, CHAR_SIZE_BITS>::insert(const KEY_T* key) noexcept
{
    // Insert with a randomly generated height.
    return insert(key, get_random_height());
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
bool SkipTrie<CHAR_T, CHAR_SIZE_BITS>::insert(const KEY_T* key, size_t height) noexcept
{
    // Ensure the skip list has enough layers for the specified height.
    // Need height + 1 levels (0 to height). m_height is number of levels.
    while (height + 1 >= m_height)
    {
        add_layer();
    }
 
    // Call the recursive insertion helper. Starts from top head node (m_head).
    // Initial direction is FORWARD, initial LCP is 0, current height starts at m_height - 1.
    // insert_recursive returns nullptr if key already exists.
    bool already_exists = insert_recursive(key, m_head.get(), height, Direction::FORWARD, 0, m_height - 1) == nullptr;
 
    if (already_exists)
    {
        return false; // Key was not inserted because it already existed.
    }
 
    // Increment size only if insertion was successful.
    ++m_size;
    return true;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
std::shared_ptr<typename SkipTrie<CHAR_T, CHAR_SIZE_BITS>::Node> SkipTrie<CHAR_T, CHAR_SIZE_BITS>::insert_recursive(const KEY_T* key, Node* curr, size_t height, Direction direction, size_t lcp, size_t current_height) noexcept
{
    // Find the predecessor node (`curr`) on the current level.
    Direction next_direction = iter_layer(key, curr, direction, lcp);
 
    // If iter_layer found an exact match, return nullptr to signal duplication.
    if (next_direction == Direction::INPLACE)
    {
        return nullptr;
    }
 
    // Recursively insert on the level below, unless we are at the base level.
    std::shared_ptr<Node> child = nullptr;
    if (current_height > 0)
    {
        child = insert_recursive(key, curr->down.get(), height, next_direction, lcp, current_height - 1);
 
        // If recursive call returned nullptr, propagate nullptr up.
        if (!child)
        {
            return nullptr;
        }
    }
 
 
    // If the current level is above the target insertion height, just return the child pointer
    // (or nullptr if base level) without inserting at this level.
    if (current_height > height)
    {
        return child;
    }
 
    // --- Insert node at the current level ---
 
    // Create the new node.
    std::shared_ptr<Node> new_node = std::make_shared<Node>();
    new_node->key = key;
    new_node->down = child; // Link down to the node created/returned from the level below.
 
    // Link the new node horizontally based on the direction determined by iter_layer.
    // `curr` is the predecessor node found by iter_layer.
    switch (next_direction)
    {
    case Direction::FORWARD: // Insert new_node after curr
        new_node->next = curr->next;
        new_node->prev = curr;
        curr->next->prev = new_node.get(); // Update next node's prev pointer
        curr->next = new_node;
 
        // Update LCP values.
        new_node->lcp_next = curr->lcp_next; // New node inherits LCP curr had with its old next.
        curr->lcp_next = lcp; // LCP between curr and new_node is the final LCP from iter_layer.
        break;
    case Direction::BACKWARD: // Insert new_node before curr (i.e., after curr->prev)
        new_node->next = curr->prev->next; // Should be `curr` itself
        new_node->prev = curr->prev;
        curr->prev->next = new_node; // Update previous node's next pointer
        curr->prev = new_node.get();
 
        // Update LCP values.
        new_node->lcp_next = lcp; // LCP between new_node and curr is the final LCP from iter_layer.
        // new_node->prev->lcp_next remains unchanged (LCP between prev and new_node).
        break;
    case Direction::INPLACE:
        // This case should not be reached here due to the check after iter_layer.
        // If it were, it would indicate an error.
        break;
    }
 
    // Return the newly created node so the level above can link its `down` pointer.
    return new_node;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
bool SkipTrie<CHAR_T, CHAR_SIZE_BITS>::remove(const KEY_T* key) noexcept
{
    Node* curr = find_first(key).node;
 
    // If key not found, return false.
    if (!curr)
    {
        return false;
    }
 
    // remove the node and all its children
    while (curr)
    {
        auto prev = curr->prev;
        auto next = curr->next;
        prev->lcp_next = std::min(prev->lcp_next, curr->lcp_next);
        curr = curr->down.get();
        
        prev->next = next;
        next->prev = prev;
    }
 
    // remove empty layers
    while (!m_head->down || !m_head->down->next->key)
    {
        remove_layer();
 
        if (m_height == 0)
        {
            break;
        }
    }
 
    --m_size;
 
    return true;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
size_t SkipTrie<CHAR_T, CHAR_SIZE_BITS>::lcp(const KEY_T* key) const noexcept
{
    // Find the node (or insertion point) and return the LCP computed during search.
    return find_first(key).lcp;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
size_t SkipTrie<CHAR_T, CHAR_SIZE_BITS>::lcp_with_others(const KEY_T* key) const noexcept
{
    auto [node, lcp] = find_first(key, true);
 
    if (!node)
    {
        return lcp;
    }
 
    return std::max(node->lcp(Direction::FORWARD), node->lcp(Direction::BACKWARD));
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
std::vector<const BitString<CHAR_T, CHAR_SIZE_BITS>*> SkipTrie<CHAR_T, CHAR_SIZE_BITS>::suffix_search(const KEY_T* key) const noexcept
{
    std::vector<const KEY_T*> keys;
 
    auto [curr, _, lcp] = find_equal_or_successor(key, true).node;
 
    if (lcp != key->size())
    {
        return keys; // No keys with this prefix found.
    }
 
    keys.push_back(curr->key);
 
    // iterate through the list until the edge lcp is less than the key length
    while (curr->next->key && curr->lcp_next >= lcp)
    {
        curr = curr->next.get();
        keys.push_back(curr->key);
    }
 
    return keys;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
typename SkipTrie<CHAR_T, CHAR_SIZE_BITS>::EqualOrSuccessor SkipTrie<CHAR_T, CHAR_SIZE_BITS>::find_equal_or_successor(const KEY_T* key, const bool require_level0) const noexcept
{
    // Handle empty list case.
    if (!m_head)
    {
        return { nullptr, false, 0 };
    }
 
    Node* curr = m_head.get();
    Node* prev = nullptr;
    Direction direction = Direction::FORWARD;
    size_t lcp = 0;
 
    // Traverse down the levels.
    while (curr)
    {
        // Find predecessor on current level.
        Direction next_direction = iter_layer(key, curr, direction, lcp);
 
        if (next_direction == Direction::INPLACE)
        {
            while (require_level0 && curr->down)
            {
                curr = curr->down.get();
            }
 
            return { curr, true, lcp };
        }
 
        prev = curr;
        curr = curr->down.get();
        direction = next_direction;
    }
 
    return { direction == Direction::FORWARD ? prev->next.get() : prev, false, lcp };
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
std::vector<const BitString<CHAR_T, CHAR_SIZE_BITS>*> SkipTrie<CHAR_T, CHAR_SIZE_BITS>::range_search(const KEY_T* key1, const KEY_T* key2) const noexcept
{
    std::vector<const KEY_T*> keys;
 
    // Find the starting node (>= key1) at the base level.
    Node* curr = find_equal_or_successor(key1, true).node;
 
    if (!curr)
    {
        return keys;
    }
 
    auto [end, is_equal] = find_equal_or_successor(key2, true);
 
    if (is_equal)
    {
        end = end->next.get();
    }
 
    while (curr != end)
    {
        keys.push_back(curr->key);
        curr = curr->next.get();
    }
 
    return keys;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
std::unordered_map<size_t, size_t> SkipTrie<CHAR_T, CHAR_SIZE_BITS>::get_lcp_group_sizes() const noexcept
{
    std::unordered_map<size_t, size_t> lcp_group_sizes;
 
    std::vector<size_t> lcps_we_are_greater_than;
    std::vector<size_t> index_first_encountered;
 
    Node* curr = m_lower_head->next.get();
 
    size_t i = 0;
 
    while (curr->key)
    {
        size_t lcp = curr->lcp_next;
 
        while (!lcps_we_are_greater_than.empty() && lcp < lcps_we_are_greater_than.back())
        {
            size_t popped_lcp = lcps_we_are_greater_than.back();
            lcps_we_are_greater_than.pop_back();
 
            size_t group_size = i - index_first_encountered.back();
            index_first_encountered.pop_back();
 
            lcp_group_sizes[popped_lcp] = std::max(lcp_group_sizes[popped_lcp], group_size);
        }
 
        if (lcps_we_are_greater_than.empty() || lcp > lcps_we_are_greater_than.back())
        {
            lcps_we_are_greater_than.push_back(lcp);
            index_first_encountered.push_back(i);
        }
        // If current_lcp is equal to top, do nothing (extend the current group).
 
        // Move to the next node and increment the gap index.
        curr = curr->next.get();
    }
 
    return lcp_group_sizes;
}
 
// --- Stream Operator Implementation ---
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
std::ostream& operator<<(std::ostream& os, const SkipTrie<CHAR_T, CHAR_SIZE_BITS>& ssl) {
    // Delegate to the print method.
    return ssl.print(os);
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
std::ostream& SkipTrie<CHAR_T, CHAR_SIZE_BITS>::print(std::ostream& os) const noexcept
{
    // print layer by layer
    auto left = m_head;
    size_t height = m_height - 1;
    while (left)
    {
        os << "Height: " << height << "|\t";
        os << "HEAD -0-> ";
        auto curr = left->next;
        while (curr->next)
        {
            // Print the key and the LCP to the next node.
            // Assumes KEY_T has an overloaded operator<<.
            os << *curr->key << " -" << curr->lcp_next << "-> ";
            curr = curr->next;
        }
 
        os << "TAIL\n";
        left = left->down;
        --height;
    }
 
    return os;
}
 
template<typename CHAR_T, unsigned CHAR_SIZE_BITS>
std::string SkipTrie<CHAR_T, CHAR_SIZE_BITS>::to_string() const noexcept
{
    // Use a string stream to capture the output of the print method.
    std::ostringstream oss;
    print(oss);
    return oss.str();
}