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
    }
}
#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,
    };
 
    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,
    };
 
    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