PointKDTree.h

Go to the documentation of this file.

#ifndef G3D_PointKDTree_h
#define G3D_PointKDTree_h

#include "G3D/platform.h"
#include "G3D/Array.h"
#include "G3D/Table.h"
#include "G3D/Vector2.h"
#include "G3D/Vector3.h"
#include "G3D/Vector4.h"
#include "G3D/AABox.h"
#include "G3D/Sphere.h"
#include "G3D/Box.h"
#include "G3D/BinaryInput.h"
#include "G3D/BinaryOutput.h"
#include "G3D/CollisionDetection.h"
#include "G3D/Frustum.h"
#include "G3D/PositionTrait.h"
#include <algorithm>

namespace G3D {

template<class T,
 class PositionFunc = PositionTrait<T>, 
 class HashFunc = HashTrait<T>, 
 class EqualsFunc = EqualsTrait<T> > 
class PointKDTree {
protected:
#define TreeType PointKDTree<T, PositionFunc, HashFunc, EqualsFunc>

 // Unlike the KDTree, the PointKDTree assumes that T elements are
 // small and keeps the handle and cached position together instead of
 // placing them in separate bounds arrays. Also note that a copy of T
 // is kept in the member table and that there is no indirection.
 class Handle {
 private:
 Vector3 m_position;

 public:
 T value;

 inline Handle() {}
 inline Handle(const T& v) : value(v) {
 PositionFunc::getPosition(v, m_position);
 }

 void setPosition(const Vector3& v) {
 m_position = v;
 }

 inline const Vector3& position() const {
 return m_position;
 }
 };

 static AABox computeBounds(
 const Array<Handle>& point) {
 
 if (point.size() == 0) {
 return AABox(Vector3::inf(), Vector3::inf());
 }

 AABox bounds(point[0].position());

 for (int p = 0; p < point.size(); ++p) {
 bounds.merge(point[p].position());
 }

 return bounds;
 }

 class Node {
 public:

 AABox splitBounds;

 Vector3::Axis splitAxis;

 float splitLocation;
 
 Node* child[2];

 Array<Handle> valueArray;

 Node() {
 splitAxis = Vector3::X_AXIS;
 splitLocation = 0;
 splitBounds = AABox(-Vector3::inf(), Vector3::inf());
 for (int i = 0; i < 2; ++i) {
 child[i] = NULL;
 }
 }

 Node(const Node& other) : valueArray(other.valueArray) {
 splitAxis = other.splitAxis;
 splitLocation = other.splitLocation;
 splitBounds = other.splitBounds; 
 for (int i = 0; i < 2; ++i) {
 child[i] = NULL;
 }
 }

 Node(const Array<Handle>& pt) {
 splitAxis = Vector3::X_AXIS;
 splitLocation = 0;
 for (int i = 0; i < 2; ++i) {
 child[i] = NULL;
 }
 valueArray = pt;
 }


 ~Node() {
 for (int i = 0; i < 2; ++i) {
 delete child[i];
 }
 }


 inline bool isLeaf() const {
 return (child[0] == NULL) && (child[1] == NULL);
 }


 void getHandles(Array<Handle>& handleArray) const {
 handleArray.append(valueArray);
 for (int i = 0; i < 2; ++i) {
 if (child[i] != NULL) {
 child[i]->getHandles(handleArray);
 }
 }
 }


 void verifyNode(const Vector3& lo, const Vector3& hi) {
 // debugPrintf("Verifying: split %d @ %f [%f, %f, %f], [%f, %f, %f]\n",
 // splitAxis, splitLocation, lo.x, lo.y, lo.z, hi.x, hi.y, hi.z);

 debugAssert(lo == splitBounds.low());
 debugAssert(hi == splitBounds.high());

# ifdef G3D_DEBUG
 for (int i = 0; i < valueArray.length(); ++i) {
 const Vector3& b = valueArray[i].position();
 debugAssert(splitBounds.contains(b));
 }
# endif

 if (child[0] || child[1]) {
 debugAssert(lo[splitAxis] < splitLocation);
 debugAssert(hi[splitAxis] > splitLocation);
 }

 Vector3 newLo = lo;
 newLo[splitAxis] = splitLocation;
 Vector3 newHi = hi;
 newHi[splitAxis] = splitLocation;

 if (child[0] != NULL) {
 child[0]->verifyNode(lo, newHi);
 }

 if (child[1] != NULL) {
 child[1]->verifyNode(newLo, hi);
 }
 }


 static void serializeStructure(const Node* n, BinaryOutput& bo) {
 if (n == NULL) {
 bo.writeUInt8(0);
 } else {
 bo.writeUInt8(1);
 n->splitBounds.serialize(bo);
 serialize(n->splitAxis, bo);
 bo.writeFloat32(n->splitLocation);
 for (int c = 0; c < 2; ++c) {
 serializeStructure(n->child[c], bo);
 }
 }
 }

 static Node* deserializeStructure(BinaryInput& bi) {
 if (bi.readUInt8() == 0) {
 return NULL;
 } else {
 Node* n = new Node();
 n->splitBounds.deserialize(bi);
 deserialize(n->splitAxis, bi);
 n->splitLocation = bi.readFloat32();
 for (int c = 0; c < 2; ++c) {
 n->child[c] = deserializeStructure(bi);
 }
 return n;
 }
 }

 Node* findDeepestContainingNode(const Vector3& point) {

 // See which side of the splitting plane the bounds are on
 if (point[splitAxis] < splitLocation) {
 // Point is on the low side. Recurse into the child
 // if it exists.
 if (child[0] != NULL) {
 return child[0]->findDeepestContainingNode(point);
 }
 } else if (point[splitAxis] > splitLocation) {
 // Point is on the high side, recurse into the child
 // if it exists.
 if (child[1] != NULL) {
 return child[1]->findDeepestContainingNode(point);
 }
 }

 // There was no containing child, so this node is the
 // deepest containing node.
 return this;
 }

 void getIntersectingMembers(
 const AABox& sphereBounds,
 const Sphere& sphere,
 Array<T>& members) const {

 // Test all values at this node. Extract the
 // underlying C array for speed
 const int N = valueArray.size();
 const Handle* handleArray = valueArray.getCArray();
 
 const float r2 = square(sphere.radius);

 // Copy the sphere center so that it is on the stack near the radius
 const Vector3 center = sphere.center; 
 for (int v = 0; v < N; ++v) {
 if ((center - handleArray[v].position()).squaredLength() <= r2) {
 members.append(handleArray[v].value);
 }
 }

 // If the left child overlaps the box, recurse into it
 if (child[0] && (sphereBounds.low()[splitAxis] < splitLocation)) {
 child[0]->getIntersectingMembers(sphereBounds, sphere, members);
 }

 // If the right child overlaps the box, recurse into it
 if (child[1] && (sphereBounds.high()[splitAxis] > splitLocation)) {
 child[1]->getIntersectingMembers(sphereBounds, sphere, members);
 }
 }

 void getIntersectingMembers(
 const AABox& box, 
 const Sphere& sphere,
 Array<T>& members,
 bool useSphere) const {

 // Test all values at this node
 for (int v = 0; v < valueArray.size(); ++v) {
 if ((useSphere && sphere.contains(valueArray[v].position())) ||
 (! useSphere && box.contains(valueArray[v].position()))) {
 members.append(valueArray[v].value);
 }
 }

 // If the left child overlaps the box, recurse into it
 if ((child[0] != NULL) && (box.low()[splitAxis] < splitLocation)) {
 child[0]->getIntersectingMembers(box, sphere, members, useSphere);
 }

 // If the right child overlaps the box, recurse into it
 if ((child[1] != NULL) && (box.high()[splitAxis] > splitLocation)) {
 child[1]->getIntersectingMembers(box, sphere, members, useSphere);
 }
 }

 void assignSplitBounds(const AABox& myBounds) {
 splitBounds = myBounds;

# ifdef G3D_DEBUG
 if (child[0] || child[1]) {
 debugAssert(splitBounds.high()[splitAxis] > splitLocation);
 debugAssert(splitBounds.low()[splitAxis] < splitLocation);
 }
# endif

 AABox childBounds[2];
 myBounds.split(splitAxis, splitLocation, childBounds[0], childBounds[1]);

 for (int c = 0; c < 2; ++c) {
 if (child[c]) {
 child[c]->assignSplitBounds(childBounds[c]);
 }
 }
 }
 };

 class AxisComparator {
 private:
 Vector3::Axis sortAxis;

 public:

 AxisComparator(Vector3::Axis s) : sortAxis(s) {}

 inline int operator()(const Handle& A, const Handle& B) const {
 if (A.position()[sortAxis] > B.position()[sortAxis]) {
 return -1;
 } else if (A.position()[sortAxis] < B.position()[sortAxis]) {
 return 1;
 } else {
 return 0;
 }
 }
 };

 Node* makeNode(
 Array<Handle>& source, 
 Array<Handle>& temp,
 int valuesPerNode, 
 int numMeanSplits) {

 Node* node = NULL;
 
 if (source.size() <= valuesPerNode) {
 // Make a new leaf node
 node = new Node(source);
 
 // Set the pointers in the memberTable
 for (int i = 0; i < source.size(); ++i) {
 memberTable.set(source[i].value, node);
 }
 
 } else {
 // Make a new internal node
 node = new Node();
 
 const AABox bounds = computeBounds(source);
 const Vector3 extent = bounds.high() - bounds.low();
 
 Vector3::Axis splitAxis = extent.primaryAxis();
 
 float splitLocation;
 
 Array<Handle> lt, gt;

 if (numMeanSplits <= 0) {
 source.medianPartition(lt, node->valueArray, gt, temp, AxisComparator(splitAxis));
 splitLocation = node->valueArray[0].position()[splitAxis];
 
 if ((node->valueArray.size() > source.size() / 2) &&
 (source.size() > 10)) {
 // Our median split put an awful lot of points on the splitting plane. Try a mean
 // split instead
 numMeanSplits = 1;
 }
 }

 if (numMeanSplits > 0) {
 // Compute the mean along the axis

 splitLocation = (bounds.high()[splitAxis] + 
 bounds.low()[splitAxis]) / 2.0f;

 Handle splitHandle;
 Vector3 v;
 v[splitAxis] = splitLocation;
 splitHandle.setPosition(v);

 source.partition(splitHandle, lt, node->valueArray, gt, AxisComparator(splitAxis));
 }

# if defined(G3D_DEBUG) && defined(VERIFY_TREE)
 for (int i = 0; i < lt.size(); ++i) {
 const Vector3& v = lt[i].position(); 
 debugAssert(v[splitAxis] < splitLocation);
 }
 for (int i = 0; i < gt.size(); ++i) {
 debugAssert(gt[i].position()[splitAxis] > splitLocation);
 }
 for (int i = 0; i < node->valueArray.size(); ++i) {
 debugAssert(node->valueArray[i].position()[splitAxis] == splitLocation);
 }
# endif

 node->splitAxis = splitAxis;
 node->splitLocation = splitLocation;

 // Throw away the source array to save memory
 source.fastClear();
 
 if (lt.size() > 0) {
 node->child[0] = makeNode(lt, temp, valuesPerNode, numMeanSplits - 1);
 }

 if (gt.size() > 0) {
 node->child[1] = makeNode(gt, temp, valuesPerNode, numMeanSplits - 1);
 }

 // Add the values stored at this interior node to the member table
 for(int i = 0; i < node->valueArray.size(); ++i) {
 memberTable.set(node->valueArray[i].value, node);
 }
 
 }
 
 return node;
 }

 Node* cloneTree(Node* src) {
 Node* dst = new Node(*src);

 // Make back pointers
 for (int i = 0; i < dst->valueArray.size(); ++i) {
 memberTable.set(dst->valueArray[i].value, dst);
 }

 // Clone children
 for (int i = 0; i < 2; ++i) {
 if (src->child[i] != NULL) {
 dst->child[i] = cloneTree(src->child[i]);
 }
 }

 return dst;
 }

 typedef Table<T, Node*, HashFunc, EqualsFunc> MemberTable;
 MemberTable memberTable;

 Node* root;

public:

 PointKDTree() : root(NULL) {}


 PointKDTree(const PointKDTree& src) : root(NULL) {
 *this = src;
 }


 PointKDTree& operator=(const PointKDTree& src) {
 delete root;
 // Clone tree takes care of filling out the memberTable.
 root = cloneTree(src.root);
 return *this;
 }


 ~PointKDTree() {
 clear();
 }

 void clear() {
 memberTable.clear();
 delete root;
 root = NULL;
 }

 void clearData() {
 memberTable.clear();
 Array<Node*> stack;
 stack.push(root);
 while (stack.size() > 0) {
 Node* node = stack.pop();
 node->valueArray.fastClear();

 for (int i = 0; i < 2; ++i) {
 if (node->child[i] != NULL) {
 stack.push(node->child[i]);
 }
 }
 }
 }


 size_t size() const {
 return memberTable.size();
 }

 void insert(const T& value) {
 if (contains(value)) {
 // Already in the set
 return;
 }

 Handle h(value);

 if (root == NULL) {
 // This is the first node; create a root node
 root = new Node();
 }

 Node* node = root->findDeepestContainingNode(h.position());

 // Insert into the node
 node->valueArray.append(h);
 
 // Insert into the node table
 memberTable.set(value, node);
 }

 void insert(const Array<T>& valueArray) {
 // Pre-size the member table to avoid multiple allocations
 memberTable.setSizeHint(valueArray.size() + size());

 if (root == NULL) {
 // Optimized case for an empty tree; don't bother
 // searching or reallocating the root node's valueArray
 // as we incrementally insert.
 root = new Node();
 root->valueArray.resize(valueArray.size());
 for (int i = 0; i < valueArray.size(); ++i) {
 // Insert in opposite order so that we have the exact same
 // data structure as if we inserted each (i.e., order is reversed
 // from array).
 root->valueArray[valueArray.size() - i - 1] = Handle(valueArray[i]);
 memberTable.set(valueArray[i], root);
 }
 } else {
 // Insert at appropriate tree depth.
 for (int i = 0; i < valueArray.size(); ++i) {
 insert(valueArray[i]);
 }
 }
 }


 bool contains(const T& value) {
 return memberTable.containsKey(value);
 }


 void remove(const T& value) {
 debugAssertM(contains(value),
 "Tried to remove an element from a "
 "PointKDTree that was not present");

 Array<Handle>& list = memberTable[value]->valueArray;

 // Find the element and remove it
 for (int i = list.length() - 1; i >= 0; --i) {
 if (list[i].value == value) {
 list.fastRemove(i);
 break;
 }
 }
 memberTable.remove(value);
 }


 void update(const T& value) {
 if (contains(value)) {
 remove(value);
 }
 insert(value);
 }


 void balance(int valuesPerNode = 40, int numMeanSplits = 3) {
 if (root == NULL) {
 // Tree is empty
 return;
 }

 Array<Handle> handleArray;
 root->getHandles(handleArray);

 // Delete the old tree
 clear();

 Array<Handle> temp;
 root = makeNode(handleArray, temp, valuesPerNode, numMeanSplits);
 temp.fastClear();

 // Walk the tree, assigning splitBounds. We start with unbounded
 // space.
 root->assignSplitBounds(AABox::maxFinite());

# ifdef _DEBUG
 root->verifyNode(Vector3::minFinite(), Vector3::maxFinite());
# endif
 }

private:

 static void getIntersectingMembers(
 const Array<Plane>& plane,
 Array<T>& members,
 Node* node,
 uint32 parentMask) {

 int dummy;

 if (parentMask == 0) {
 // None of these planes can cull anything
 for (int v = node->valueArray.size() - 1; v >= 0; --v) {
 members.append(node->valueArray[v].value);
 }

 // Iterate through child nodes
 for (int c = 0; c < 2; ++c) {
 if (node->child[c]) {
 getIntersectingMembers(plane, members, node->child[c], 0);
 }
 }
 } else {

 if (node->valueArray.size() > 0) {
 // This is a leaf; check the points
 debugAssertM(node->child[0] == NULL, "Malformed Point tree");
 debugAssertM(node->child[1] == NULL, "Malformed Point tree");

 // Test values at this node against remaining planes
 for (int p = 0; p < plane.size(); ++p) {
 if (((parentMask >> p) & 1) != 0) {
 // Test against this plane
 const Plane& curPlane = plane[p];
 for (int v = node->valueArray.size() - 1; v >= 0; --v) {
 if (curPlane.halfSpaceContains(node->valueArray[v].position())) {
 members.append(node->valueArray[v].value);
 }
 }
 }
 }
 } else {

 uint32 childMask = 0xFFFFFF;

 // Iterate through child nodes
 for (int c = 0; c < 2; ++c) {
 if (node->child[c] &&
 ! node->child[c]->splitBounds.culledBy(plane, dummy, parentMask, childMask)) {
 // This node was not culled
 getIntersectingMembers(plane, members, node->child[c], childMask);
 }
 }
 }
 }
 }

public:

 void getIntersectingMembers(const Array<Plane>& plane, Array<T>& members) const {
 if (root == NULL) {
 return;
 }

 getIntersectingMembers(plane, members, root, 0xFFFFFF);
 }

 void getIntersectingMembers(const Frustum& frustum, Array<T>& members) const {
 Array<Plane> plane;
 
 for (int i = 0; i < frustum.faceArray.size(); ++i) {
 plane.append(frustum.faceArray[i].plane);
 }

 getIntersectingMembers(plane, members);
 }

 // This iterator turns Node::getIntersectingMembers into a
 // coroutine. It first translates that method from recursive to
 // stack based, then captures the system state (analogous to a Scheme
 // continuation) after each element is appended to the member array,
 // and allowing the computation to be restarted.
 class BoxIntersectionIterator {
 private:
 friend class TreeType;

 bool isEnd;

 AABox box;

 Node* node;

 // We could use backpointers within the tree and careful
 // state management to avoid ever storing the stack-- but
 // it is much easier this way and only inefficient if the
 // caller uses post increment (which they shouldn't!).
 Array<Node*> stack;

 int nextValueArrayIndex;

 BoxIntersectionIterator() : isEnd(true) {}
 
 BoxIntersectionIterator(const AABox& b, const Node* root) : 
 isEnd(root == NULL), box(b), 
 node(const_cast<Node*>(root)), nextValueArrayIndex(-1) {

 // We intentionally start at the "-1" index of the current
 // node so we can use the preincrement operator to move
 // ourselves to element 0 instead of repeating all of the
 // code from the preincrement method. Note that this might
 // cause us to become the "end" instance.
 ++(*this);
 }

 public:

 inline bool operator!=(const BoxIntersectionIterator& other) const {
 return ! (*this == other);
 }

 bool operator==(const BoxIntersectionIterator& other) const {
 if (isEnd) {
 return other.isEnd;
 } else if (other.isEnd) {
 return false;
 } else {
 // Two non-end iterators; see if they match. This is kind of 
 // silly; users shouldn't call == on iterators in general unless
 // one of them is the end iterator.
 if ((box != other.box) || (node != other.node) || 
 (nextValueArrayIndex != other.nextValueArrayIndex) ||
 (stack.length() != other.stack.length())) {
 return false;
 }

 // See if the stacks are the same
 for (int i = 0; i < stack.length(); ++i) {
 if (stack[i] != other.stack[i]) {
 return false;
 }
 }

 // We failed to find a difference; they must be the same
 return true;
 }
 }

 BoxIntersectionIterator& operator++() {
 ++nextValueArrayIndex;

 bool foundIntersection = false;
 while (! isEnd && ! foundIntersection) {

 // Search for the next node if we've exhausted this one
 while ((! isEnd) && (nextValueArrayIndex >= node->valueArray.length())) {
 // If we entered this loop, then the iterator has exhausted the elements at 
 // node (possibly because it just switched to a child node with no members).
 // This loop continues until it finds a node with members or reaches
 // the end of the whole intersection search.

 // If the right child overlaps the box, push it onto the stack for
 // processing.
 if ((node->child[1] != NULL) &&
 (box.high()[node->splitAxis] > node->splitLocation)) {
 stack.push(node->child[1]);
 }
 
 // If the left child overlaps the box, push it onto the stack for
 // processing.
 if ((node->child[0] != NULL) &&
 (box.low()[node->splitAxis] < node->splitLocation)) {
 stack.push(node->child[0]);
 }

 if (stack.length() > 0) {
 // Go on to the next node (which may be either one of the ones we 
 // just pushed, or one from farther back the tree).
 node = stack.pop();
 nextValueArrayIndex = 0;
 } else {
 // That was the last node; we're done iterating
 isEnd = true;
 }
 }

 // Search for the next intersection at this node until we run out of children
 while (! isEnd && ! foundIntersection && (nextValueArrayIndex < node->valueArray.length())) {
 if (box.intersects(node->valueArray[nextValueArrayIndex].bounds)) {
 foundIntersection = true;
 } else {
 ++nextValueArrayIndex;
 // If we exhaust this node, we'll loop around the master loop 
 // to find a new node.
 }
 }
 }

 return *this;
 }

 BoxIntersectionIterator operator++(int) {
 BoxIntersectionIterator old = *this;
 ++this;
 return old;
 }

 const T& operator*() const {
 alwaysAssertM(! isEnd, "Can't dereference the end element of an iterator");
 return node->valueArray[nextValueArrayIndex].value;
 }

 T const * operator->() const {
 alwaysAssertM(! isEnd, "Can't dereference the end element of an iterator");
 return &(stack.last()->valueArray[nextValueArrayIndex].value);
 }

 operator T*() const {
 alwaysAssertM(! isEnd, "Can't dereference the end element of an iterator");
 return &(stack.last()->valueArray[nextValueArrayIndex].value);
 }
 };


 BoxIntersectionIterator beginBoxIntersection(const AABox& box) const {
 return BoxIntersectionIterator(box, root);
 }

 BoxIntersectionIterator endBoxIntersection() const {
 // The "end" iterator instance
 return BoxIntersectionIterator();
 }

 void getIntersectingMembers(const AABox& box, Array<T>& members) const {
 if (root == NULL) {
 return;
 }
 root->getIntersectingMembers(box, Sphere(Vector3::zero(), 0), members, false);
 }


 void getIntersectingMembers(const Sphere& sphere, Array<T>& members) const {
 if (root == NULL) {
 return;
 }

 AABox box;
 sphere.getBounds(box);
 root->getIntersectingMembers(box, sphere, members);

 }


 void serializeStructure(BinaryOutput& bo) const {
 Node::serializeStructure(root, bo);
 }

 void deserializeStructure(BinaryInput& bi) {
 clear();
 root = Node::deserializeStructure(bi);
 }

 void getMembers(Array<T>& members) const {
 memberTable.getKeys(members);
 }


 class Iterator {
 private:
 friend class TreeType;

 // Note: this is a Table iterator, we are currently defining
 // Set iterator
 typename MemberTable::Iterator it;

 Iterator(const typename MemberTable::Iterator& it) : it(it) {}

 public:
 inline bool operator!=(const Iterator& other) const {
 return !(*this == other);
 }

 bool operator==(const Iterator& other) const {
 return it == other.it;
 }

 Iterator& operator++() {
 ++it;
 return *this;
 }

 Iterator operator++(int) {
 Iterator old = *this;
 ++(*this);
 return old;
 }

 const T& operator*() const {
 return it->key;
 }

 T* operator->() const {
 return &(it->key);
 }

 operator T*() const {
 return &(it->key);
 }
 };


 Iterator begin() const {
 return Iterator(memberTable.begin());
 }


 Iterator end() const {
 return Iterator(memberTable.end());
 }
#undef TreeType
};

}

#endif