key.hpp

 
#ifndef UFO_CONTAINER_TREE_KEY_HPP
#define UFO_CONTAINER_TREE_KEY_HPP
 
// UFO
#include <ufo/morton/morton.hpp>
#include <ufo/numeric/vec.hpp>
 
// STL
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <limits>
#include <ostream>
 
namespace ufo
{
template <std::size_t Dim>
class TreeKey : public Vec<Dim, std::uint32_t>
{
 public:
    using key_type   = std::uint32_t;
    using Key        = Vec<Dim, key_type>;
    using depth_type = key_type;
    using size_type  = std::size_t;
 
 public:
    /**************************************************************************************
    |                                                                                     |
    |                                    Constructors                                     |
    |                                                                                     |
    **************************************************************************************/
 
    constexpr TreeKey() noexcept               = default;
    constexpr TreeKey(TreeKey const&) noexcept = default;
 
    constexpr TreeKey(Key key, depth_type depth) : Key(key), depth_(depth) {}
 
    constexpr explicit TreeKey(Key key) : TreeKey(key, 0) {}
 
    /**************************************************************************************
    |                                                                                     |
    |                                 Assignment operator                                 |
    |                                                                                     |
    **************************************************************************************/
 
    constexpr TreeKey& operator=(TreeKey const&) noexcept = default;
 
    /**************************************************************************************
    |                                                                                     |
    |                                      Capacity                                       |
    |                                                                                     |
    **************************************************************************************/
 
    [[nodiscard]] static constexpr depth_type maxDepth() noexcept
    {
        // All the time you have to leave the space.
        // Shifting with `maxDepth()` causes problems when
        // `std::numeric_limits<key_type>::digits <= maxDepth()` because you are trying to
        // shift more bits than are allowed and has a well-defined behaviour in C++.
        // Therefore, we have this check; otherwise, would cause "shift count overflow".
        return std::numeric_limits<key_type>::digits - 1;
    }
 
    /**************************************************************************************
    |                                                                                     |
    |                                     Operations                                      |
    |                                                                                     |
    **************************************************************************************/
 
    [[nodiscard]] constexpr bool valid() const noexcept { return maxDepth() >= depth_; }
 
    [[nodiscard]] constexpr depth_type depth() const noexcept { return depth_; }
 
    [[nodiscard]] constexpr TreeKey toDepth(depth_type depth) const
    {
        assert(maxDepth() >= depth);
 
        TreeKey ret = *this;
        ret.setDepth(depth);
        return ret;
    }
 
    constexpr void setDepth(depth_type depth) noexcept
    {
        assert(maxDepth() >= depth);
 
        if (depth_ > depth) {
            *this <<= depth_ - depth;
        } else {
            *this >>= depth - depth_;
        }
        this->depth_ = depth;
    }
 
    // TODO: Add methods from Code
 
    [[nodiscard]] constexpr key_type offset() const noexcept { return offset(depth_); }
 
    [[nodiscard]] constexpr key_type offset(depth_type depth) const noexcept
    {
        assert(maxDepth() >= depth);
        assert(depth_ <= depth);
 
        auto     v   = (*this >> (depth - depth_)) & key_type(1);
        key_type ret = v[0];
        for (std::size_t i = 1; Dim > i; ++i) {
            ret |= v[i] << i;
        }
        return ret;
    }
 
    [[nodiscard]] constexpr TreeKey parent() const
    {
        assert(maxDepth() > depth_);
 
        TreeKey ret = *this;
        ret >>= 1;
        ++ret.depth_;
        return ret;
    }
 
    void swap(TreeKey& other) noexcept
    {
        using std::swap;
        static_cast<Key&>(*this).swap(static_cast<Key&>(other));
        swap(depth_, other.depth_);
    }
 
 private:
    depth_type depth_{};
};
 
using BinaryKey = TreeKey<1>;
using QuadKey   = TreeKey<2>;
using OctKey    = TreeKey<3>;
using HexKey    = TreeKey<4>;
 
template <std::size_t Dim>
std::ostream& operator<<(std::ostream& out, TreeKey<Dim> const& key)
{
    return out << static_cast<typename TreeKey<Dim>::Key const&>(key)
               << " d: " << key.depth();
}
 
template <std::size_t Dim>
void swap(TreeKey<Dim>& lhs, TreeKey<Dim>& rhs) noexcept
{
    lhs.swap(rhs);
}
 
/**************************************************************************************
|                                                                                     |
|                                       Compare                                       |
|                                                                                     |
**************************************************************************************/
 
template <std::size_t Dim>
[[nodiscard]] constexpr bool operator==(TreeKey<Dim> const& lhs,
                                        TreeKey<Dim> const& rhs) noexcept
{
    using Key = typename TreeKey<Dim>::Key;
    return static_cast<Key const&>(lhs) == static_cast<Key const&>(rhs) &&
           lhs.depth() == rhs.depth();
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr bool operator!=(TreeKey<Dim> const& lhs,
                                        TreeKey<Dim> const& rhs) noexcept
{
    return !(lhs == rhs);
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr bool operator<(TreeKey<Dim> const& lhs,
                                       TreeKey<Dim> const& rhs) noexcept
{
    for (std::size_t i{}; Dim > i; ++i) {
        if (lhs[i] < rhs[i]) {
            return true;
        } else if (lhs[i] > rhs[i]) {
            return false;
        }
    }
    return false;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr bool operator<=(TreeKey<Dim> const& lhs,
                                        TreeKey<Dim> const& rhs) noexcept
{
    for (std::size_t i{}; Dim > i; ++i) {
        if (lhs[i] < rhs[i]) {
            return true;
        } else if (lhs[i] > rhs[i]) {
            return false;
        }
    }
    return true;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr bool operator>(TreeKey<Dim> const& lhs,
                                       TreeKey<Dim> const& rhs) noexcept
{
    for (std::size_t i{}; Dim > i; ++i) {
        if (lhs[i] > rhs[i]) {
            return true;
        } else if (lhs[i] < rhs[i]) {
            return false;
        }
    }
    return false;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr bool operator>=(TreeKey<Dim> const& lhs,
                                        TreeKey<Dim> const& rhs) noexcept
{
    for (std::size_t i{}; Dim > i; ++i) {
        if (lhs[i] > rhs[i]) {
            return true;
        } else if (lhs[i] < rhs[i]) {
            return false;
        }
    }
    return true;
}
 
/**************************************************************************************
|                                                                                     |
|                                  Binary operators                                   |
|                                                                                     |
**************************************************************************************/
 
template <std::size_t Dim>
[[nodiscard]] constexpr TreeKey<Dim> operator+(TreeKey<Dim> key, std::uint32_t v) noexcept
{
    key += v;
    return key;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr TreeKey<Dim> operator+(TreeKey<Dim>                   key,
                                               Vec<Dim, std::uint32_t> const& v) noexcept
{
    key += v;
    return key;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr TreeKey<Dim> operator+(std::uint32_t v, TreeKey<Dim> key) noexcept
{
    key += v;
    return key;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr TreeKey<Dim> operator+(Vec<Dim, std::uint32_t> const& v,
                                               TreeKey<Dim> key) noexcept
{
    key += v;
    return key;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr TreeKey<Dim> operator-(TreeKey<Dim> key, std::uint32_t v) noexcept
{
    key -= v;
    return key;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr TreeKey<Dim> operator-(TreeKey<Dim>                   key,
                                               Vec<Dim, std::uint32_t> const& v) noexcept
{
    key -= v;
    return key;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr TreeKey<Dim> operator-(std::uint32_t v, TreeKey<Dim> key) noexcept
{
    key -= v;
    return key;
}
 
template <std::size_t Dim>
[[nodiscard]] constexpr TreeKey<Dim> operator-(Vec<Dim, std::uint32_t> const& v,
                                               TreeKey<Dim> key) noexcept
{
    key -= v;
    return key;
}
}  // namespace ufo
 
namespace std
{
template <std::size_t Dim>
struct hash<ufo::TreeKey<Dim>> {
    std::size_t operator()(ufo::TreeKey<Dim> const& key) const
    {
        using Key = typename ufo::TreeKey<Dim>::Key;
        return static_cast<std::size_t>(
            ufo::Morton<Dim>::encode(static_cast<Key const&>(key)));
    }
};
}  // namespace std
 
#endif  // UFO_CONTAINER_TREE_KEY_HPP