kdtree.cpp

Go to the documentation of this file.

/*
 * kdtree.cpp
 *
 * Copyright (C) 2009  Thomas A. Vaughan
 * All rights reserved.
 *
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the <organization> nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THOMAS A. VAUGHAN ''AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THOMAS A. VAUGHAN BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 *
 * Implementation of the kdtree API.  See kdtree.h
 */

// includes --------------------------------------------------------------------
#include "kdtree.h"             // always include our own header first!

#include "perf/perf.h"
#include "util/distribution.h"


namespace kdtree {



// interface destructors
Entry::~Entry(void) throw() { }

static const float s_eps                = 1.0e-9;       // epsilon

////////////////////////////////////////////////////////////////////////////////
//
//      static helper methods
//
////////////////////////////////////////////////////////////////////////////////

eSort
sortPoint
(
IN const plane_t& plane,
IN const point3d_t& point
)
throw()
{
        float delta = 0.0;
        switch (plane.axis) {
        case eAxis_X:
                delta = point.x - plane.value;
                break;

        case eAxis_Y:
                delta = point.y - plane.value;
                break;

        case eAxis_Z:
                delta = point.z - plane.value;
                break;
        
        default:
                return eSort_Invalid;
        }

        if (delta < -s_eps)
                return eSort_Left;
        if (delta > s_eps)
                return eSort_Right;
        
        // must be in plane!
        return eSort_Straddle;
}



eSort
sortRect
(
IN const plane_t& plane,
IN const rect3d_t& r
)
throw()
{
        // either all corners are in same sort, or it straddles
        // (we ignore straddling corners)
        eSort sort = eSort_Invalid;
        for (int i = 0; i < rect3d_t::eMaxCorners; ++i) {
                point3d_t p = r.getCorner(i);
                eSort val = sortPoint(plane, p);
                if (eSort_Straddle == val)
                        continue;       // skip this corner

                if (eSort_Invalid == sort) {
                        sort = val;
                } else if (sort != val) {
                        return eSort_Straddle;
                }
        }

        // return the result
        if (eSort_Invalid == sort) {
                // we can get here for thin rectangles that entirely straddle
                //   the sort plane!  For instance, if a very flat object (2D
                //   rect) was chosen as the sort plane, the 2D rect will be
                //   entirely contained in the sort plane itself.
                return eSort_Straddle;
        }
        return sort;
}



static eAxis
getLongAxis
(
IN const rect3d_t& r
)
throw()
{
        float dx = r.x1 - r.x0;
        float dy = r.y1 - r.y0;
        float dz = r.z1 - r.z0;

        eAxis axis = eAxis_Invalid;

        if (dx > dy) {
                if (dx > dz) {
                        axis = eAxis_X;
                } else {
                        axis = eAxis_Z;
                }
        } else if (dy > dz) {
                axis = eAxis_Y;
        } else {
                axis = eAxis_Z;
        }

        return axis;
}




static void
simpleSplit
(
IN const rect3d_t& r,
OUT plane_t& plane
)
throw()
{
        plane.axis = getLongAxis(r);
        switch (plane.axis) {
        case eAxis_X:
                plane.value = 0.5 * (r.x0 + r.x1);
                break;

        case eAxis_Y:
                plane.value = 0.5 * (r.y0 + r.y1);
                break;

        case eAxis_Z:
                plane.value = 0.5 * (r.z0 + r.z1);
                break;

        default:
                ASSERT(false, "Bad axis: %d", plane.axis);
        }
}



static void
measureSplit
(
IN const plane_t& plane,
IN const Node::vec_entries_t& vec,
OUT int& nLeft,
OUT int& nRight
)
throw()
{
        // walk all entries and find nLeft & nRight
        nLeft = nRight = 0;
        for (Node::vec_entries_t::const_iterator i = vec.begin();
             i != vec.end(); ++i) {
                const Node::entry_rec_t& er = *i;
                eSort sort = sortRect(plane, er.rect);
                if (eSort_Left == sort) {
                        ++nLeft;
                } else if (eSort_Right == sort) {
                        ++nRight;
                }
        }
}



static void
bestSplit
(
IN int offset,
IN const Node::vec_entries_t& vec,
OUT plane_t& plane,
OUT int& nLeft,
OUT int& nRight
)
{
        ASSERT(offset >= 0 && offset < 3, "Bad offset: %d", offset);
        plane.clear();
        plane.axis = (eAxis) (offset + 1);
        //DPRINTF("Finding best split using axis = %d", plane.axis);

        // walk through all entries and populate distribution with relevant
        //   coordinate
        distribution_t<float> dist;
        for (Node::vec_entries_t::const_iterator i = vec.begin();
             i != vec.end(); ++i) {
                const Node::entry_rec_t& er = *i;
                const float * p0 = &er.rect.x0;
                p0 += offset;
                //DPRINTF("  sample: %f", *p0);
                dist.addSample(*p0);
        }

        // a naive approach is to use the median value, but that only
        // works for randomly-distributed objects.  For degenerate cases
        // (many objects clustered at a particular coordinate value) the
        // median breaks down.  This is especially true when there are a
        // small number of objects.
        //
        // So in general we need to perform a binary search to find the best
        // value.
        //
        // We search for the first point (walking right) where the nLeft
        // count is larger than nRight.  Then we either use that point, or the
        // point immediately to the left.
        int N = dist.size();
        ASSERT(N > 1, "Not enough entries");
        int start = 0;          // starting edge for binary search
        int end = N - 1;        // ending edge for binary search

        int nLeftStart, nLeftEnd;
        plane.value = dist.getValueAtIndex(start);
        measureSplit(plane, vec, nLeftStart, nRight);
        plane.value = dist.getValueAtIndex(end);
        measureSplit(plane, vec, nLeftEnd, nRight);

        int mid = end / 2;
        //DPRINTF("Performing binary search [0 - %d]:", end);
        //DPRINTF("  nLeftStart=%d  nLeftEnd=%d", nLeftStart, nLeftEnd);
        while (true) {

                // how many objects are here?
                int nLeftMid;
                plane.value = dist.getValueAtIndex(mid);
                measureSplit(plane, vec, nLeftMid, nRight);

                //DPRINTF("  start=%d(%d)  end=%d(%d)  mid=%d(%d)",
                //    start, nLeftStart, end, nLeftEnd, mid, nLeftMid);

                // go left or right--or terminate!
                if (nLeftMid > nRight) {
                        //DPRINTF("Go left!");
                        // need to go left
                        // terminate?
                        if (mid - start < 2) {
                                // we are adjacent to the start!
                                if (nLeftStart > 0) {
                                        mid = start;
                                }
                                break;
                        }

                        // midpoint is the new end point
                        end = mid;
                        nLeftEnd = nLeftMid;
                } else if (nLeftMid < nRight) {
                        //DPRINTF("Go right!");
                        // need to go right
                        // terminate?
                        if (end - mid < 2) {
                                // we are adjacent to the end!
                                if (!nLeftMid) {
                                        mid = end;
                                }
                                break;
                        }

                        // midpoint is the new start point
                        start = mid;
                        nLeftStart = nLeftMid;
                } else {
                        //DPRINTF("Equal!");
                        // nLeft == nRight!
                        break;
                }

                // new midpoint
                mid = (start + end) / 2;
        }

        // exited the search--what did we find?
        //DPRINTF("  Stopped search with idx=%d", mid);
        plane.value = dist.getValueAtIndex(mid);
        //DPRINTF("  Coordinate value: %f", plane.value);
        measureSplit(plane, vec, nLeft, nRight);
        //DPRINTF("  Final nLeft=%d, nRight=%d", nLeft, nRight);
}



static void
balancedSplit
(
IN const rect3d_t& r,
OUT plane_t& plane,
IN const Node::vec_entries_t& vec
)
throw()
{
        // loop through all 3 axes and try splitting along each one
        plane_t planes[3];
        int N = vec.size();
        ASSERT(N > 1, "Only works with 2 or more entries!");
        int iPreferred = getLongAxis(r) - 1;
        int iBest = 0;          // default is to split along x-direction!
        float dBest = 1.0;
        for (int i = 0; i < 3; ++i) {
                int nLeft, nRight;
                bestSplit(i, vec, planes[i], nLeft, nRight);

                // rLeft is the ratio of entities cleanly on the left
                float rLeft = (1.0 * nLeft) / (1.0 * N);
                float rRight = (1.0 * nRight) / (1.0 * N);

                //DPRINTF("For offset %d, rLeft=%f and rRight=%f",
                //   i, rLeft, rRight);

                // we want both rLeft and rRight to be 0.5
                // find the distance^2 from the (0.5, 0.5) value
                float dLeft = rLeft - 0.5;
                float dRight = rRight - 0.5;
                float d2 = (dLeft * dLeft) + (dRight * dRight);
                //DPRINTF("  d2 = %f", d2);
                if (d2 < dBest) {
                        iBest = i;
                        dBest = d2;
                } else if (d2 == dBest && i == iPreferred) {
                        // use preferred axis if tiebreaker
                        iBest = i;
                }
        }

        // which was best?
        // DPRINTF("Best offset/axis = %d", iBest);
        ASSERT(iBest >= 0 && iBest < 3, "Bad best offset: %d", iBest);

        // return that plane
        plane = planes[iBest];
}



/*
 * Saving these for posterity because they were nontrivial.
 *
 * But I've decided not to support arbitrary sort planes for these kd-trees,
 * and instead follow convention (and performance) by only allowing axis-aligned
 * sort planes.
 *
static void
checkCorners
(
IN const rect3d_t&r,
IN const plane_t& plane,
IN const eSort sort,
OUT rect3d_t& out
)
throw()
{
        out.clear();
        bool set = false;

        for (int i = 0; i < 8; ++i) {
                point3d_t p = r.getCorner(i);
                eSort psort = sortPoint(plane, p);
                if (eSort_Straddle == psort || sort == psort) {
                        // point qualifies!
                        if (!set) {
                                // this is the first point we've found
                                set = true;
                                out.setToPoint(p);
                        } else {
                                // make sure rect includes
                                out.includePoint(p);
                        }
                }
        }

        ASSERT(set, "none of the rectangle corners are in the partition?");
}



static void
getLeftRightBounds
(
IN const rect3d_t& r,
IN const plane_t& plane,
OUT rect3d_t& left,
OUT rect3d_t& right
)
throw()
{
        // this was a surprisingly hard function!  That is, given an axis-
        // aligned bounding box (rect3d_r) and an arbirary, split the rect into
        // left and right subrects ("left" and "right" as determined by
        // sorting).

        // A subrect must be another axis-aligned bounding box that completely
        // contains all points originally in the parent box that are now
        // partitioned on one side of the plane.  In some cases this can mean
        // that both subrects are equal to the parent rect!

        // r.dump("Splitting this rect");
        // plane.dump("with this plane");

        // first, populate with any original rect corners now partitioned
        checkCorners(r, plane, eSort_Left,  left);
        checkCorners(r, plane, eSort_Right, right);

        // include any points on rect edges
        for (int i = 0; i < 12; ++i) {
                point3d_t p0, p1;
                r.getEdge(i, p0, p1);

                // assume p(t) = (1 - t) p0 + t p1.  For what t does
                //  p(t) lie in the plane?

                float d = plane.distance(p0);
                float denom = dotProduct(plane.n, p1 - p0);
                if (denom > -s_eps && denom < s_eps) {
                        // edge is parallel to plane!
                        continue;
                }

                float t = - d / denom;
                if (t < 0 || t > 1) {
                        // t is outside edge!
                        continue;
                }

                // got here?  then we have a valid point on the edge
                point3d_t p = (1 - t) * p0 + t * p1;
                // p.dump("  edge point!");
                left.includePoint(p);
                right.includePoint(p);
        }
}
*/



static void
getLeftRightBounds
(
IN const rect3d_t& r,
IN const plane_t& plane,
OUT rect3d_t& left,
OUT rect3d_t& right
)
throw()
{
        left = r;
        right = r;

        switch (plane.axis) {
        case eAxis_X:
                left.x1 = right.x0 = plane.value;
                break;

        case eAxis_Y:
                left.y1 = right.y0 = plane.value;
                break;

        case eAxis_Z:
                left.z1 = right.z0 = plane.value;
                break;

        default:
                ASSERT(false, "Invalid axis: %d", plane.axis);
        }
}



////////////////////////////////////////////////////////////////////////////////
//
//      Node implementation
//
////////////////////////////////////////////////////////////////////////////////

Node::Node(void)
throw()
{
        m_user = NULL;
        m_left = m_right = NULL;
        m_bounds.clear();
        m_sortPlane.clear();
}



Node::~Node(void)
throw()
{
}



void
Node::initialize
(
IN const rect3d_t& bounds
)
{
        m_bounds = bounds;
}



int
Node::getTreeDepth
(
void
)
const
throw()
{
        int maxChild = (m_left) ? m_left->getTreeDepth() : 0;
        if (m_right) {
                int n = m_right->getTreeDepth();
                if (n > maxChild) {
                        maxChild = n;
                }
        }
        return maxChild + 1;
}



int
Node::getNodeCount
(
void
)
const
throw()
{
        int iLeft = (m_left) ? m_left->getNodeCount() : 0;
        int iRight = (m_right) ? m_right->getNodeCount() : 0;
        return iLeft + iRight + 1;
}



int
Node::getStaticEntryCount
(
void
)
const
throw()
{
        int iLeft = (m_left) ? m_left->getStaticEntryCount() : 0;
        int iRight = (m_right) ? m_right->getStaticEntryCount() : 0;
        return iLeft + iRight + m_entries.size();
}



////////////////////////////////////////////////////////////////////////////////
//
//      Node -- kdtree::Node class interface methods
//
////////////////////////////////////////////////////////////////////////////////

void
Node::addStaticEntry
(
IN smart_ptr<Entry>& entry,
IN const rect3d_t& r
)
{
        // ASSERT(entry) -- we don't care if it is null!

        // this had better be contained!
        if (!m_bounds.containsRect(r)) {
                m_bounds.dump("Node boundary");
                r.dump("Entity box");
                ASSERT_THROW(false,
                    "Incoming entry is not contained by kd-tree node");
        }

        if (m_left || m_right) {
                // sort!
                eSort sort = sortRect(m_sortPlane, r);
                switch (sort) {
                case eSort_Left:
                        m_left->addStaticEntry(entry, r);
                        return;
                case eSort_Right:
                        m_right->addStaticEntry(entry, r);
                        return;
                case eSort_Straddle:
                        break;  // add to this node
                default:
                        ASSERT(false, "bad sort");
                }
        }

        // no child nodes or straddle--just add to our local collection
        entry_rec_t er;
        er.entry = entry;
        er.rect = r;
        m_entries.push_back(er);
}



void
Node::addDynamicEntry
(
IN smart_ptr<dynamic_entry_t>& dyn
)
{
        ASSERT(dyn, "null");
        ASSERT(!dyn->next, "should only pass in an unlinked entry");

        // if we are a parent, attempt to recurse
        if (m_left) {
                switch (sortRect(m_sortPlane, dyn->rect)) {
                case eSort_Left:
                        m_left->addDynamicEntry(dyn);
                        return;

                case eSort_Right:
                        m_right->addDynamicEntry(dyn);
                        return;

                default:
                        // any other result, we add to our list
                        break;
                }
        }

        // incoming node is the new head of our list
        dyn->next = m_dynamic;
        m_dynamic = dyn;
}



bool
Node::removeStaticEntry
(
IN Entry * entry
)
throw()
{
        ASSERT(entry, "null");

        // TODO: make this more efficient?  Set?  List?
        // Client use cases vary.
        for (vec_entries_t::iterator i = m_entries.begin();
             i != m_entries.end(); ++i) {
                entry_rec_t& er = *i;
                if (entry == er.entry) {
                        m_entries.erase(i);
                        return true;
                }
        }

        // not found!
        return false;
}



smart_ptr<Node::dynamic_entry_t>
Node::popDynamicEntry
(
void
)
throw()
{
        smart_ptr<dynamic_entry_t> dyn = m_dynamic;
        if (dyn) {
                m_dynamic = dyn->next;
                dyn->next = NULL;
        }
        return dyn;
}



eAction
Node::iterateStaticEntries
(
IN const rect3d_t& r,
IN iteration_callback callback,
IN void * context
)
{
        ASSERT(r.isValid(), "invalid");
        ASSERT(callback, "null");
        // ASSERT(context) -- we don't care

        // always iterate through any entities here
        for (vec_entries_t::iterator i = m_entries.begin();
             i != m_entries.end();) {
                vec_entries_t::iterator next = i;
                ++next;

                entry_rec_t& er = *i;
                if (r.intersectsRect(er.rect)) {
                        // e.bounds.dump("Intersects!");
                        // entity overlaps with search rect!
                        eAction act = callback(context, er.entry, er.rect);
                        if (eAction_Remove & act) {
                                m_entries.erase(i);
                        }
                        if (eAction_Halt & act)
                                return eAction_Halt;
                }

                // next!
                i = next;
        }

        // anything to recurse?
        if (!m_left)
                return eAction_Continue;

        // see what to recurse
        eSort sort = sortRect(m_sortPlane, r);
//      DPRINTF("Sorted: %s", getSortName(sort));
        if (eSort_Left == sort || eSort_Straddle == sort) {
                eAction act = m_left->iterateStaticEntries(r, callback, context);
                if (eAction_Halt & act)
                        return eAction_Halt;
        }
        if (eSort_Right == sort || eSort_Straddle == sort) {
                eAction act = m_right->iterateStaticEntries(r, callback, context);
                if (eAction_Halt & act)
                        return eAction_Halt;
        }

        // success!
        return eAction_Continue;
}



////////////////////////////////////////////////////////////////////////////////
//
//      Node -- private helper methods
//
////////////////////////////////////////////////////////////////////////////////

void
Node::subdivide
(
IN eStrategy strategy,
IN float param
)
{
        // boring leaf?  then return right away
        int minSize = (param < 2) ? 2 : param;
        if (!m_left && (int) m_entries.size() <= minSize) {
                return;
        }

        // do we need to split this node?
        if (!m_left) {
                this->subdivideInternal(strategy, param);
        }

        // sort!
        vec_entries_t vecNew;
        for (vec_entries_t::iterator i = m_entries.begin();
             i != m_entries.end(); ++i) {
                entry_rec_t& er = *i;

                // does this entry go in a child?
                eSort sort = sortRect(m_sortPlane, er.rect);
                if (eSort_Left == sort) {
                        m_left->addStaticEntry(er.entry, er.rect);
                } else if (eSort_Right == sort) {
                        m_right->addStaticEntry(er.entry, er.rect);
                } else if (eSort_Straddle == sort) {
                         vecNew.push_back(er);
                } else {
                        ASSERT(false, "bad sort result");
                }
        }

        // nuke old collection, reset to new
        m_entries.clear();
        m_entries = vecNew;

        // subdivide if it will help
        int nLeft = m_left->getStaticEntryCount();
        int nRight = m_right->getStaticEntryCount();
        if (!nLeft || !nRight) {
                // degenerate--no point in subdividing
                return;
        }

        // may be worth subdividing
        m_left->subdivide(strategy, param);
        m_right->subdivide(strategy, param);

//      DPRINTF("After subdivide:");
//      DPRINTF("  Tree depth: %d", this->getTreeDepth());
//      DPRINTF("  Node count: %d", this->getNodeCount());
//      DPRINTF("  Entry count: %d", this->getEntryCount());
//      DPRINTF("  Our own entry list %d", m_entries.size());
}



void
Node::subdivideInternal
(
IN eStrategy strategy,
IN float param
)
{
        ASSERT(!m_left && !m_right, "already have children?");

        // determine boundary
        switch (strategy) {
        case eStrategy_Naive:
                simpleSplit(m_bounds, m_sortPlane);
                break;

        case eStrategy_Balance:
                balancedSplit(m_bounds, m_sortPlane, m_entries);
                break;

        default:
                ASSERT(false, "Unknown strategy: %d", strategy);
        }
        //m_sortPlane.dump("sort plane");

        // get left and right child boundaries
        rect3d_t rLeft, rRight;
        getLeftRightBounds(m_bounds, m_sortPlane, rLeft, rRight);       
        //rLeft.dump("  left");
        //rRight.dump("  right");
        ASSERT(rLeft.isValid(), "invalid");
        ASSERT(rRight.isValid(), "invalid");

        // create left and right child nodes
        m_left = Node::create(rLeft);
        ASSERT(m_left, "failed to create left child");
        m_right = Node::create(rRight);
        ASSERT(m_right, "failed to create right child");
}



////////////////////////////////////////////////////////////////////////////////
//
//      public API
//
////////////////////////////////////////////////////////////////////////////////

smart_ptr<Node>
Node::create
(
IN const rect3d_t& bounds
)
{
        ASSERT(bounds.isValid(), "invalid");

        smart_ptr<Node> local = new Node;
        ASSERT(local, "out of memory");

        local->initialize(bounds);

        return local;
}



};      // kdtree namespace