EBGeometry_TreeBVH.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2022 Robert Marskar <robert.marskar@sintef.no>
//
// SPDX-License-Identifier: GPL-3.0-or-later
 
#ifndef EBGEOMETRY_TREEBVH_HPP
#define EBGEOMETRY_TREEBVH_HPP
 
// Std includes
#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <iterator>
#include <memory>
#include <type_traits>
#include <utility>
#include <vector>
 
// Our includes
#include "EBGeometry_Array.hpp"
#include "EBGeometry_BoundingVolumes.hpp"
#include "EBGeometry_Macros.hpp"
#include "EBGeometry_Math.hpp"
#include "EBGeometry_PackedBVH.hpp"
#include "EBGeometry_SFC.hpp"
#include "EBGeometry_Vec.hpp"
 
namespace EBGeometry {
 
namespace BVH {
 
namespace detail {
 
template <class P>
EBGEOMETRY_HOST
inline void
appendTreeLeaf(std::vector<P>& a_dst, const PrimitiveList<P>& a_leafPrims)
{
  // NB: do NOT reserve(a_dst.size() + a_leafPrims.size()) here. appendTreeLeaf is called once per
  // leaf as the whole tree is packed, and reserving to an each-time-slightly-larger *exact* size
  // defeats std::vector's geometric growth -- it forces a fresh reallocation (copying every
  // element already appended) on every leaf, making a whole build O(N^2) in the primitive count.
  // Plain push_back grows the buffer geometrically and keeps construction linear.
  for (const auto& p : a_leafPrims) {
    a_dst.push_back(*p);
  }
}
 
template <class List, size_t K>
EBGEOMETRY_HOST
inline void
requireNonEmptyPartitions(const Array<List, K>& a_partitions, const size_t a_numInput, const char* a_who)
{
  for (size_t k = 0; k < K; k++) {
    EBGEOMETRY_REQUIRE(!a_partitions[k].empty(),
                       "%s: the partitioner returned an empty partition (%zu of %zu, from %zu primitives). A leaf "
                       "must hold at least one primitive; a partitioner must return %zu non-empty partitions.",
                       a_who,
                       k,
                       K,
                       a_numInput,
                       K);
  }
}
 
template <class P>
EBGEOMETRY_HOST
inline void
appendAliased(std::vector<P>& a_dst, const std::shared_ptr<std::vector<P>>& a_block)
{
  if (a_dst.empty()) {
    a_dst = std::move(*a_block);
  }
  else {
    a_dst.insert(a_dst.end(), std::make_move_iterator(a_block->begin()), std::make_move_iterator(a_block->end()));
  }
}
 
inline void
requireLeafSetting(const char* a_who, const Construction a_construction, const ConstructionOptions& a_options) noexcept
{
  switch (a_construction) {
  case Construction::SAH:
  case Construction::CentroidSplit:
  case Construction::MidpointSplit: {
    EBGEOMETRY_REQUIRE(a_options.maxLeafSize > 0,
                       "%s: ConstructionOptions::maxLeafSize must be positive for a top-down construction",
                       a_who);
 
    return;
  }
  case Construction::Morton:
  case Construction::Nested:
  case Construction::Hilbert: {
    EBGEOMETRY_REQUIRE(
      a_options.targetLeafSize > 0,
      "%s: ConstructionOptions::targetLeafSize must be positive for a space-filling-curve construction",
      a_who);
 
    return;
  }
  case Construction::ClusterSAH:
  default: {
    return;
  }
  }
}
 
} // namespace detail
 
template <class X, size_t K>
auto EqualCounts = [](std::vector<X> a_primitives) -> Array<std::vector<X>, K> {
  static_assert(K >= 2, "EqualCounts<X, K>: branching factor K must be at least 2");
 
  EBGEOMETRY_EXPECT(!a_primitives.empty());
 
  const int length = a_primitives.size() / K;
  int       remain = a_primitives.size() % K;
 
  int begin = 0;
  int end   = 0;
 
  Array<std::vector<X>, K> chunks;
 
  for (size_t k = 0; k < K; k++) {
    end += (remain > 0) ? length + 1 : length;
    remain--;
 
    // Move the [begin, end) slice out -- the input is taken by value and not reused, so the elements
    // (e.g. shared_ptr primitive handles) transfer without copying or refcount churn.
    chunks[k] = std::vector<X>(std::make_move_iterator(a_primitives.begin() + begin),
                               std::make_move_iterator(a_primitives.begin() + end));
 
    begin = end;
  }
 
  return chunks;
};
 
template <class T, class P, class BV, size_t K>
auto PrimitiveCentroidPartitioner = [](PrimAndBVList<P, BV> a_primsAndBVs) -> Array<PrimAndBVList<P, BV>, K> {
  EBGEOMETRY_EXPECT(!a_primsAndBVs.empty());
 
  Vec3T<T> lo = +Vec3T<T>::max();
  Vec3T<T> hi = -Vec3T<T>::max();
 
  for (const auto& pbv : a_primsAndBVs) {
    lo = min(lo, pbv.first->getCentroid());
    hi = max(hi, pbv.first->getCentroid());
  }
 
  const size_t splitDir = (hi - lo).maxDir(true);
 
  // The input is taken by value; sort it in place (no working copy) and move it into EqualCounts.
  std::sort(a_primsAndBVs.begin(),
            a_primsAndBVs.end(),
            [splitDir](const PrimAndBV<P, BV>& pbv1, const PrimAndBV<P, BV>& pbv2) -> bool {
              return pbv1.first->getCentroid(splitDir) < pbv2.first->getCentroid(splitDir);
            });
 
  return BVH::EqualCounts<PrimAndBV<P, BV>, K>(std::move(a_primsAndBVs));
};
 
template <class T, class P, class BV, size_t K>
auto BVCentroidPartitioner = [](PrimAndBVList<P, BV> a_primsAndBVs) -> Array<PrimAndBVList<P, BV>, K> {
  EBGEOMETRY_EXPECT(!a_primsAndBVs.empty());
 
  Vec3T<T> lo = +Vec3T<T>::max();
  Vec3T<T> hi = -Vec3T<T>::max();
 
  for (const auto& pbv : a_primsAndBVs) {
    lo = min(lo, pbv.second.getCentroid());
    hi = max(hi, pbv.second.getCentroid());
  }
 
  const size_t splitDir = (hi - lo).maxDir(true);
 
  // The input is taken by value; sort it in place (no working copy) and move it into EqualCounts.
  std::sort(a_primsAndBVs.begin(),
            a_primsAndBVs.end(),
            [splitDir](const PrimAndBV<P, BV>& pbv1, const PrimAndBV<P, BV>& pbv2) -> bool {
              return pbv1.second.getCentroid()[splitDir] < pbv2.second.getCentroid()[splitDir];
            });
 
  return BVH::EqualCounts<PrimAndBV<P, BV>, K>(std::move(a_primsAndBVs));
};
 
template <class T, class P, class BV>
inline size_t
SAH2WaySplit(PrimAndBVList<P, BV>& a_list,
             const size_t          a_begin,
             const size_t          a_end,
             const bool            a_longestAxisOnly = false) noexcept
{
  static_assert(std::is_same_v<BV, EBGeometry::BoundingVolumes::AABBT<T>>, "SAH2WaySplit requires BV == AABBT<T>");
 
  EBGEOMETRY_EXPECT(a_begin < a_end);
 
  constexpr int BINS = 32;
 
  const size_t N = a_end - a_begin;
 
  // Centroid bounding box
  Vec3T<T> clo = +Vec3T<T>::max();
  Vec3T<T> chi = -Vec3T<T>::max();
 
  for (size_t i = a_begin; i < a_end; i++) {
    const auto& c = a_list[i].second.getCentroid();
 
    clo = min(clo, c);
    chi = max(chi, c);
  }
 
  T   bestCost  = Math::Limits<T>::max();
  T   bestPlane = T(0);
  int bestAxis  = -1;
 
  Vec3T<T> binLo[BINS];
  Vec3T<T> binHi[BINS];
 
  int binCnt[BINS];
 
  T   leftArea[BINS - 1];
  int leftCnt[BINS - 1];
 
  const int firstAxis = a_longestAxisOnly ? (chi - clo).maxDir(true) : 0;
  const int lastAxis  = a_longestAxisOnly ? firstAxis : 2;
 
  for (int axis = firstAxis; axis <= lastAxis; axis++) {
    const T lo  = clo[axis];
    const T hi  = chi[axis];
    const T ext = hi - lo;
 
    // An extent so small that BINS / ext overflows cannot be binned; leave that axis out.
    const T scale = T(BINS) / ext;
 
    if (!(ext > T(0)) || !std::isfinite(scale)) {
      continue;
    }
 
    for (int b = 0; b < BINS; b++) {
      binLo[b]  = Vec3T<T>::max();
      binHi[b]  = -Vec3T<T>::max();
      binCnt[b] = 0;
    }
 
    for (size_t i = a_begin; i < a_end; i++) {
      // Clamped in floating point before the conversion, which is undefined past int's range.
      const T   x = (a_list[i].second.getCentroid()[axis] - lo) * scale;
      const int b = (x > T(0)) ? ((x < T(BINS - 1)) ? static_cast<int>(x) : BINS - 1) : 0;
      binLo[b]    = min(binLo[b], a_list[i].second.getLowCorner());
      binHi[b]    = max(binHi[b], a_list[i].second.getHighCorner());
      binCnt[b]   = binCnt[b] + 1;
    }
 
    // Left prefix: accumulated AABB and count for bins [0..b]
    Vec3T<T> rlo  = Vec3T<T>::max();
    Vec3T<T> rhi  = -Vec3T<T>::max();
    int      rcnt = 0;
 
    for (int b = 0; b < BINS - 1; b++) {
      if (binCnt[b] > 0) {
        rlo = min(rlo, binLo[b]);
        rhi = max(rhi, binHi[b]);
      }
 
      rcnt        = rcnt + binCnt[b];
      leftArea[b] = (rcnt > 0) ? BV(rlo, rhi).getArea() : T(0);
      leftCnt[b]  = rcnt;
    }
 
    // Right suffix sweep; evaluate SAH cost at each candidate split boundary
    Vec3T<T> rrlo  = Vec3T<T>::max();
    Vec3T<T> rrhi  = -Vec3T<T>::max();
    int      rrcnt = 0;
 
    for (int b = BINS - 1; b >= 1; b--) {
      if (binCnt[b] > 0) {
        rrlo = min(rrlo, binLo[b]);
        rrhi = max(rrhi, binHi[b]);
      }
 
      rrcnt += binCnt[b];
 
      if (leftCnt[b - 1] > 0 && rrcnt > 0) {
        const T cost = leftArea[b - 1] * T(leftCnt[b - 1]) + BV(rrlo, rrhi).getArea() * T(rrcnt);
 
        if (cost < bestCost) {
          bestCost  = cost;
          bestAxis  = axis;
          bestPlane = lo + T(b) / scale;
        }
      }
    }
  }
 
  // All axes degenerate — equal split
  if (bestAxis < 0) {
    return a_begin + N / 2;
  }
 
  auto mid = std::partition(
    a_list.begin() + a_begin, a_list.begin() + a_end, [bestAxis, bestPlane](const PrimAndBV<P, BV>& pbv) noexcept {
      return pbv.second.getCentroid()[bestAxis] < bestPlane;
    });
 
  const size_t splitIdx = static_cast<size_t>(std::distance(a_list.begin(), mid));
 
  // Guard: if everything ended up on one side, fall back to equal split
  return (splitIdx == a_begin || splitIdx == a_end) ? a_begin + N / 2 : splitIdx;
}
 
template <class T, class P, class BV>
inline void
SAHKWaySplit(PrimAndBVList<P, BV>&                   a_list,
             const size_t                            a_begin,
             const size_t                            a_end,
             const size_t                            a_K,
             std::vector<std::pair<size_t, size_t>>& a_groups,
             const bool                              a_longestAxisOnly = false)
{
  static_assert(std::is_same_v<BV, EBGeometry::BoundingVolumes::AABBT<T>>, "SAHKWaySplit requires BV == AABBT<T>");
 
  if (a_K <= 1 || a_begin >= a_end) {
    a_groups.emplace_back(a_begin, a_end);
 
    return;
  }
 
  const size_t K1 = a_K / 2;
  const size_t K2 = a_K - K1;
 
  // Clamp the SAH split so the left half has >= K1 elements and the right has >= K2.
  // This guarantees neither recursive call receives an under-populated range, which
  // would cause empty sub-lists to reach the TreeBVH constructor and crash on
  // AABBT(vector::front()) when the vector is empty.
  // The clamp is valid whenever a_end - a_begin >= a_K = K1 + K2.
  const size_t rawMid = SAH2WaySplit<T, P, BV>(a_list, a_begin, a_end, a_longestAxisOnly);
  const size_t mid    = Math::max(a_begin + K1, Math::min(a_end - K2, rawMid));
 
  SAHKWaySplit<T, P, BV>(a_list, a_begin, mid, K1, a_groups, a_longestAxisOnly);
  SAHKWaySplit<T, P, BV>(a_list, mid, a_end, K2, a_groups, a_longestAxisOnly);
}
 
template <class T, class P, class BV, size_t K, bool LongestAxisOnly = false>
auto BinnedSAHPartitioner = [](PrimAndBVList<P, BV> a_primsAndBVs) -> Array<PrimAndBVList<P, BV>, K> {
  EBGEOMETRY_EXPECT(!a_primsAndBVs.empty());
 
  // The input is taken by value; partition it in place (no working copy). SAHKWaySplit reorders it
  // and yields K [begin, end) index ranges, which we move out (disjoint, each moved once).
  std::vector<std::pair<size_t, size_t>> groups;
  groups.reserve(K);
 
  SAHKWaySplit<T, P, BV>(a_primsAndBVs, 0, a_primsAndBVs.size(), K, groups, LongestAxisOnly);
 
  Array<PrimAndBVList<P, BV>, K> result;
  for (size_t k = 0; k < K; k++) {
    const auto [b, e] = groups[k];
    result[k]         = PrimAndBVList<P, BV>(std::make_move_iterator(a_primsAndBVs.begin() + b),
                                     std::make_move_iterator(a_primsAndBVs.begin() + e));
  }
 
  return result;
};
 
template <class T, class P, class BV>
inline size_t
Midpoint2WaySplit(PrimAndBVList<P, BV>& a_list, const size_t a_begin, const size_t a_end) noexcept
{
  EBGEOMETRY_EXPECT(a_begin < a_end);
 
  const size_t N = a_end - a_begin;
 
  Vec3T<T> lo = +Vec3T<T>::max();
  Vec3T<T> hi = -Vec3T<T>::max();
 
  for (size_t i = a_begin; i < a_end; i++) {
    const auto& c = a_list[i].second.getCentroid();
 
    lo = min(lo, c);
    hi = max(hi, c);
  }
 
  const size_t splitDir = (hi - lo).maxDir(true);
  const T      midpoint = T(0.5) * (lo[splitDir] + hi[splitDir]);
 
  auto mid = std::partition(
    a_list.begin() + a_begin, a_list.begin() + a_end, [splitDir, midpoint](const PrimAndBV<P, BV>& pbv) noexcept {
      return pbv.second.getCentroid()[splitDir] < midpoint;
    });
 
  const size_t splitIdx = static_cast<size_t>(std::distance(a_list.begin(), mid));
 
  // Guard: if every centroid ended up on one side (e.g. all coincide on the split axis, or the
  // distribution is skewed entirely to one side of the midpoint), fall back to an equal-count
  // split so neither resulting group is ever empty.
  return (splitIdx == a_begin || splitIdx == a_end) ? a_begin + N / 2 : splitIdx;
}
 
template <class T, class P, class BV>
inline void
MidpointKWaySplit(PrimAndBVList<P, BV>&                   a_list,
                  const size_t                            a_begin,
                  const size_t                            a_end,
                  const size_t                            a_K,
                  std::vector<std::pair<size_t, size_t>>& a_groups)
{
  if (a_K <= 1 || a_begin >= a_end) {
    a_groups.emplace_back(a_begin, a_end);
 
    return;
  }
 
  const size_t K1 = a_K / 2;
  const size_t K2 = a_K - K1;
 
  // Clamp so the left half has >= K1 elements and the right has >= K2, guaranteeing neither
  // recursive call receives an under-populated range (see SAHKWaySplit's identical clamp for why:
  // an empty sub-list reaching the TreeBVH constructor crashes on AABBT(vector::front())).
  const size_t rawMid = Midpoint2WaySplit<T, P, BV>(a_list, a_begin, a_end);
  const size_t mid    = Math::max(a_begin + K1, Math::min(a_end - K2, rawMid));
 
  MidpointKWaySplit<T, P, BV>(a_list, a_begin, mid, K1, a_groups);
  MidpointKWaySplit<T, P, BV>(a_list, mid, a_end, K2, a_groups);
}
 
template <class T, class P, class BV, size_t K>
auto MidpointPartitioner = [](PrimAndBVList<P, BV> a_primsAndBVs) -> Array<PrimAndBVList<P, BV>, K> {
  EBGEOMETRY_EXPECT(!a_primsAndBVs.empty());
 
  // The input is taken by value; partition it in place (no working copy). MidpointKWaySplit reorders
  // it and yields K [begin, end) index ranges, which we move out (disjoint, each moved once). This
  // is the crux of the partitioner's cost: the midpoint split itself is a couple of std::partition
  // passes, so moving rather than copying the lists is what keeps it cheap.
  std::vector<std::pair<size_t, size_t>> groups;
  groups.reserve(K);
 
  MidpointKWaySplit<T, P, BV>(a_primsAndBVs, 0, a_primsAndBVs.size(), K, groups);
 
  Array<PrimAndBVList<P, BV>, K> result;
  for (size_t k = 0; k < K; k++) {
    const auto [b, e] = groups[k];
    result[k]         = PrimAndBVList<P, BV>(std::make_move_iterator(a_primsAndBVs.begin() + b),
                                     std::make_move_iterator(a_primsAndBVs.begin() + e));
  }
 
  return result;
};
 
template <class T, class P, class BV, size_t K>
auto DefaultLeafPredicate =
  [](const BVH::TreeBVH<T, P, BV, K>& a_node) noexcept -> bool { return (a_node.getPrimitives()).size() < K; };
 
template <class T, class P, class BV, size_t K>
class TreeBVH
{
  static_assert(std::is_floating_point_v<T>, "TreeBVH: T must be a floating-point type");
  static_assert(K >= 2, "TreeBVH: branching factor K must be at least 2");
 
public:
  using PrimitiveList = BVH::PrimitiveList<P>;
 
  using Vec3 = Vec3T<T>;
 
  using Node = TreeBVH<T, P, BV, K>;
 
  using NodePtr = std::shared_ptr<Node>;
 
  using Partitioner = BVH::Partitioner<P, BV, K>;
 
  using LeafPredicate = BVH::LeafPredicate<T, P, BV, K>;
 
  TreeBVH();
 
  TreeBVH(const std::vector<PrimAndBV<P, BV>>& a_primsAndBVs);
 
  TreeBVH(std::vector<PrimAndBV<P, BV>>&& a_primsAndBVs);
 
  ~TreeBVH() noexcept = default;
 
  TreeBVH(const TreeBVH& a_other) = delete;
 
  TreeBVH&
  operator=(const TreeBVH& a_other) = delete;
 
  TreeBVH(TreeBVH&& a_other) noexcept = default;
 
  TreeBVH&
  operator=(TreeBVH&& a_other) noexcept = default;
 
  inline void
  topDownSortAndPartition(const Partitioner&   a_partitioner = BVCentroidPartitioner<T, P, BV, K>,
                          const LeafPredicate& a_stopCrit    = DefaultLeafPredicate<T, P, BV, K>);
 
  template <typename S>
  inline void
  bottomUpSortAndPartition(size_t a_targetLeafSize = 1);
 
  [[nodiscard]] inline bool
  isLeaf() const noexcept;
 
  [[nodiscard]] inline bool
  isPartitioned() const noexcept;
 
  [[nodiscard]] inline const BV&
  getBoundingVolume() const noexcept;
 
  [[nodiscard]] inline const PrimitiveList&
  getPrimitives() const noexcept;
 
  [[nodiscard]] inline const std::vector<BV>&
  getBoundingVolumes() const noexcept;
 
  [[nodiscard]] inline T
  getDistanceToBoundingVolume(const Vec3& a_point) const noexcept;
 
  [[nodiscard]] inline const Array<std::shared_ptr<TreeBVH<T, P, BV, K>>, K>&
  getChildren() const noexcept;
 
  [[nodiscard]] inline std::shared_ptr<TreeBVH<T, P, BV, K>>
  deepCopy() const;
 
  template <class NodeKey>
  inline void
  traverse(const BVH::LeafEvaluator<P>&               a_leafEvaluator,
           const BVH::PrunePredicate<Node, NodeKey>&  a_prunePredicate,
           const BVH::ChildOrderer<Node, NodeKey, K>& a_childOrderer,
           const BVH::NodeKeyFactory<Node, NodeKey>&  a_nodeKeyFactory) const;
 
  template <class BVConstructor>
  inline const BV&
  refit(const BVConstructor& a_bvConstructor);
 
  [[nodiscard]] EBGEOMETRY_HOST
  inline PackedBVH<T, P, K>
  pack(Pool& a_pool) const;
 
  template <class Q, class Converter>
  [[nodiscard]] EBGEOMETRY_HOST
  inline PackedBVH<T, Q, K>
  packWith(Pool& a_pool, Converter&& a_converter) const;
 
protected:
  BV m_boundingVolume;
 
  bool m_partitioned;
 
  std::vector<std::shared_ptr<const P>> m_primitives;
 
  std::vector<BV> m_boundingVolumes;
 
  Array<std::shared_ptr<TreeBVH<T, P, BV, K>>, K> m_children;
 
  inline PrimitiveList&
  getPrimitives() noexcept;
 
  inline std::vector<BV>&
  getBoundingVolumes() noexcept;
 
  inline void
  setChildren(const Array<std::shared_ptr<TreeBVH<T, P, BV, K>>, K>& a_children);
};
 
} // namespace BVH
 
} // namespace EBGeometry
 
#include "EBGeometry_TreeBVHImplem.hpp"
 
// pack() and packWith() call the PackedBVH constructor that takes a TreeBVH, defined there.
#include "EBGeometry_BVHBuild.hpp"
 
#endif