xref: /aosp_15_r20/external/eigen/unsupported/Eigen/src/BVH/KdBVH.h (revision bf2c37156dfe67e5dfebd6d394bad8b2ab5804d4)
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2009 Ilya Baran <[email protected]>
5 //
6 // This Source Code Form is subject to the terms of the Mozilla
7 // Public License v. 2.0. If a copy of the MPL was not distributed
8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9 
10 #ifndef KDBVH_H_INCLUDED
11 #define KDBVH_H_INCLUDED
12 
13 namespace Eigen {
14 
15 namespace internal {
16 
17 //internal pair class for the BVH--used instead of std::pair because of alignment
18 template<typename Scalar, int Dim>
19 struct vector_int_pair
20 {
21 EIGEN_MAKE_ALIGNED_OPERATOR_NEW_IF_VECTORIZABLE_FIXED_SIZE(Scalar, Dim)
22   typedef Matrix<Scalar, Dim, 1> VectorType;
23 
vector_int_pairvector_int_pair24   vector_int_pair(const VectorType &v, int i) : first(v), second(i) {}
25 
26   VectorType first;
27   int second;
28 };
29 
30 //these templates help the tree initializer get the bounding boxes either from a provided
31 //iterator range or using bounding_box in a unified way
32 template<typename ObjectList, typename VolumeList, typename BoxIter>
33 struct get_boxes_helper {
operatorget_boxes_helper34   void operator()(const ObjectList &objects, BoxIter boxBegin, BoxIter boxEnd, VolumeList &outBoxes)
35   {
36     outBoxes.insert(outBoxes.end(), boxBegin, boxEnd);
37     eigen_assert(outBoxes.size() == objects.size());
38     EIGEN_ONLY_USED_FOR_DEBUG(objects);
39   }
40 };
41 
42 template<typename ObjectList, typename VolumeList>
43 struct get_boxes_helper<ObjectList, VolumeList, int> {
44   void operator()(const ObjectList &objects, int, int, VolumeList &outBoxes)
45   {
46     outBoxes.reserve(objects.size());
47     for(int i = 0; i < (int)objects.size(); ++i)
48       outBoxes.push_back(bounding_box(objects[i]));
49   }
50 };
51 
52 } // end namespace internal
53 
54 
55 /** \class KdBVH
56  *  \brief A simple bounding volume hierarchy based on AlignedBox
57  *
58  *  \param _Scalar The underlying scalar type of the bounding boxes
59  *  \param _Dim The dimension of the space in which the hierarchy lives
60  *  \param _Object The object type that lives in the hierarchy.  It must have value semantics.  Either bounding_box(_Object) must
61  *                 be defined and return an AlignedBox<_Scalar, _Dim> or bounding boxes must be provided to the tree initializer.
62  *
63  *  This class provides a simple (as opposed to optimized) implementation of a bounding volume hierarchy analogous to a Kd-tree.
64  *  Given a sequence of objects, it computes their bounding boxes, constructs a Kd-tree of their centers
65  *  and builds a BVH with the structure of that Kd-tree.  When the elements of the tree are too expensive to be copied around,
66  *  it is useful for _Object to be a pointer.
67  */
68 template<typename _Scalar, int _Dim, typename _Object> class KdBVH
69 {
70 public:
71   enum { Dim = _Dim };
72   typedef _Object Object;
73   typedef std::vector<Object, aligned_allocator<Object> > ObjectList;
74   typedef _Scalar Scalar;
75   typedef AlignedBox<Scalar, Dim> Volume;
76   typedef std::vector<Volume, aligned_allocator<Volume> > VolumeList;
77   typedef int Index;
78   typedef const int *VolumeIterator; //the iterators are just pointers into the tree's vectors
79   typedef const Object *ObjectIterator;
80 
81   KdBVH() {}
82 
83   /** Given an iterator range over \a Object references, constructs the BVH.  Requires that bounding_box(Object) return a Volume. */
84   template<typename Iter> KdBVH(Iter begin, Iter end) { init(begin, end, 0, 0); } //int is recognized by init as not being an iterator type
85 
86   /** Given an iterator range over \a Object references and an iterator range over their bounding boxes, constructs the BVH */
87   template<typename OIter, typename BIter> KdBVH(OIter begin, OIter end, BIter boxBegin, BIter boxEnd) { init(begin, end, boxBegin, boxEnd); }
88 
89   /** Given an iterator range over \a Object references, constructs the BVH, overwriting whatever is in there currently.
90     * Requires that bounding_box(Object) return a Volume. */
91   template<typename Iter> void init(Iter begin, Iter end) { init(begin, end, 0, 0); }
92 
93   /** Given an iterator range over \a Object references and an iterator range over their bounding boxes,
94     * constructs the BVH, overwriting whatever is in there currently. */
95   template<typename OIter, typename BIter> void init(OIter begin, OIter end, BIter boxBegin, BIter boxEnd)
96   {
97     objects.clear();
98     boxes.clear();
99     children.clear();
100 
101     objects.insert(objects.end(), begin, end);
102     int n = static_cast<int>(objects.size());
103 
104     if(n < 2)
105       return; //if we have at most one object, we don't need any internal nodes
106 
107     VolumeList objBoxes;
108     VIPairList objCenters;
109 
110     //compute the bounding boxes depending on BIter type
111     internal::get_boxes_helper<ObjectList, VolumeList, BIter>()(objects, boxBegin, boxEnd, objBoxes);
112 
113     objCenters.reserve(n);
114     boxes.reserve(n - 1);
115     children.reserve(2 * n - 2);
116 
117     for(int i = 0; i < n; ++i)
118       objCenters.push_back(VIPair(objBoxes[i].center(), i));
119 
120     build(objCenters, 0, n, objBoxes, 0); //the recursive part of the algorithm
121 
122     ObjectList tmp(n);
123     tmp.swap(objects);
124     for(int i = 0; i < n; ++i)
125       objects[i] = tmp[objCenters[i].second];
126   }
127 
128   /** \returns the index of the root of the hierarchy */
129   inline Index getRootIndex() const { return (int)boxes.size() - 1; }
130 
131   /** Given an \a index of a node, on exit, \a outVBegin and \a outVEnd range over the indices of the volume children of the node
132     * and \a outOBegin and \a outOEnd range over the object children of the node */
133   EIGEN_STRONG_INLINE void getChildren(Index index, VolumeIterator &outVBegin, VolumeIterator &outVEnd,
134                                        ObjectIterator &outOBegin, ObjectIterator &outOEnd) const
135   { //inlining this function should open lots of optimization opportunities to the compiler
136     if(index < 0) {
137       outVBegin = outVEnd;
138       if(!objects.empty())
139         outOBegin = &(objects[0]);
140       outOEnd = outOBegin + objects.size(); //output all objects--necessary when the tree has only one object
141       return;
142     }
143 
144     int numBoxes = static_cast<int>(boxes.size());
145 
146     int idx = index * 2;
147     if(children[idx + 1] < numBoxes) { //second index is always bigger
148       outVBegin = &(children[idx]);
149       outVEnd = outVBegin + 2;
150       outOBegin = outOEnd;
151     }
152     else if(children[idx] >= numBoxes) { //if both children are objects
153       outVBegin = outVEnd;
154       outOBegin = &(objects[children[idx] - numBoxes]);
155       outOEnd = outOBegin + 2;
156     } else { //if the first child is a volume and the second is an object
157       outVBegin = &(children[idx]);
158       outVEnd = outVBegin + 1;
159       outOBegin = &(objects[children[idx + 1] - numBoxes]);
160       outOEnd = outOBegin + 1;
161     }
162   }
163 
164   /** \returns the bounding box of the node at \a index */
165   inline const Volume &getVolume(Index index) const
166   {
167     return boxes[index];
168   }
169 
170 private:
171   typedef internal::vector_int_pair<Scalar, Dim> VIPair;
172   typedef std::vector<VIPair, aligned_allocator<VIPair> > VIPairList;
173   typedef Matrix<Scalar, Dim, 1> VectorType;
174   struct VectorComparator //compares vectors, or more specifically, VIPairs along a particular dimension
175   {
176     VectorComparator(int inDim) : dim(inDim) {}
177     inline bool operator()(const VIPair &v1, const VIPair &v2) const { return v1.first[dim] < v2.first[dim]; }
178     int dim;
179   };
180 
181   //Build the part of the tree between objects[from] and objects[to] (not including objects[to]).
182   //This routine partitions the objCenters in [from, to) along the dimension dim, recursively constructs
183   //the two halves, and adds their parent node.  TODO: a cache-friendlier layout
184   void build(VIPairList &objCenters, int from, int to, const VolumeList &objBoxes, int dim)
185   {
186     eigen_assert(to - from > 1);
187     if(to - from == 2) {
188       boxes.push_back(objBoxes[objCenters[from].second].merged(objBoxes[objCenters[from + 1].second]));
189       children.push_back(from + (int)objects.size() - 1); //there are objects.size() - 1 tree nodes
190       children.push_back(from + (int)objects.size());
191     }
192     else if(to - from == 3) {
193       int mid = from + 2;
194       std::nth_element(objCenters.begin() + from, objCenters.begin() + mid,
195                         objCenters.begin() + to, VectorComparator(dim)); //partition
196       build(objCenters, from, mid, objBoxes, (dim + 1) % Dim);
197       int idx1 = (int)boxes.size() - 1;
198       boxes.push_back(boxes[idx1].merged(objBoxes[objCenters[mid].second]));
199       children.push_back(idx1);
200       children.push_back(mid + (int)objects.size() - 1);
201     }
202     else {
203       int mid = from + (to - from) / 2;
204       nth_element(objCenters.begin() + from, objCenters.begin() + mid,
205                   objCenters.begin() + to, VectorComparator(dim)); //partition
206       build(objCenters, from, mid, objBoxes, (dim + 1) % Dim);
207       int idx1 = (int)boxes.size() - 1;
208       build(objCenters, mid, to, objBoxes, (dim + 1) % Dim);
209       int idx2 = (int)boxes.size() - 1;
210       boxes.push_back(boxes[idx1].merged(boxes[idx2]));
211       children.push_back(idx1);
212       children.push_back(idx2);
213     }
214   }
215 
216   std::vector<int> children; //children of x are children[2x] and children[2x+1], indices bigger than boxes.size() index into objects.
217   VolumeList boxes;
218   ObjectList objects;
219 };
220 
221 } // end namespace Eigen
222 
223 #endif //KDBVH_H_INCLUDED
224