trie_map.hpp

/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

// SPDX-FileCopyrightText: 2015 - 2025 Kohei Yoshida
//
// SPDX-License-Identifier: MIT

#pragma once

#include "trie_map_itr.hpp"

#include <vector>
#include <string>
#include <deque>
#include <map>
#include <memory>
#include <limits>
#include <cstdint>

namespace mdds {

namespace trie {

namespace detail {

struct copy_to_pack
{
};

struct move_to_pack
{
};

template<typename TrieT, typename PackedT>
struct dump_packed_buffer;

} // namespace detail

template<typename T>
struct numeric_value_serializer
{
    static constexpr bool variable_size = false;

    static constexpr size_t value_size = sizeof(T);

    static void write(std::ostream& os, const T& v);

    static void read(std::istream& is, size_t n, T& v);
};

template<typename T>
struct variable_value_serializer
{
    static constexpr bool variable_size = true;

    static void write(std::ostream& os, const T& v);

    static void read(std::istream& is, size_t n, T& v);
};

template<typename T>
struct numeric_sequence_value_serializer
{
    using element_serializer = numeric_value_serializer<typename T::value_type>;

    static constexpr bool variable_size = true;

    static void write(std::ostream& os, const T& v);

    static void read(std::istream& is, size_t n, T& v);
};

template<typename T, typename U = void>
struct value_serializer : numeric_value_serializer<T>
{
};

template<typename T>
struct value_serializer<T, typename std::enable_if<mdds::detail::has_value_type<T>::value>::type>
    : numeric_sequence_value_serializer<T>
{
};

template<>
struct value_serializer<std::string> : variable_value_serializer<std::string>
{
};

struct default_traits
{
    using pack_value_type = uintptr_t;
};

enum class dump_structure_type
{
    packed_buffer,

    trie_traversal,
};

} // namespace trie

template<typename KeyT, typename ValueT, typename TraitsT>
class packed_trie_map;

template<typename KeyT, typename ValueT, typename TraitsT = trie::default_traits>
class trie_map
{
    friend class packed_trie_map<KeyT, ValueT, TraitsT>;
    friend class trie::detail::iterator_base<trie_map, true>;
    friend class trie::detail::iterator_base<trie_map, false>;
    friend class trie::detail::const_iterator<trie_map>;
    friend class trie::detail::iterator<trie_map>;
    friend class trie::detail::search_results<trie_map>;
    friend trie::detail::get_node_stack_type<trie_map, std::true_type>;
    friend trie::detail::get_node_stack_type<trie_map, std::false_type>;

public:
    using traits_type = TraitsT;
    typedef packed_trie_map<KeyT, ValueT, TraitsT> packed_type;
    typedef KeyT key_type;
    typedef typename key_type::value_type key_unit_type;
    typedef ValueT value_type;
    typedef size_t size_type;

    using const_iterator = trie::detail::const_iterator<trie_map>;
    using iterator = trie::detail::iterator<trie_map>;
    typedef trie::detail::search_results<trie_map> search_results;

private:
    struct trie_node
    {
        typedef std::map<key_unit_type, trie_node> children_type;

        children_type children;
        value_type value;
        bool has_value;

        trie_node();
        trie_node(const trie_node& other);
        trie_node(trie_node&& other);

        void swap(trie_node& other);
    };

    template<bool IsConst>
    struct stack_item
    {
        using _is_const = std::bool_constant<IsConst>;

        using child_pos_type =
            typename mdds::detail::get_iterator_type<typename trie_node::children_type, _is_const>::type;

        using trie_node_type = typename mdds::detail::const_or_not<trie_node, _is_const>::type;

        trie_node_type* node;
        child_pos_type child_pos;

        stack_item(trie_node_type* _node, const child_pos_type& _child_pos) : node(_node), child_pos(_child_pos)
        {}

        bool operator==(const stack_item& r) const
        {
            return node == r.node && child_pos == r.child_pos;
        }

        bool operator!=(const stack_item& r) const
        {
            return !operator==(r);
        }
    };

    using const_node_stack_type = std::vector<stack_item<true>>;
    using node_stack_type = std::vector<stack_item<false>>;

public:
    class const_node_type
    {
        friend class trie_map;

        const trie_node* m_ref_node = nullptr;

        const_node_type(const trie_node* ref_node);

    public:
        const_node_type();
        const_node_type(const const_node_type& other);

        const_node_type& operator=(const const_node_type& other);

        bool valid() const;

        bool has_child() const;

        bool has_value() const;

        const value_type& value() const;

        const_node_type child(key_unit_type c) const;
    };

    trie_map();

    trie_map(const trie_map& other);

    trie_map(trie_map&& other);

    const_iterator begin() const;

    iterator begin();

    const_iterator end() const;

    iterator end();

    trie_map& operator=(trie_map other);

    bool operator==(const trie_map& other) const;
    bool operator!=(const trie_map& other) const;

    void swap(trie_map& other);

    const_node_type root_node() const;

    void insert(const key_type& key, value_type value);

    void insert(const key_unit_type* key, size_type len, value_type value);

    bool erase(const key_unit_type* key, size_type len);

    const_iterator find(const key_type& key) const;

    const_iterator find(const key_unit_type* input, size_type len) const;

    iterator find(const key_type& key);

    iterator find(const key_unit_type* input, size_type len);

    search_results prefix_search(const key_type& prefix) const;

    search_results prefix_search(const key_unit_type* prefix, size_type len) const;

    size_type size() const;

    bool empty() const noexcept;

    void clear();

    packed_type pack();

private:
    void insert_into_tree(trie_node& node, const key_unit_type* key, const key_unit_type* key_end, value_type value);

    const trie_node* find_prefix_node(
        const trie_node& node, const key_unit_type* prefix, const key_unit_type* prefix_end) const;

    template<bool IsConst>
    void find_prefix_node_with_stack(
        std::vector<stack_item<IsConst>>& node_stack, mdds::detail::const_t<trie_node, IsConst>& node,
        const key_unit_type* prefix, const key_unit_type* prefix_end) const;

    template<bool IsConst>
    key_type build_key_buffer_from_node_stack(const std::vector<stack_item<IsConst>>& node_stack) const;

    void count_values(size_type& n, const trie_node& node) const;

    bool descend_for_equality(const trie_node& left, const trie_node& right) const;

private:
    trie_node m_root;
};

template<typename KeyT, typename ValueT, typename TraitsT = trie::default_traits>
class packed_trie_map
{
    using pack_value_type = typename TraitsT::pack_value_type;
    using packed_type = std::vector<pack_value_type>;

    friend class trie::detail::packed_iterator_base<packed_trie_map>;
    friend class trie::detail::packed_search_results<packed_trie_map>;
    friend class trie_map<KeyT, ValueT, TraitsT>;
    friend struct trie::detail::dump_packed_buffer<packed_trie_map, packed_type>;

public:
    using traits_type = TraitsT;
    typedef KeyT key_type;
    typedef typename key_type::value_type key_unit_type;
    typedef ValueT value_type;
    typedef size_t size_type;
    typedef trie::detail::packed_iterator_base<packed_trie_map> const_iterator;
    typedef trie::detail::packed_search_results<packed_trie_map> search_results;

    struct entry
    {
        const key_unit_type* key;
        size_type keylen;
        value_type value;

        entry(const key_unit_type* _key, size_type _keylen, value_type _value)
            : key(_key), keylen(_keylen), value(_value)
        {}
    };

private:
    using value_store_type = std::deque<value_type>;

    static constexpr auto null_value = std::numeric_limits<pack_value_type>::max();
    static constexpr auto max_value_pos = null_value - 1u;

    struct trie_node
    {
        key_unit_type key;
        const value_type* value;

        std::deque<trie_node*> children;

        trie_node(key_unit_type _key) : key(_key), value(nullptr)
        {}
    };

    struct stack_item
    {
        const value_store_type* value_store = nullptr;

        const pack_value_type* node_pos = nullptr;
        const pack_value_type* child_pos = nullptr;
        const pack_value_type* child_end = nullptr;

        stack_item(
            const value_store_type* _value_store, const pack_value_type* _node_pos, const pack_value_type* _child_pos,
            const pack_value_type* _child_end)
            : value_store(_value_store), node_pos(_node_pos), child_pos(_child_pos), child_end(_child_end)
        {}

        bool operator==(const stack_item& other) const
        {
            return value_store == other.value_store && node_pos == other.node_pos && child_pos == other.child_pos &&
                   child_end == other.child_end;
        }

        bool operator!=(const stack_item& other) const
        {
            return !operator==(other);
        }

        bool has_value() const
        {
            return *node_pos != null_value;
        }

        pack_value_type get_value_pos() const
        {
            return *node_pos;
        }
    };

    typedef std::vector<stack_item> node_stack_type;
    typedef std::deque<trie_node> node_pool_type;
    typedef std::vector<std::tuple<size_t, key_unit_type>> child_offsets_type;

    packed_trie_map(trie::detail::move_to_pack, trie_map<KeyT, ValueT, TraitsT>& from);

public:
    class const_node_type
    {
        friend class packed_trie_map;

        const packed_type* m_packed = nullptr;
        const value_store_type* m_value_store = nullptr;
        const pack_value_type* m_pos = nullptr;

        const_node_type(const packed_type* packed, const value_store_type* value_store, const pack_value_type* p);

    public:
        const_node_type() = default;
        const_node_type(const const_node_type& other) = default;

        const_node_type& operator=(const const_node_type& other);

        bool valid() const;

        bool has_child() const;

        bool has_value() const;

        const value_type& value() const;

        const_node_type child(key_unit_type c) const;
    };

    packed_trie_map();

    packed_trie_map(const entry* entries, size_type entry_size);

    packed_trie_map(const trie_map<KeyT, ValueT, TraitsT>& other);

    packed_trie_map(const packed_trie_map& other);

    packed_trie_map(packed_trie_map&& other);

    packed_trie_map& operator=(packed_trie_map other);

    bool operator==(const packed_trie_map& other) const;

    bool operator!=(const packed_trie_map& other) const;

    const_iterator begin() const;

    const_iterator end() const;

    const_iterator cbegin() const;

    const_iterator cend() const;

    const_iterator find(const key_type& key) const;

    const_iterator find(const key_unit_type* input, size_type len) const;

    search_results prefix_search(const key_type& prefix) const;

    search_results prefix_search(const key_unit_type* prefix, size_type len) const;

    size_type size() const noexcept;

    bool empty() const noexcept;

    void swap(packed_trie_map& other);

    const_node_type root_node() const;

    template<typename FuncT = trie::value_serializer<value_type>>
    void save_state(std::ostream& os) const;

    template<typename FuncT = trie::value_serializer<value_type>>
    void load_state(std::istream& is);

    std::string dump_structure(trie::dump_structure_type type) const;

private:
    void dump_trie_traversal(std::ostream& os) const;

    node_stack_type get_root_stack() const;

    void traverse_range(
        trie_node& root, node_pool_type& node_pool, const entry* start, const entry* end, size_type pos);

    size_type compact_node(const trie_node& node);

    template<typename ModeT, typename NodeT>
    size_type compact_node(ModeT, NodeT& node);

    void check_value_size_or_throw() const;

    size_type push_value_to_store(trie::detail::copy_to_pack, const value_type& value);
    size_type push_value_to_store(trie::detail::move_to_pack, value_type& value);

    void push_child_offsets(size_type offset, const child_offsets_type& child_offsets);

    void init_pack();
    void compact(const trie_node& root);

    template<typename ModeT, typename NodeT>
    void compact(ModeT, NodeT& root);

    template<typename _Handler>
    void traverse_tree(_Handler hdl) const;

private:
    value_store_type m_value_store;
    packed_type m_packed;
};

} // namespace mdds

#include "trie_map_def.inl"

/* vim:set shiftwidth=4 softtabstop=4 expandtab: */