1 /* Copyright Jukka Jyl�nki
2
3    Licensed under the Apache License, Version 2.0 (the "License");
4    you may not use this file except in compliance with the License.
5    You may obtain a copy of the License at
6
7        http://www.apache.org/licenses/LICENSE-2.0
8
9    Unless required by applicable law or agreed to in writing, software
10    distributed under the License is distributed on an "AS IS" BASIS,
11    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12    See the License for the specific language governing permissions and
13    limitations under the License. */
14
15 /** @file KdTree.h
16         @author Jukka Jyl�nki
17         @brief A KD-tree acceleration structure for static geometry. */
18 #pragma once
19
20 #include "../Math/MathTypes.h"
21 #include "../Math/myassert.h"
22 #include "Triangle.h"
23
24 #ifdef MATH_CONTAINERLIB_SUPPORT
25 #include "Container/MaxHeap.h"
26 #endif
27
28 MATH_BEGIN_NAMESPACE
29
30 enum CardinalAxis
31 {
32         AxisX = 0,
33         AxisY,
34         AxisZ,
35         AxisNone
36 };
37
38 struct KdTreeNode
39 {
40         /// If this is an inner node, specifies along which axis this node is split. Of type CardinalAxis.
41         /// If this is a leaf, has the value AxisNone.
42         unsigned splitAxis : 2;
43         /// If this is an inner node, specifies the index/offset to the child node pair.
44         /// If this is a leaf, the value is undefined.
45         unsigned childIndex : 30;
46         union
47         {
48                 float splitPos; ///< If this is an inner node, specifies the position along the cardinal axis of the split.
49                 /// If this is a leaf, specifies the index/ofset to the object bucket for this leaf.
50                 /// If zero, this leaf does not have a bucket associated with it. (empty leaf)
51                 u32 bucketIndex;
52         };
53
54         /// If true, this leaf does not contain any objects.
55         bool IsEmptyLeaf() const { assert(IsLeaf()); return bucketIndex == 0; }
56         bool IsLeaf() const { return splitAxis == AxisNone; }
57         int LeftChildIndex() const { return (int)childIndex; }
58         int RightChildIndex() const { return (int)childIndex+1; }
59         CardinalAxis SplitAxis() const { return (CardinalAxis)splitAxis; }
60 };
61
62 /// Type T must have a member function bool T.Intersects(const AABB &) const;
63 template<typename T>
64 class KdTree
65 {
66 public:
67         /// Represents the end of list in the index list of a bucket.
68         static const u32 BUCKET_SENTINEL = 0xFFFFFFFF;
69
70         /// Constructs an empty kD-tree.
71         KdTree()
72 #ifdef _DEBUG
73         :needsBuilding(false)
74 #endif
75         {}
76
77         ~KdTree();
78         
79         /// Adds a given number of objects to this kD-tree.
80         /// Call this function repeatedly as many times as necessary to prepare the data. Then
81         /// call Build() to create the tree data structure.
82         void AddObjects(const T *objects, int numObjects);
83
84         /// Creates the kD-tree data structure based on all the objects added to the tree.
85         /// After Build() has been called, do *not* call AddObjects() again.
86         void Build();
87
88         /// Empties the whole kD-tree of all objects.
89         /// Call this function if you want to reuse this structure for rebuilding another kD-tree, after first
90         /// having called AddObjects/Build to build a previous tree.
91         void Clear();
92
93         /// Returns an object bucket by the given bucket index.
94         /// An object bucket is a contiguous C array of object indices, terminated with a sentinel value BUCKET_SENTINEL.
95         /// To fetch the actual object based on an object index, call the Object() method.
96         u32 *Bucket(int bucketIndex);
97         const u32 *Bucket(int bucketIndex) const;
98
99         /// Returns an object by the given object index.
100         T &Object(int objectIndex);
101         const T &Object(int objectIndex) const;
102
103         /// Returns the total number of objects stored in this tree.
104         int NumObjects() const;
105
106         /// Returns the total number of nodes (all nodes, i.e. inner nodes + leaves) in the tree.
107         int NumNodes() const;
108
109         /// Returns the total number of leaf nodes in the tree.
110         /// Warning: This function iterates over the whole tree, so the running time is linear to the number of nodes, and not constant.
111         int NumLeaves() const;
112
113         /// Returns the total number of inner nodes in the tree.
114         /// Warning: This function iterates over the whole tree, so the running time is linear to the number of nodes, and not constant.
115         int NumInnerNodes() const;
116
117         /// Returns the maximum height of the tree (the path from the root to the farthest leaf node).
118         int TreeHeight() const;
119
120         /// Returns the root node.
121         KdTreeNode *Root();
122         const KdTreeNode *Root() const;
123
124         /// Returns true if the given node belongs to this kD-tree data structure.
125         /// Use only for debugging!
126         bool IsPartOfThisTree(const KdTreeNode *node) const;
127
128         bool IsPartOfThisTree(const KdTreeNode *root, const KdTreeNode *node) const;
129
130         /// Returns an AABB that tightly encloses all geometry in this kD-tree. Calling this is only valid after
131         /// Build() has been called.
132         const AABB &BoundingAABB() const { return rootAABB; }
133
134         /// Traverses a ray through this kD-tree, and calls the given leafCallback function for each leaf of the tree.
135         /// Uses the "recursive B" method from Vlastimil Havran's thesis.
136         /** @param r The ray to query through this kD-tree.
137                 @param leafCallback A function or a function object of prototype
138                         bool LeafCallbackFunction(KdTree<T> &tree, const KdTreeNode &leaf, const Ray &ray, float tNear, float tFar);
139                         If the callback function returns true, the execution of the query is stopped and this function immediately
140                         returns afterwards. If the callback function returns false, the execution of the query continues. */
141         template<typename Func>
142         inline void RayQuery(const Ray &r, Func &leafCallback);
143
144         /// Performs an AABB intersection query in this kD-tree, and calls the given leafCallback function for each leaf
145         /// of the tree which intersects the given AABB.
146         /** @param aabb The axis-aligned bounding box to query through this kD-tree.
147                 @param leafCallback A function or a function object of prototype
148                         bool LeafCallbackFunction(KdTree<T> &tree, KdTreeNode &leaf, const AABB &aabb);
149                         If the callback function returns true, the execution of the query is stopped and this function immediately
150                         returns afterwards. If the callback function returns false, the execution of the query continues. */
151         template<typename Func>
152         inline void AABBQuery(const AABB &aabb, Func &leafCallback);
153
154 #if 0 ///\bug Doesn't work properly. Fix up!
155         /// Performs an intersection query of this kD-tree against a given kD-tree, and calls the given
156         /// leafCallback function for each leaf pair that intersect each other.
157         /// @param leafCallback A function or a function object of prototype
158         ///    bool LeafCallbackFunction(KdTree<T> &thisTree, KdTreeNode &thisLeaf, const AABB &thisLeafAABB,
159         ///                              KdTree<T> &tree2, KdTreeNode &tree2Leaf, const OBB &tree2LeafOBB);
160         ///    If the callback function returns true, the execution of the query is stopped and this function immediately
161         ///    returns afterwards. If the callback function returns false, the execution of the query continues.
162         template<typename Func>
163         inline void KdTreeQuery(KdTree<T> &tree2, const float3x4 &thisWorldTransform, const float3x4 &tree2WorldTransform, Func &leafCallback);
164 #endif
165
166 #ifdef MATH_CONTAINERLIB_SUPPORT
167         /// Performs a nearest neighbor search on this kD-tree.
168         /// @param leafCallback A function or a function object of prototype
169         ///    bool LeafCallbackFunction(KdTree<T> &tree, const float3 &point, KdTreeNode &leaf, const AABB &aabb, float minDistance);
170         ///    If the callback function returns true, the execution of the query is stopped and this function immediately
171         ///    returns afterwards. If the callback function returns false, the execution of the query continues.
172         ///        minDistance is the minimum distance the objects in this leaf (and all future leaves to be passed to the
173         ///    callback) have to the point that is being queried.
174         template<typename Func>
175         inline void NearestObjects(const vec &point, Func &leafCallback);
176 #endif
177
178 private:
179         static const int maxNodes = 256 * 1024;
180         static const int maxTreeDepth = 30;
181
182         std::vector<KdTreeNode> nodes;
183         std::vector<u8, AlignedAllocator<u8, 16> > objects;
184         std::vector<u32*> buckets;
185
186         int AllocateNodePair();
187
188         void FreeBuckets();
189
190         AABB BoundingAABB(const u32 *bucket) const;
191
192         void SplitLeaf(int nodeIndex, const AABB &nodeAABB, int numObjectsInBucket, int leafDepth);
193
194         ///\todo Implement support for deep copying.
195         KdTree(const KdTree &);
196         void operator =(const KdTree &);
197
198         AABB rootAABB;
199
200 #ifdef _DEBUG
201         bool needsBuilding; ///< If true, new objects have been added to this tree, but it has not yet been rebuilt with Build();
202 #endif
203
204         int TreeHeight(int nodeIndex) const;
205 };
206
207 /// Finds the nearestray hit to a KdTree<Triangle>.
208 struct TriangleKdTreeRayQueryNearestHitVisitor
209 {
210         float rayT;
211         vec pos;
212         u32 triangleIndex;
213         float2 barycentricUV;
214
215         TriangleKdTreeRayQueryNearestHitVisitor()
216         {
217                 rayT = FLOAT_INF;
218                 triangleIndex = KdTree<Triangle>::BUCKET_SENTINEL;
219                 pos = vec::nan;
220                 barycentricUV = float2::nan;
221         }
222         bool operator()(KdTree<Triangle> &tree, const KdTreeNode &leaf, const Ray &ray, float tNear, float tFar)
223         {
224                 u32 *bucket = tree.Bucket(leaf.bucketIndex);
225                 assert(bucket);
226                 while(*bucket != KdTree<Triangle>::BUCKET_SENTINEL)
227                 {
228                         const Triangle &tri = tree.Object(*bucket);
229                         float u, v, t;
230                         t = Triangle::IntersectLineTri(ray.pos, ray.dir, tri.a, tri.b, tri.c, u, v);
231                         bool intersects = t < std::numeric_limits<float>::infinity();
232                         if (intersects && t >= tNear && t <= tFar && t < rayT)
233                         {
234                                 rayT = t;
235                                 pos = ray.GetPoint(t);
236                                 barycentricUV = float2(u,v);
237                                 triangleIndex = *bucket;
238                         }
239                         ++bucket;
240                 }
241                 return rayT < FLOAT_INF; // If we hit a triangle, no need to visit any farther nodes, since we are only interested in the nearest hit.
242         }
243 };
244
245 MATH_END_NAMESPACE
246
247 #include "KDTree.inl"