1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 #pragma once19 20 #include "[AABB.h]"21 #include "[OBB.h]"22 #include "[Ray.h]"23 #include "../Math/assume.h"24 #include "../Math/MathFunc.h"25 26 [MATH_BEGIN_NAMESPACE]27 28 template<typename T>29 int [KdTree<T>::AllocateNodePair]()30 {31 int index = (int)nodes.size();32 [KdTreeNode] n;33 n.[splitAxis] = [AxisNone]; 34 n.[bucketIndex] = 0;35 nodes.push_back(n);36 nodes.push_back(n);37 return index;38 }39 40 template<typename T>41 void [KdTree<T>::FreeBuckets]()42 {43 for(size_t i = 0; i < buckets.size(); ++i)44 delete[] buckets[i];45 buckets.clear();46 }47 48 template<typename T>49 [AABB] [KdTree<T>::BoundingAABB](const u32 *bucket) const50 {51 [assert](bucket);52 53 [AABB] a;54 a.[SetNegativeInfinity]();55 56 while(*bucket != BUCKET_SENTINEL)57 a.[Enclose](Object(*bucket++).BoundingAABB());58 59 return a;60 }61 62 template<typename T>63 void [KdTree<T>::SplitLeaf](int nodeIndex, const [AABB] &nodeAABB, int numObjectsInBucket, int leafDepth)64 {65 if (leafDepth >= maxTreeDepth)66 return; 67 68 [KdTreeNode] *node = &nodes[nodeIndex];69 [assert](node->[IsLeaf]());70 71 int curBucketIndex = node->[bucketIndex]; 72 [assert](curBucketIndex != 0); 73 [CardinalAxis] splitAxis = ([CardinalAxis])nodeAABB.[Size]().[MaxElementIndex]();74 float splitPos = nodeAABB.[CenterPoint]()[splitAxis];75 76 77 [AABB] leftAABB = nodeAABB;78 [AABB] rightAABB = nodeAABB;79 leftAABB.[maxPoint][splitAxis] = splitPos;80 rightAABB.[minPoint][splitAxis] = splitPos;81 82 83 u32 *leftBucket = new u32[numObjectsInBucket+1];84 u32 *rightBucket = new u32[numObjectsInBucket+1];85 86 u32 *curObject = buckets[curBucketIndex];87 u32 *l = leftBucket;88 u32 *r = rightBucket;89 int numObjectsLeft = 0;90 int numObjectsRight = 0;91 while(*curObject != BUCKET_SENTINEL)92 {93 [AABB] aabb = Object(*curObject).BoundingAABB();94 bool left = leftAABB.[Intersects](aabb);95 bool right = rightAABB.[Intersects](aabb);96 if (!left && !right)97 left = right = true; 98 if (left)99 {100 *l++ = *curObject;101 ++numObjectsLeft;102 }103 if (right)104 {105 *r++ = *curObject;106 ++numObjectsRight;107 }108 ++curObject;109 }110 *l = BUCKET_SENTINEL;111 *r = BUCKET_SENTINEL;112 113 114 if (numObjectsLeft == numObjectsInBucket || numObjectsRight == numObjectsInBucket)115 {116 delete[] leftBucket;117 delete[] rightBucket;118 return;119 }120 121 122 node->[splitAxis] = splitAxis;123 node->[splitPos] = splitPos;124 125 126 int childIndex = AllocateNodePair();127 node = &nodes[nodeIndex]; 128 node->[childIndex] = childIndex;129 130 131 leftAABB = BoundingAABB(leftBucket);132 rightAABB = BoundingAABB(rightBucket);133 134 135 [KdTreeNode] *leftChild = &nodes[childIndex];136 delete[] buckets[curBucketIndex];137 buckets[curBucketIndex] = leftBucket;138 leftChild->[bucketIndex] = curBucketIndex;139 140 141 [KdTreeNode] *rightChild = &nodes[childIndex+1];142 rightChild->[bucketIndex] = (u32)buckets.size();143 buckets.push_back(rightBucket);144 145 [assert](numObjectsLeft < numObjectsInBucket && numObjectsRight < numObjectsInBucket);146 147 148 if (numObjectsLeft > 16)149 SplitLeaf(childIndex, leftAABB, numObjectsLeft, leafDepth + 1);150 if (numObjectsRight > 16)151 SplitLeaf(childIndex+1, rightAABB, numObjectsRight, leafDepth + 1);152 }153 154 template<typename T>155 [KdTree<T>::~KdTree]()156 {157 FreeBuckets();158 }159 160 template<typename T>161 u32 *[KdTree<T>::Bucket](int bucketIndex)162 {163 return buckets[bucketIndex];164 }165 166 template<typename T>167 const u32 *[KdTree<T>::Bucket](int bucketIndex) const168 {169 return buckets[bucketIndex];170 }171 172 template<typename T>173 T &[KdTree<T>::Object](int objectIndex)174 {175 return (T&)objects[sizeof(T)*objectIndex];176 }177 178 template<typename T>179 const T &[KdTree<T>::Object](int objectIndex) const180 {181 return (const T&)objects[sizeof(T)*objectIndex];182 }183 184 185 template<typename T>186 int [KdTree<T>::NumNodes]() const187 {188 return nodes.size() - 1;189 }190 191 192 template<typename T>193 int [KdTree<T>::NumLeaves]() const194 {195 int numLeaves = 0;196 for(size_t i = 1; i < nodes.size(); ++i)197 if (nodes[i].IsLeaf())198 ++numLeaves;199 200 return numLeaves;201 }202 203 template<typename T>204 int [KdTree<T>::NumObjects]() const205 {206 return (int)(objects.size() / sizeof(T));207 }208 209 210 template<typename T>211 int [KdTree<T>::NumInnerNodes]() const212 {213 int numInnerNodes = 0;214 for(size_t i = 1; i < nodes.size(); ++i)215 if (!nodes[i].IsLeaf())216 ++numInnerNodes;217 218 return numInnerNodes;219 }220 221 template<typename T>222 int [KdTree<T>::TreeHeight](int nodeIndex) const223 {224 const [KdTreeNode] &node = nodes[nodeIndex];225 if (node.[IsLeaf]())226 return 1;227 return 1 + std::max(TreeHeight(node.[LeftChildIndex]()), TreeHeight(node.[RightChildIndex]()));228 }229 230 template<typename T>231 int [KdTree<T>::TreeHeight]() const232 {233 return TreeHeight(1);234 }235 236 template<typename T>237 void [KdTree<T>::AddObjects](const T *objects_, int numObjects)238 {239 objects.insert(objects.end(), (const [u8]*)objects_, (const [u8]*)(objects_ + numObjects));240 #ifdef _DEBUG241 needsBuilding = true;242 #endif243 }244 245 template<typename T>246 void [KdTree<T>::Build]()247 {248 nodes.clear();249 FreeBuckets();250 251 252 [KdTreeNode] dummy;253 dummy.[splitAxis] = [AxisNone];254 dummy.[childIndex] = 0;255 dummy.[bucketIndex] = 0;256 nodes.push_back(dummy); 257 258 259 buckets.push_back(0);260 261 262 [KdTreeNode] rootNode;263 rootNode.[splitAxis] = [AxisNone];264 rootNode.[childIndex] = 0;265 rootNode.[bucketIndex] = 1;266 nodes.push_back(rootNode);267 268 269 u32 *rootBucket = new u32[NumObjects()+1];270 for(int i = 0; i < NumObjects(); ++i)271 rootBucket[i] = (u32)i;272 rootBucket[NumObjects()] = BUCKET_SENTINEL;273 buckets.push_back(rootBucket);274 275 rootAABB = BoundingAABB(rootBucket);276 277 278 279 SplitLeaf(1, rootAABB, NumObjects(), 1);280 281 #ifdef _DEBUG282 needsBuilding = false;283 #endif284 }285 286 template<typename T>287 void [KdTree<T>::Clear]()288 {289 nodes.clear();290 objects.clear();291 buckets.clear();292 #ifdef _DEBUG293 needsBuilding = false;294 #endif295 }296 297 template<typename T>298 [KdTreeNode] *[KdTree<T>::Root]() { return nodes.size() > 1 ? &nodes[1] : 0; }299 300 template<typename T>301 const [KdTreeNode] *[KdTree<T>::Root]() const { return nodes.size() > 1 ? &nodes[1] : 0; }302 303 template<typename T>304 bool [KdTree<T>::IsPartOfThisTree](const [KdTreeNode] *node) const305 {306 if (!Root())307 return false;308 if (!node)309 return false;310 return IsPartOfThisTree(Root(), node);311 }312 313 template<typename T>314 bool [KdTree<T>::IsPartOfThisTree](const [KdTreeNode] *root, const [KdTreeNode] *node) const315 {316 [assert](root);317 [assert](node);318 if (root == node)319 return true;320 if (root->[IsLeaf]())321 return false;322 return IsPartOfThisTree(&nodes[root->[LeftChildIndex]()], node) || IsPartOfThisTree(&nodes[root->[RightChildIndex]()], node);323 }324 325 326 template<typename T>327 template<typename Func>328 inline void [KdTree<T>::RayQuery](const [Ray] &r, Func &nodeProcessFunc)329 {330 float tNear = 0.f, tFar = [FLOAT_INF];331 332 [assume](rootAABB.IsFinite());333 [assume](!rootAABB.IsDegenerate());334 #ifdef _DEBUG335 [assume](!needsBuilding);336 #endif337 338 if (!rootAABB.IntersectLineAABB(r.[pos], r.[dir], tNear, tFar))339 return; 340 341 342 343 344 345 tNear = 0.f;346 tFar = [FLOAT_INF];347 348 static const [CardinalAxis] axes[] = { [AxisX], [AxisY], [AxisZ], [AxisX], AxisY };349 350 typedef int StackPtr; 351 struct StackElem352 {353 [KdTreeNode] *node;354 float t;355 vec pos; 356 StackPtr prev; 357 };358 359 const int cMaxStackItems = 50*2;360 StackElem stack[cMaxStackItems];361 362 [KdTreeNode] *farChild;363 [KdTreeNode] *currentNode = Root();364 StackPtr entryPoint = 0;365 stack[entryPoint].t = tNear;366 367 368 stack[entryPoint].pos = r.[pos] + [Max](tNear, 0.f) * r.[dir];369 370 StackPtr exitPoint = 1; 371 stack[exitPoint].t = tFar;372 stack[exitPoint].pos = r.[pos] + r.[dir] * tFar;373 stack[exitPoint].node = 0;374 375 376 while(currentNode)377 {378 while(!currentNode->[IsLeaf]()) 379 {380 381 382 383 384 385 386 387 const float splitPos = currentNode->[splitPos];388 const [CardinalAxis] axis = ([CardinalAxis])currentNode->[splitAxis];389 if (stack[entryPoint].pos[axis] <= splitPos)390 {391 if (stack[exitPoint].pos[axis] <= splitPos)392 { 393 currentNode = &nodes[currentNode->[LeftChildIndex]()];394 continue;395 }396 if ([EqualAbs](stack[exitPoint].pos[axis], splitPos))397 { 398 currentNode = &nodes[currentNode->[RightChildIndex]()];399 continue;400 }401 402 farChild = &nodes[currentNode->[RightChildIndex]()];403 currentNode = &nodes[currentNode->[LeftChildIndex]()];404 }405 else406 {407 if (splitPos < stack[exitPoint].pos[axis])408 { 409 currentNode = &nodes[currentNode->[RightChildIndex]()];410 continue;411 }412 413 farChild = &nodes[currentNode->[LeftChildIndex]()];414 currentNode = &nodes[currentNode->[RightChildIndex]()];415 }416 417 const float t = (splitPos - r.[pos][axis]) / r.[dir][axis];418 StackPtr tempPtr = exitPoint++;419 if (exitPoint == entryPoint) 420 ++exitPoint;421 [assert](exitPoint < cMaxStackItems);422 423 stack[exitPoint].prev = tempPtr;424 stack[exitPoint].t = t;425 stack[exitPoint].node = farChild;426 stack[exitPoint].pos[axis] = splitPos;427 const [CardinalAxis] axis2 = axes[axis+1];428 const [CardinalAxis] axis3 = axes[axis+2];429 stack[exitPoint].pos[axis2] = r.[pos][axis2] + t * r.[dir][axis2];430 stack[exitPoint].pos[axis3] = r.[pos][axis3] + t * r.[dir][axis3];431 }432 433 const float dNear = stack[entryPoint].t;434 const float dFar = stack[exitPoint].t;435 436 bool traversalFinished = nodeProcessFunc(*this, *currentNode, r, dNear, dFar);437 438 439 440 441 if (traversalFinished)442 return;443 444 445 entryPoint = exitPoint;446 currentNode = stack[exitPoint].node;447 exitPoint = stack[entryPoint].prev;448 }449 }450 451 template<typename T>452 template<typename Func>453 inline void [KdTree<T>::AABBQuery](const [AABB] &aabb, Func &leafCallback)454 {455 const int cMaxStackItems = maxTreeDepth*2;456 457 [KdTreeNode] *stack[cMaxStackItems];458 int stackSize = 1;459 stack[0] = Root();460 461 462 if (!aabb.[Intersects](BoundingAABB()))463 return;464 465 while(stackSize > 0)466 {467 [KdTreeNode] *cur = stack[--stackSize];468 [assert](!cur->[IsLeaf]());469 470 471 472 473 474 if (aabb.[minPoint][cur->[splitAxis]] <= cur->[splitPos])475 {476 [KdTreeNode] *leftChild = &nodes[cur->[LeftChildIndex]()];477 if (leftChild->[IsLeaf]()) 478 {479 if (leafCallback(*this, *leftChild, aabb))480 return; 481 }482 else 483 stack[stackSize++] = leftChild;484 }485 486 487 if (aabb.[maxPoint][cur->[splitAxis]] >= cur->[splitPos])488 {489 [KdTreeNode] *rightChild = &nodes[cur->[RightChildIndex]()];490 if (rightChild->[IsLeaf]()) 491 {492 if (leafCallback(*this, *rightChild, aabb))493 return; 494 }495 else 496 stack[stackSize++] = rightChild;497 }498 }499 }500 501 #if 0 502 503 struct StackElem504 {505 [KdTreeNode] *thisNode;506 [KdTreeNode] *tree2Node;507 [AABB] thisAABB;508 [AABB] tree2AABB;509 [OBB] tree2OBB;510 };511 512 513 const int cMaxStackItems = 50*2;514 StackElem stack[cMaxStackItems*1000];515 516 517 template<typename T>518 template<typename Func>519 inline void [KdTree<T>::KdTreeQuery]([KdTree<T>] &tree2, const [float3x4] &thisWorldTransform, const [float3x4] &tree2WorldTransform, Func &leafCallback)520 {521 [float3x4] tree2Transform = thisWorldTransform.[Inverted]() * tree2WorldTransform;522 523 int stackSize = 1;524 stack[0].thisNode = Root();525 stack[0].thisAABB = BoundingAABB();526 stack[0].tree2Node = tree2.[Root]();527 stack[0].tree2AABB = tree2.[BoundingAABB](); 528 [OBB] tree2OBB = stack[0].tree2AABB.Transform(tree2Transform);529 stack[0].tree2OBB = tree2OBB;530 531 if (!stack[0].thisAABB.Intersects(tree2OBB))532 return;533 534 while(stackSize > 0)535 {536 [assert](stackSize < cMaxStackItems*100);537 --stackSize;538 [KdTreeNode] *thisNode = stack[stackSize].thisNode;539 [KdTreeNode] *tree2Node = stack[stackSize].tree2Node;540 [AABB] thisAABB = stack[stackSize].thisAABB;541 [AABB] tree2AABB = stack[stackSize].tree2AABB;542 [OBB] tree2OBB = stack[stackSize].tree2OBB;543 544 if (thisNode->[IsLeaf]() && tree2Node->[IsLeaf]())545 {546 if (leafCallback(*this, *thisNode, thisAABB, tree2, *tree2Node, tree2OBB, tree2Transform))547 return;548 }549 else if (thisNode->[IsLeaf]())550 {551 if (thisNode->[IsEmptyLeaf]())552 continue; 553 554 555 [KdTreeNode] *tree2LeftChild = &tree2.nodes[tree2Node->[LeftChildIndex]()];556 [KdTreeNode] *tree2RightChild = &tree2.nodes[tree2Node->[RightChildIndex]()];557 [AABB] tree2LeftAABB = tree2AABB;558 [AABB] tree2RightAABB = tree2AABB;559 tree2LeftAABB.[maxPoint][tree2Node->[splitAxis]] = tree2Node->[splitPos];560 tree2RightAABB.[minPoint][tree2Node->[splitAxis]] = tree2Node->[splitPos];561 [OBB] tree2LeftOBB = tree2LeftAABB.[Transform](tree2Transform);562 [OBB] tree2RightOBB = tree2RightAABB.[Transform](tree2Transform);563 564 {565 stack[stackSize].tree2Node = tree2LeftChild;566 stack[stackSize].tree2AABB = tree2LeftAABB;567 stack[stackSize].tree2OBB = tree2LeftOBB;568 ++stackSize;569 }570 571 {572 stack[stackSize].thisNode = thisNode;573 stack[stackSize].thisAABB = thisAABB;574 stack[stackSize].tree2Node = tree2RightChild;575 stack[stackSize].tree2AABB = tree2RightAABB;576 stack[stackSize].tree2OBB = tree2RightOBB;577 ++stackSize;578 }579 }580 else if (tree2Node->[IsLeaf]())581 {582 if (tree2Node->[IsEmptyLeaf]())583 continue; 584 585 586 [KdTreeNode] *leftChild = &nodes[thisNode->[LeftChildIndex]()];587 [KdTreeNode] *rightChild = &nodes[thisNode->[RightChildIndex]()];588 [AABB] leftAABB = thisAABB;589 [AABB] rightAABB = thisAABB;590 leftAABB.[maxPoint][thisNode->[splitAxis]] = thisNode->[splitPos];591 rightAABB.[minPoint][thisNode->[splitAxis]] = thisNode->[splitPos];592 593 {594 stack[stackSize].thisNode = leftChild;595 stack[stackSize].thisAABB = leftAABB;596 ++stackSize;597 }598 599 {600 stack[stackSize].thisNode = rightChild;601 stack[stackSize].thisAABB = rightAABB;602 stack[stackSize].tree2Node = tree2Node;603 stack[stackSize].tree2AABB = tree2AABB;604 stack[stackSize].tree2OBB = tree2OBB;605 ++stackSize;606 }607 }608 else609 {610 611 [KdTreeNode] *leftChild = &nodes[thisNode->[LeftChildIndex]()];612 [KdTreeNode] *rightChild = &nodes[thisNode->[RightChildIndex]()];613 [AABB] leftAABB = thisAABB;614 [AABB] rightAABB = thisAABB;615 leftAABB.[maxPoint][thisNode->[splitAxis]] = thisNode->[splitPos];616 rightAABB.[minPoint][thisNode->[splitAxis]] = thisNode->[splitPos];617 618 [KdTreeNode] *tree2LeftChild = &tree2.nodes[tree2Node->[LeftChildIndex]()];619 [KdTreeNode] *tree2RightChild = &tree2.nodes[tree2Node->[RightChildIndex]()];620 [AABB] tree2LeftAABB = tree2AABB;621 [AABB] tree2RightAABB = tree2AABB;622 tree2LeftAABB.[maxPoint][tree2Node->[splitAxis]] = tree2Node->[splitPos];623 tree2RightAABB.[minPoint][tree2Node->[splitAxis]] = tree2Node->[splitPos];624 [OBB] tree2LeftOBB = tree2LeftAABB.[Transform](tree2Transform);625 [OBB] tree2RightOBB = tree2RightAABB.[Transform](tree2Transform);626 627 if (leftAABB.[Intersects](tree2LeftOBB))628 {629 stack[stackSize].thisNode = leftChild;630 stack[stackSize].thisAABB = leftAABB;631 stack[stackSize].tree2Node = tree2LeftChild;632 stack[stackSize].tree2AABB = tree2LeftAABB;633 stack[stackSize].tree2OBB = tree2LeftOBB;634 ++stackSize;635 }636 if (rightAABB.[Intersects](tree2LeftOBB))637 {638 stack[stackSize].thisNode = rightChild;639 stack[stackSize].thisAABB = rightAABB;640 stack[stackSize].tree2Node = tree2LeftChild;641 stack[stackSize].tree2AABB = tree2LeftAABB;642 stack[stackSize].tree2OBB = tree2LeftOBB;643 ++stackSize;644 }645 if (leftAABB.[Intersects](tree2RightOBB))646 {647 stack[stackSize].thisNode = leftChild;648 stack[stackSize].thisAABB = leftAABB;649 stack[stackSize].tree2Node = tree2RightChild;650 stack[stackSize].tree2AABB = tree2RightAABB;651 stack[stackSize].tree2OBB = tree2RightOBB;652 ++stackSize;653 }654 if (rightAABB.[Intersects](tree2RightOBB))655 {656 stack[stackSize].thisNode = rightChild;657 stack[stackSize].thisAABB = rightAABB;658 stack[stackSize].tree2Node = tree2RightChild;659 stack[stackSize].tree2AABB = tree2RightAABB;660 stack[stackSize].tree2OBB = tree2RightOBB;661 ++stackSize;662 }663 }664 }665 }666 667 #endif668 669 #ifdef MATH_CONTAINERLIB_SUPPORT670 struct NearestObjectsTraversalNode671 {672 float [d];673 [AABB] aabb;674 [KdTreeNode] *node;675 676 bool [operator <](const NearestObjectsTraversalNode &t) const { return [d] > t.d; }677 };678 679 template<typename T>680 template<typename Func>681 inline void [KdTree<T>::NearestObjects](const vec &point, Func &leafCallback)682 {683 MaxHeap<NearestObjectsTraversalNode> queue;684 NearestObjectsTraversalNode t;685 t.d = 0.f;686 t.aabb = BoundingAABB();687 t.node = Root();688 queue.Insert(t);689 690 while(queue.Size() > 0)691 {692 t = queue.Front();693 queue.PopFront();694 695 if (t.node->IsLeaf())696 leafCallback(*this, point, *t.node, t.aabb, t.d);697 else698 {699 NearestObjectsTraversalNode n;700 701 702 n.aabb = t.aabb;703 n.aabb.maxPoint[t.node->splitAxis] = t.node->splitPos;704 n.node = &nodes[t.node->LeftChildIndex()];705 n.d = n.aabb.Distance(point);706 queue.Insert(n);707 708 709 n.aabb.maxPoint[t.node->splitAxis] = t.aabb.maxPoint[t.node->splitAxis]; 710 n.aabb.minPoint[t.node->splitAxis] = t.node->splitPos;711 n.node = &nodes[t.node->RightChildIndex()];712 n.d = n.aabb.Distance(point);713 queue.Insert(n);714 }715 }716 }717 #endif718 719 [MATH_END_NAMESPACE] Go back to previous page