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: */