rtree.hpp

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

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

#pragma once

#include <deque>
#include <vector>
#include <cstdlib>
#include <string>
#include <unordered_set>
#include <unordered_map>
#include <functional>

namespace mdds {

namespace detail { namespace rtree {

struct default_rtree_traits
{
    static constexpr size_t dimensions = 2;

    static constexpr size_t min_node_size = 40;

    static constexpr size_t max_node_size = 100;

    static constexpr size_t max_tree_depth = 100;

    static constexpr bool enable_forced_reinsertion = true;

    static constexpr size_t reinsertion_size = 30;
};

struct integrity_check_properties
{
    bool throw_on_first_error = true;

    bool error_on_min_node_size = true;
};

enum class node_type
{
    unspecified,
    directory_leaf,
    directory_nonleaf,
    value
};

enum class export_tree_type
{
    formatted_node_properties,

    extent_as_obj,

    extent_as_svg
};

enum class search_type
{
    overlap,

    match,
};

template<typename _NodePtrT>
struct ptr_to_string;

}} // namespace detail::rtree

template<typename KeyT, typename ValueT, typename Traits = detail::rtree::default_rtree_traits>
class rtree
{
    using traits_type = Traits;

    constexpr static size_t max_dist_size = traits_type::max_node_size - traits_type::min_node_size * 2 + 2;

public:
    using key_type = KeyT;
    using value_type = ValueT;

    struct point_type
    {
        key_type d[traits_type::dimensions];

        point_type();
        point_type(std::initializer_list<key_type> vs);

        std::string to_string() const;

        bool operator==(const point_type& other) const;
        bool operator!=(const point_type& other) const;
    };

    struct extent_type
    {
        point_type start;
        point_type end;

        extent_type();
        extent_type(const point_type& _start, const point_type& _end);

        std::string to_string() const;

        bool is_point() const;

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

        bool contains(const point_type& pt) const;

        bool contains(const extent_type& bb) const;

        bool intersects(const extent_type& bb) const;

        bool contains_at_boundary(const extent_type& other) const;
    };

    using node_type = detail::rtree::node_type;
    using export_tree_type = detail::rtree::export_tree_type;
    using search_type = detail::rtree::search_type;
    using integrity_check_properties = detail::rtree::integrity_check_properties;

    struct node_properties
    {
        node_type type;
        extent_type extent;
    };

private:
    struct node;
    struct directory_node;

    struct node_store
    {
        node_type type;
        extent_type extent;
        node_store* parent;
        node* node_ptr;
        size_t count;

        bool valid_pointer;

        node_store(const node_store&) = delete;
        node_store& operator=(const node_store&) = delete;

        node_store();
        node_store(node_store&& r);
        node_store(node_type _type, const extent_type& _extent, node* _node_ptr);
        ~node_store();

        node_store clone() const;

        static node_store create_leaf_directory_node();
        static node_store create_nonleaf_directory_node();
        static node_store create_value_node(const extent_type& extent, value_type&& v);
        static node_store create_value_node(const extent_type& extent, const value_type& v);

        node_store& operator=(node_store&& other);

        bool pack();

        void pack_upward();

        bool is_directory() const;
        bool is_root() const;
        bool exceeds_capacity() const;

        void swap(node_store& other);

        void reset_parent_pointers_of_children();

        void reset_parent_pointers();

        directory_node* get_directory_node();
        const directory_node* get_directory_node() const;

        bool erase_child(const node_store* p);
    };

    using dir_store_type = std::deque<node_store>;

    struct dir_store_segment
    {
        typename dir_store_type::iterator begin;
        typename dir_store_type::iterator end;
        size_t size;

        dir_store_segment() : size(0)
        {}

        dir_store_segment(
            typename dir_store_type::iterator _begin, typename dir_store_type::iterator _end, size_t _size)
            : begin(std::move(_begin)), end(std::move(_end)), size(_size)
        {}
    };

    struct distribution
    {
        dir_store_segment g1;
        dir_store_segment g2;

        distribution(size_t dist, dir_store_type& nodes)
        {
            g1.begin = nodes.begin();
            g1.end = g1.begin;
            g1.size = traits_type::min_node_size - 1 + dist;
            std::advance(g1.end, g1.size);

            g2.begin = g1.end;
            g2.end = nodes.end();
            g2.size = nodes.size() - g1.size;
        }
    };

    struct node
    {
        node();
        node(const node&) = delete;
        ~node();
    };

    struct value_node : public node
    {
        value_type value;

        value_node() = delete;
        value_node(value_type&& _value);
        value_node(const value_type& _value);
        ~value_node();
    };

    struct directory_node : public node
    {
        dir_store_type children;

        directory_node(const directory_node&) = delete;
        directory_node& operator=(const directory_node&) = delete;

        directory_node();
        ~directory_node();

        bool erase(const node_store* p);

        node_store* get_child_with_minimal_overlap(const extent_type& bb);
        node_store* get_child_with_minimal_area_enlargement(const extent_type& bb);

        extent_type calc_extent() const;
        key_type calc_overlap_cost(const extent_type& bb) const;
        bool has_leaf_directory() const;
    };

    struct orphan_node_entry
    {
        node_store ns;
        size_t depth;
    };

    using orphan_node_entries_type = std::deque<orphan_node_entry>;

    struct insertion_point
    {
        node_store* ns;
        size_t depth;
    };

public:
    template<typename NS>
    class search_results_base
    {
        friend class rtree;

    protected:
        using node_store_type = NS;

        struct entry
        {
            node_store_type* ns;
            size_t depth;

            entry(node_store_type* _ns, size_t _depth);
        };

        using store_type = std::vector<entry>;
        store_type m_store;

        void add_node_store(node_store_type* ns, size_t depth);
    };

    class const_iterator;
    class iterator;
    class search_results;

    class const_search_results : public search_results_base<const node_store>
    {
        using base_type = search_results_base<const node_store>;
        using base_type::m_store;

    public:
        const_iterator cbegin() const;
        const_iterator cend() const;
        const_iterator begin() const;
        const_iterator end() const;
    };

    class search_results : public search_results_base<node_store>
    {
        using base_type = search_results_base<node_store>;
        using base_type::m_store;

    public:
        iterator begin();
        iterator end();
    };

    template<typename _SelfIter, typename _StoreIter, typename _ValueT>
    class iterator_base
    {
    public:
        using store_iterator_type = _StoreIter;
        using self_iterator_type = _SelfIter;

    protected:
        store_iterator_type m_pos;

        iterator_base(store_iterator_type pos);

    public:
        // iterator traits
        using value_type = _ValueT;
        using pointer = value_type*;
        using reference = value_type&;
        using difference_type = std::ptrdiff_t;
        using iterator_category = std::bidirectional_iterator_tag;

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

        self_iterator_type& operator++();
        self_iterator_type operator++(int);

        self_iterator_type& operator--();
        self_iterator_type operator--(int);

        const extent_type& extent() const;
        size_t depth() const;
    };

    class const_iterator
        : public iterator_base<
              const_iterator, typename const_search_results::store_type::const_iterator, const rtree::value_type>
    {
        using base_type = iterator_base<
            const_iterator, typename const_search_results::store_type::const_iterator, const rtree::value_type>;

        using store_type = typename const_search_results::store_type;
        using base_type::m_pos;
        using typename base_type::store_iterator_type;

        friend class rtree;

        const_iterator(store_iterator_type pos);

    public:
        using typename base_type::value_type;

        value_type& operator*() const
        {
            return static_cast<const value_node*>(m_pos->ns->node_ptr)->value;
        }

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

    class iterator : public iterator_base<iterator, typename search_results::store_type::iterator, rtree::value_type>
    {
        using base_type = iterator_base<iterator, typename search_results::store_type::iterator, rtree::value_type>;

        using store_type = typename const_search_results::store_type;
        using base_type::m_pos;
        using typename base_type::store_iterator_type;

        friend class rtree;

        iterator(store_iterator_type pos);

    public:
        using typename base_type::value_type;

        value_type& operator*()
        {
            return static_cast<value_node*>(m_pos->ns->node_ptr)->value;
        }

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

    class bulk_loader
    {
        dir_store_type m_store;

    public:
        bulk_loader();

        void insert(const extent_type& extent, value_type&& value);
        void insert(const extent_type& extent, const value_type& value);

        void insert(const point_type& position, value_type&& value);
        void insert(const point_type& position, const value_type& value);

        rtree pack();

    private:
        void pack_level(dir_store_type& store, size_t depth);

        void insert_impl(const extent_type& extent, value_type&& value);
        void insert_impl(const extent_type& extent, const value_type& value);
    };

    rtree();
    rtree(rtree&& other);
    rtree(const rtree& other);

private:
    rtree(node_store&& root); // only used internally by bulk_loader.

public:
    ~rtree();

    rtree& operator=(const rtree& other);

    rtree& operator=(rtree&& other);

    void insert(const extent_type& extent, value_type&& value);

    void insert(const extent_type& extent, const value_type& value);

    void insert(const point_type& position, value_type&& value);

    void insert(const point_type& position, const value_type& value);

    const_search_results search(const point_type& pt, search_type st) const;

    search_results search(const point_type& pt, search_type st);

    const_search_results search(const extent_type& extent, search_type st) const;

    search_results search(const extent_type& extent, search_type st);

    void erase(const_iterator pos);

    void erase(iterator pos);

    const extent_type& extent() const;

    bool empty() const;

    size_t size() const;

    void swap(rtree& other);

    void clear();

    template<typename FuncT>
    void walk(FuncT func) const;

    void check_integrity(const integrity_check_properties& props) const;

    std::string export_tree(export_tree_type mode) const;

private:
    void insert_impl(const point_type& start, const point_type& end, value_type&& value);
    void insert_impl(const point_type& start, const point_type& end, const value_type& value);

    void erase_impl(const node_store* ns, size_t depth);

    detail::rtree::ptr_to_string<const node_store*> build_ptr_to_string_map() const;

    std::string export_tree_formatted() const;
    std::string export_tree_extent_as_obj() const;
    std::string export_tree_extent_as_svg() const;

    void insert(node_store&& ns, std::unordered_set<size_t>* reinserted_depths);
    void insert_dir(node_store&& ns, size_t max_depth);

    void split_node(node_store* ns);

    void perform_forced_reinsertion(node_store* ns, std::unordered_set<size_t>& reinserted_depth);

    void sort_dir_store_by_split_dimension(dir_store_type& children);

    void sort_dir_store_by_dimension(size_t dim, dir_store_type& children);

    size_t pick_optimal_distribution(dir_store_type& children) const;

    insertion_point find_leaf_directory_node_for_insertion(const extent_type& bb);
    node_store* find_nonleaf_directory_node_for_insertion(const extent_type& bb, size_t max_depth);

    template<typename FuncT>
    void descend_with_func(FuncT func) const;

    using search_condition_type = std::function<bool(const node_store&)>;

    template<typename _ResT>
    void search_descend(
        size_t depth, const search_condition_type& dir_cond, const search_condition_type& value_cond,
        typename _ResT::node_store_type& ns, _ResT& results) const;

    void shrink_tree_upward(node_store* ns, const extent_type& bb_affected);

private:
    node_store m_root;
};

} // namespace mdds

#include "rtree_def.inl"

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