kdTreeTraits.h

/*
 This Source Code Form is subject to the terms of the Mozilla Public
 License, v. 2.0. If a copy of the MPL was not distributed with this
 file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/
 
#pragma once
 
#include "../defines.h"
#include "../../Common/Macro.h"
 
#include <cstddef>
 
#include <Eigen/Geometry>
 
namespace Ponca
{
#ifndef PARSED_WITH_DOXYGEN
    namespace internal
    {
        constexpr int clz(unsigned int value)
        {
#    if PONCA_HAS_BUILTIN_CLZ
            return __builtin_clz(value);
#    else
            if (value == 0)
            {
                return 0;
            }
 
            unsigned int msb_mask = 1 << (sizeof(unsigned int) * 8 - 1);
            int count             = 0;
            for (; (value & msb_mask) == 0; value <<= 1, ++count)
            {
            }
            return count;
#    endif
        }
    } // namespace internal
#endif
 
    template <typename NodeIndex, typename Scalar, int DIM>
    struct KdTreeDefaultInnerNode
    {
    private:
        enum
        {
            // The minimum bit width required to store the split dimension.
            //
            // Equal to the index of DIM's most significant bit (starting at 1), e.g.:
            // With DIM = 4,
            //             -------------
            // DIM =    0b | 1 | 0 | 0 |
            //             -------------
            // Bit index =  #3  #2  #1
            //
            // The MSB has an index of 3, so we store the dimension on 3 bits.
            DIM_BITS = sizeof(unsigned int) * 8 - internal::clz((unsigned int)DIM),
        };
 
        // The node stores bitfields as unsigned indices.
        using UIndex = typename std::make_unsigned<NodeIndex>::type;
 
    public:
        enum
        {
            INDEX_BITS = sizeof(UIndex) * 8 - DIM_BITS,
        };
 
        Scalar split_value{0};
        UIndex first_child_id : INDEX_BITS;
        UIndex split_dim : DIM_BITS;
    };
 
    template <typename Index, typename Size>
    struct KdTreeDefaultLeafNode
    {
        Index start{0};
        Size size{0};
    };
 
    template <typename Index, typename NodeIndex, typename DataPoint, typename LeafSize = Index,
              typename _InnerNodeType = KdTreeDefaultInnerNode<NodeIndex, typename DataPoint::Scalar, DataPoint::Dim>,
              typename _LeafNodeType  = KdTreeDefaultLeafNode<Index, LeafSize>>
    class KdTreeCustomizableNode
    {
    private:
        using Scalar    = typename DataPoint::Scalar;
        using InnerType = _InnerNodeType;
        using LeafType  = _LeafNodeType;
 
    public:
        enum : std::size_t
        {
            MAX_COUNT = std::size_t(2) << InnerType::INDEX_BITS,
        };
 
        using AabbType = Eigen::AlignedBox<Scalar, DataPoint::Dim>;
 
        PONCA_MULTIARCH KdTreeCustomizableNode() = default;
 
        PONCA_MULTIARCH constexpr KdTreeCustomizableNode(KdTreeCustomizableNode&& n) : m_is_leaf(n.m_is_leaf)
        {
            if (m_is_leaf)
            {
                data.m_leaf = n.data.m_leaf;
            }
            else
            {
                data.m_inner = n.data.m_inner;
            }
        }
 
        PONCA_MULTIARCH constexpr KdTreeCustomizableNode& operator=(KdTreeCustomizableNode&& n)
        {
            if (&n != this)
            {
                m_is_leaf = n.m_is_leaf;
                if (m_is_leaf)
                {
                    data.m_leaf = n.data.m_leaf;
                }
                else
                {
                    data.m_inner = n.data.m_inner;
                }
            }
 
            return *this;
        }
 
        PONCA_MULTIARCH constexpr KdTreeCustomizableNode(const KdTreeCustomizableNode& n) : m_is_leaf(n.m_is_leaf)
        {
            if (n.m_is_leaf)
            {
                data.m_leaf = n.data.m_leaf;
            }
            else
            {
                data.m_inner = n.data.m_inner;
            }
        }
 
        PONCA_MULTIARCH constexpr KdTreeCustomizableNode& operator=(const KdTreeCustomizableNode& n)
        {
            if (&n != this)
            {
                m_is_leaf = n.m_is_leaf;
                if (n.m_is_leaf)
                {
                    data.m_leaf = n.data.m_leaf;
                }
                else
                {
                    data.m_inner = n.data.m_inner;
                }
            }
 
            return *this;
        }
 
        PONCA_MULTIARCH_HOST ~KdTreeCustomizableNode() {}
 
        PONCA_MULTIARCH [[nodiscard]] bool is_leaf() const { return m_is_leaf; }
        PONCA_MULTIARCH void set_is_leaf(bool is_leaf) { m_is_leaf = is_leaf; }
 
        PONCA_MULTIARCH void configure_range(Index start, Index size, const AabbType& aabb)
        {
            if (m_is_leaf)
            {
                data.m_leaf.start = start;
                data.m_leaf.size  = (LeafSize)size;
            }
        }
 
        PONCA_MULTIARCH void configure_inner(Scalar split_value, Index first_child_id, Index split_dim)
        {
            if (!m_is_leaf)
            {
                data.m_inner.split_value    = split_value;
                data.m_inner.first_child_id = first_child_id;
                data.m_inner.split_dim      = split_dim;
            }
        }
 
        PONCA_MULTIARCH [[nodiscard]] Index leaf_start() const { return data.m_leaf.start; }
 
        PONCA_MULTIARCH [[nodiscard]] LeafSize leaf_size() const { return data.m_leaf.size; }
 
        PONCA_MULTIARCH [[nodiscard]] Scalar inner_split_value() const { return data.m_inner.split_value; }
 
        PONCA_MULTIARCH [[nodiscard]] int inner_split_dim() const { return (int)data.m_inner.split_dim; }
 
        PONCA_MULTIARCH [[nodiscard]] Index inner_first_child_id() const { return (Index)data.m_inner.first_child_id; }
 
    protected:
        PONCA_MULTIARCH [[nodiscard]] inline LeafType& getAsLeaf() { return data.m_leaf; }
        PONCA_MULTIARCH [[nodiscard]] inline InnerType& getAsInner() { return data.m_inner; }
        PONCA_MULTIARCH [[nodiscard]] inline const LeafType& getAsLeaf() const { return data.m_leaf; }
        PONCA_MULTIARCH [[nodiscard]] inline const InnerType& getAsInner() const { return data.m_inner; }
 
    private:
        bool m_is_leaf{true};
        union Data {
            // We need an explicit constructor here, see https://stackoverflow.com/a/70428826
            constexpr Data() : m_leaf() {}
 
            ~Data() {}
            LeafType m_leaf;
            InnerType m_inner;
        };
        Data data;
    };
 
    template <typename Index, typename NodeIndex, typename DataPoint, typename LeafSize = Index>
    struct KdTreeDefaultNode
        : public KdTreeCustomizableNode<Index, NodeIndex, DataPoint, LeafSize,
                                        KdTreeDefaultInnerNode<NodeIndex, typename DataPoint::Scalar, DataPoint::Dim>,
                                        KdTreeDefaultLeafNode<Index, LeafSize>>
    {
        using Base =
            KdTreeCustomizableNode<Index, NodeIndex, DataPoint, LeafSize,
                                   KdTreeDefaultInnerNode<NodeIndex, typename DataPoint::Scalar, DataPoint::Dim>,
                                   KdTreeDefaultLeafNode<Index, LeafSize>>;
    };
 
    template <typename _DataPoint, template <typename /*Index*/, typename /*NodeIndex*/, typename /*DataPoint*/,
                                             typename /*LeafSize*/> typename _NodeType = KdTreeDefaultNode>
    struct KdTreeDefaultTraits
    {
        enum
        {
            MAX_DEPTH    = 32,
            MAX_KNN_SIZE = 128 
        };
 
        using DataPoint    = _DataPoint;
        using IndexType    = int;
        using LeafSizeType = unsigned short;
 
        // Containers
        using PointContainer = std::vector<DataPoint>;
        using IndexContainer = std::vector<IndexType>;
 
        // Nodes
        using NodeIndexType = std::size_t;
        using NodeType      = _NodeType<IndexType, NodeIndexType, DataPoint, LeafSizeType>;
        using NodeContainer = std::vector<NodeType>;
    };
 
    template <typename _DataPoint, template <typename /*Index*/, typename /*NodeIndex*/, typename /*DataPoint*/,
                                             typename /*LeafSize*/> typename _NodeType = Ponca::KdTreeDefaultNode>
    struct KdTreePointerTraits
    {
        enum
        {
            MAX_DEPTH    = 32,
            MAX_KNN_SIZE = 128 
        };
 
        using DataPoint    = _DataPoint;
        using IndexType    = int;
        using LeafSizeType = unsigned short;
 
        // Containers
        using PointContainer = DataPoint*;
        using IndexContainer = IndexType*;
 
        // Nodes
        using NodeIndexType = std::size_t;
        using NodeType      = _NodeType<IndexType, NodeIndexType, DataPoint, LeafSizeType>;
        using NodeContainer = NodeType*;
    };
} // namespace Ponca