segment_tree.hpp

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

#pragma once

#include "./node.hpp"
#include "./global.hpp"

#include <vector>
#include <deque>
#include <iostream>
#include <map>
#include <unordered_map>
#include <memory>
#include <sstream>
#include <type_traits>
#include <utility>

namespace mdds {

template<typename KeyT, typename ValueT>
class segment_tree
{
public:
    using key_type = KeyT;
    using value_type = ValueT;
    using size_type = std::size_t;

private:
    struct segment_type
    {
        key_type start;
        key_type end;
        value_type value;

        static constexpr bool is_nothrow_less_v = noexcept(std::declval<key_type>() < std::declval<key_type>()) &&
                                                  noexcept(std::declval<value_type>() < std::declval<value_type>());

        static constexpr bool is_nothrow_equal_v = noexcept(std::declval<key_type>() == std::declval<key_type>()) &&
                                                   noexcept(std::declval<value_type>() == std::declval<value_type>());

        segment_type() = default;
        segment_type(key_type _start, key_type _end, value_type _value);
        segment_type(const segment_type&) = default;
        segment_type(segment_type&&) = default;

        segment_type& operator=(const segment_type& r) = default;
        segment_type& operator=(segment_type&& r) = default;
        bool operator<(const segment_type& r) const noexcept(is_nothrow_less_v && is_nothrow_equal_v);
        bool operator==(const segment_type& r) const noexcept(is_nothrow_equal_v);
        bool operator!=(const segment_type& r) const noexcept(is_nothrow_equal_v);
    };

    using segment_store_type = std::deque<segment_type>;
    using value_pos_type = typename segment_store_type::size_type;
    using value_chain_type = std::vector<value_pos_type>;

    using node_chain_type = std::vector<st::detail::node_base*>;
    using value_to_nodes_type = std::map<value_pos_type, node_chain_type>;

    struct nonleaf_value_type
    {
        std::unique_ptr<value_chain_type> data_chain;

        nonleaf_value_type()
        {}
        nonleaf_value_type(const nonleaf_value_type& r)
        {
            if (r.data_chain)
                data_chain = std::make_unique<value_chain_type>(*r.data_chain);
        }

        bool operator==(const nonleaf_value_type& r) const
        {
            return data_chain == r.data_chain;
        }

        ~nonleaf_value_type()
        {}
    };

    struct leaf_value_type
    {
        std::unique_ptr<value_chain_type> data_chain;

        leaf_value_type()
        {}
        leaf_value_type(const leaf_value_type& r)
        {
            if (r.data_chain)
                data_chain = std::make_unique<value_chain_type>(*r.data_chain);
        }

        bool operator==(const leaf_value_type& r) const
        {
            return data_chain == r.data_chain;
        }

        ~leaf_value_type()
        {}
    };

public:
    using node = st::detail::node<key_type, leaf_value_type>;
    using node_ptr = typename node::node_ptr;
    using nonleaf_node = typename st::detail::nonleaf_node<key_type, nonleaf_value_type>;

private:
    class search_result_inserter;

    class search_results_base
    {
        friend class search_result_inserter;

    public:
        typedef std::vector<value_chain_type*> res_chains_type;
        typedef std::shared_ptr<res_chains_type> res_chains_ptr;

    protected:
        search_results_base(const segment_store_type& segment_store) : m_segment_store(&segment_store)
        {}

        search_results_base(const search_results_base& r)
            : m_segment_store(r.m_segment_store), mp_res_chains(r.mp_res_chains)
        {}

        bool empty() const;

        size_type size() const;

        void push_back_chain(value_chain_type* chain);

        const segment_store_type* get_segment_store() const
        {
            return m_segment_store;
        }

        const res_chains_ptr& get_res_chains() const
        {
            return mp_res_chains;
        }

    private:
        const segment_store_type* m_segment_store = nullptr;
        res_chains_ptr mp_res_chains;
    };

    class const_iterator_base
    {
        friend class segment_tree;

    protected:
        typedef typename search_results_base::res_chains_type res_chains_type;
        typedef typename search_results_base::res_chains_ptr res_chains_ptr;

        const_iterator_base(const segment_store_type* segment_store, const res_chains_ptr& p)
            : m_segment_store(segment_store), mp_res_chains(p), m_end_pos(true)
        {}

    public:
        using iterator_category = std::bidirectional_iterator_tag;
        using value_type = const segment_tree::segment_type;
        using pointer = value_type*;
        using reference = value_type&;
        using difference_type = std::ptrdiff_t;

        const_iterator_base() = default;
        const_iterator_base(const const_iterator_base& r) = default;
        const_iterator_base& operator=(const const_iterator_base& r) = default;

        value_type* operator++();
        value_type* operator--();

        bool operator==(const const_iterator_base& r) const;
        bool operator!=(const const_iterator_base& r) const;

        value_type& operator*()
        {
            return cur_value();
        }

        const value_type* operator->()
        {
            return &cur_value();
        }

    protected:
        void move_to_front();
        void move_to_end();

    private:
        value_type& cur_value()
        {
            auto pos = *m_cur_pos_in_chain;
            return (*m_segment_store)[pos];
        }

        size_type cur_pos() const
        {
            return *m_cur_pos_in_chain;
        }

        const segment_store_type* m_segment_store = nullptr;
        res_chains_ptr mp_res_chains;
        typename res_chains_type::const_iterator m_cur_chain;
        typename value_chain_type::const_iterator m_cur_pos_in_chain;
        bool m_end_pos = true;
    };

    class search_result_inserter
    {
    public:
        search_result_inserter(search_results_base& results) : m_results(results)
        {}
        void operator()(value_chain_type* node_data)
        {
            if (!node_data)
                return;

            m_results.push_back_chain(node_data);
        }

    private:
        search_results_base& m_results;
    };

    static constexpr bool nothrow_default_constructible_v =
        std::is_nothrow_default_constructible_v<std::vector<nonleaf_node>> &&
        std::is_nothrow_default_constructible_v<segment_store_type> &&
        std::is_nothrow_default_constructible_v<value_to_nodes_type> &&
        std::is_nothrow_default_constructible_v<node_ptr>;

    static constexpr bool nothrow_move_constructible_v =
        std::is_nothrow_move_constructible_v<std::vector<nonleaf_node>> &&
        std::is_nothrow_move_constructible_v<segment_store_type> &&
        std::is_nothrow_move_constructible_v<value_to_nodes_type> && std::is_nothrow_move_constructible_v<node_ptr>;

    static constexpr bool nothrow_swappable_v =
        std::is_nothrow_swappable_v<std::vector<nonleaf_node>> && std::is_nothrow_swappable_v<segment_store_type> &&
        std::is_nothrow_swappable_v<value_to_nodes_type> && std::is_nothrow_swappable_v<node_ptr>;

    static constexpr bool nothrow_move_assignable_v = nothrow_move_constructible_v && nothrow_swappable_v;

public:
    class search_results : public search_results_base
    {
        friend class segment_tree;

        typedef typename search_results_base::res_chains_type res_chains_type;
        typedef typename search_results_base::res_chains_ptr res_chains_ptr;

        search_results(const segment_store_type& segment_tree) : search_results_base(segment_tree)
        {}

    public:
        class const_iterator : public const_iterator_base
        {
            friend class segment_tree<KeyT, ValueT>::search_results;

        private:
            const_iterator(const segment_store_type* segment_store, const res_chains_ptr& p)
                : const_iterator_base(segment_store, p)
            {}

        public:
            const_iterator() : const_iterator_base()
            {}

            const_iterator(const const_iterator& r) : const_iterator_base(r)
            {}

            const_iterator& operator=(const const_iterator& r)
            {
                const_iterator_base::operator=(r);
                return *this;
            }
        };

        bool empty() const
        {
            return search_results_base::empty();
        }

        size_type size() const
        {
            return search_results_base::size();
        }

        typename search_results::const_iterator begin() const
        {
            typename search_results::const_iterator it(
                search_results_base::get_segment_store(), search_results_base::get_res_chains());
            it.move_to_front();
            return it;
        }

        typename search_results::const_iterator end() const
        {
            typename search_results::const_iterator it(
                search_results_base::get_segment_store(), search_results_base::get_res_chains());
            it.move_to_end();
            return it;
        }
    };

    segment_tree() noexcept(nothrow_default_constructible_v);
    segment_tree(const segment_tree& r);
    segment_tree(segment_tree&& r) = default;
    ~segment_tree();

    segment_tree& operator=(const segment_tree& r);
    segment_tree& operator=(segment_tree&& r) noexcept(nothrow_move_assignable_v);

    bool operator==(const segment_tree& r) const;

    bool operator!=(const segment_tree& r) const;

    bool valid_tree() const noexcept
    {
        return m_valid_tree;
    }

    void build_tree();

    void insert(key_type start_key, key_type end_key, value_type value);

    search_results search(const key_type& point) const;

    void erase(const typename search_results::const_iterator& pos);

    template<typename Pred>
    size_type erase_if(Pred pred);

    void swap(segment_tree& r) noexcept(nothrow_swappable_v);

    void clear();

    size_type size() const;

    bool empty() const;

    size_type leaf_size() const
        noexcept(noexcept(st::detail::count_leaf_nodes<size_type>(m_left_leaf.get(), m_right_leaf.get())))
    {
        return st::detail::count_leaf_nodes<size_type>(m_left_leaf.get(), m_right_leaf.get());
    }

    std::string to_string() const;

    std::vector<key_type> boundary_keys() const;

    struct integrity_check_properties
    {
        struct leaf_node
        {
            key_type key = {};
            std::vector<value_type> value_chain;
        };

        std::vector<leaf_node> leaf_nodes;
    };

    void check_integrity(const integrity_check_properties& props) const;

private:
    static void create_leaf_node_instances(std::vector<key_type> keys, node_ptr& left, node_ptr& right);

    void descend_tree_and_mark(
        st::detail::node_base* pnode, value_pos_type value, const key_type& start_key, const key_type& end_key,
        node_chain_type& nodelist);

    void build_leaf_nodes();

    void remove_data_from_nodes(node_chain_type& nodelist, value_pos_type value);
    void remove_data_from_chain(value_chain_type& chain, value_pos_type value);
    void remove_value_pos(size_type pos);

    void clear_all_nodes();

private:
    struct tree_dumper_traits
    {
        using leaf_type = node;
        using nonleaf_type = nonleaf_node;
        using tree_type = segment_tree;

        struct to_string
        {
            const tree_type& tree;

            to_string(const tree_type& _tree);
            std::string operator()(const leaf_type& leaf) const;
            std::string operator()(const nonleaf_type& nonleaf) const;
        };
    };

    std::vector<nonleaf_node> m_nonleaf_node_pool;

    segment_store_type m_segment_store;

    value_to_nodes_type m_tagged_nodes_map;

    nonleaf_node* m_root_node;
    node_ptr m_left_leaf;
    node_ptr m_right_leaf;
    bool m_valid_tree;
};

} // namespace mdds

#include "segment_tree_def.inl"