Fast implementation of QuadTree

src/core/SkQuadTree.cpp

Index: src/core/SkQuadTree.cpp
diff --git a/src/core/SkQuadTree.cpp b/src/core/SkQuadTree.cpp
index 97f86cfb5add9257af5684e857b1720bd84ded70..8531823278b111b29d9f7beef8029236a9500e29 100644
--- a/src/core/SkQuadTree.cpp
+++ b/src/core/SkQuadTree.cpp
@@ -10,188 +10,157 @@
 #include <stdio.h>
 #include <vector>
 
-class SkQuadTree::QuadTreeNode {
-public:
-    struct Data {
-        Data(const SkIRect& bounds, void* data) : fBounds(bounds), fInnerBounds(bounds), fData(data) {}
-        SkIRect fBounds;
-        SkIRect fInnerBounds;
-        void* fData;
-    };
-
-    QuadTreeNode(const SkIRect& bounds)
-     : fBounds(bounds)
-     , fTopLeft(NULL)
-     , fTopRight(NULL)
-     , fBottomLeft(NULL)
-     , fBottomRight(NULL)
-     , fCanSubdivide((fBounds.width() * fBounds.height()) > 0) {}
-
-    ~QuadTreeNode() {
-        clear();
-    }
-
-    void clear() {
-        SkDELETE(fTopLeft);
-        fTopLeft = NULL;
-        SkDELETE(fTopRight);
-        fTopRight = NULL;
-        SkDELETE(fBottomLeft);
-        fBottomLeft = NULL;
-        SkDELETE(fBottomRight);
-        fBottomRight = NULL;
-        fData.reset();
-    }
-
-    const SkIRect& getBounds() const { return fBounds; }
-
-    // Insert data into the QuadTreeNode
-    bool insert(Data& data) {
-        // Ignore objects which do not belong in this quad tree
-        return data.fInnerBounds.intersect(fBounds) && doInsert(data);
-    }
-
-    // Find all data which appear within a range
-    void queryRange(const SkIRect& range, SkTDArray<void*>* dataInRange) const {
-        // Automatically abort if the range does not collide with this quad
-        if (!SkIRect::Intersects(fBounds, range)) {
-            return; // nothing added to the list
-        }
+static const int kSplitThreshold = 8;
+static const int kMinDimensions = 128;
 
-        // Check objects at this quad level
-        for (int i = 0; i < fData.count(); ++i) {
-            if (SkIRect::Intersects(fData[i].fBounds, range)) {
-                dataInRange->push(fData[i].fData);
-            }
-        }
+SkQuadTree::SkQuadTree(const SkIRect& bounds)
+    : fEntryCount(0)
+    , fRoot(NULL) {
+    SkASSERT((bounds.width() * bounds.height()) > 0);
+    fRoot = fNodePool.acquire();
+    fRoot->fBounds = bounds;
+}
 
-        // Terminate here, if there are no children
-        if (!hasChildren()) {
-            return;
-        }
+SkQuadTree::~SkQuadTree() {
+}
 
-        // Otherwise, add the data from the children
-        fTopLeft->queryRange(range, dataInRange);
-        fTopRight->queryRange(range, dataInRange);
-        fBottomLeft->queryRange(range, dataInRange);
-        fBottomRight->queryRange(range, dataInRange);
+SkQuadTree::Node* SkQuadTree::pickChild(Node* node,
+                                        const SkIRect& bounds) const {
+    // is it entirely to the left?
+    int index = 0;
+    if (bounds.fRight < node->fSplitPoint.fX) {
+        // Inside the left side
+    } else if(bounds.fLeft >= node->fSplitPoint.fX) {
+        // Inside the right side
+        index |= 1;
+    } else {
+        // Not inside any children
+        return NULL;
     }
-
-    int getDepth(int i = 1) const {
-        if (hasChildren()) {
-            int depthTL = fTopLeft->getDepth(++i);
-            int depthTR = fTopRight->getDepth(i);
-            int depthBL = fBottomLeft->getDepth(i);
-            int depthBR = fBottomRight->getDepth(i);
-            return SkTMax(SkTMax(depthTL, depthTR), SkTMax(depthBL, depthBR));
-        }
-        return i;
+    if (bounds.fBottom < node->fSplitPoint.fY) {
+        // Inside the top side
+    } else if(bounds.fTop >= node->fSplitPoint.fY) {
+        // Inside the bottom side
+        index |= 2;
+    } else {
+        // Not inside any children
+        return NULL;
     }
+    return node->fChildren[index];
+}
 
-    void rewindInserts(SkBBoxHierarchyClient* client) {
-        for (int i = fData.count() - 1; i >= 0; --i) {
-            if (client->shouldRewind(fData[i].fData)) {
-                fData.remove(i);
-            }
-        }
-        if (hasChildren()) {
-            fTopLeft->rewindInserts(client);
-            fTopRight->rewindInserts(client);
-            fBottomLeft->rewindInserts(client);
-            fBottomRight->rewindInserts(client);
+void SkQuadTree::insert(Node* node, Entry* entry) {
+    // does it belong in a child?
+    if (NULL != node->fChildren[0]) {
+        Node* child = pickChild(node, entry->fBounds);
+        if (NULL != child) {
+            insert(child, entry);
+        } else {
+            node->fEntries.push(entry);
         }
+        return;
     }
-
-private:
-    // create four children which fully divide this quad into four quads of equal area
-    void subdivide() {
-        if (!hasChildren() && fCanSubdivide) {
-            SkIPoint center = SkIPoint::Make(fBounds.centerX(), fBounds.centerY());
-            fTopLeft = SkNEW_ARGS(QuadTreeNode, (SkIRect::MakeLTRB(
-                fBounds.fLeft, fBounds.fTop, center.fX, center.fY)));
-            fTopRight = SkNEW_ARGS(QuadTreeNode, (SkIRect::MakeLTRB(
-                center.fX, fBounds.fTop, fBounds.fRight, center.fY)));
-            fBottomLeft = SkNEW_ARGS(QuadTreeNode, (SkIRect::MakeLTRB(
-                fBounds.fLeft, center.fY, center.fX, fBounds.fBottom)));
-            fBottomRight = SkNEW_ARGS(QuadTreeNode, (SkIRect::MakeLTRB(
-                center.fX, center.fY, fBounds.fRight, fBounds.fBottom)));
-
-            // If any of the data can fit entirely into a subregion, move it down now
-            for (int i = fData.count() - 1; i >= 0; --i) {
-                // If the data fits entirely into one of the 4 subregions, move that data
-                // down to that subregion.
-                if (fTopLeft->doInsert(fData[i]) ||
-                    fTopRight->doInsert(fData[i]) ||
-                    fBottomLeft->doInsert(fData[i]) ||
-                    fBottomRight->doInsert(fData[i])) {
-                    fData.remove(i);
-                }
-            }
-        }
+    // No children yet, add to this node
+    node->fEntries.push(entry);
+    // should I split?
+    if (node->fEntries.getCount() < kSplitThreshold) {
+        return;
     }
 
-    bool doInsert(const Data& data) {
-        if (!fBounds.contains(data.fInnerBounds)) {
-            return false;
-        }
+    if ((node->fBounds.width() < kMinDimensions) ||
+        (node->fBounds.height() < kMinDimensions)) {
+        return;
+    }
 
-        if (fData.count() > kQuadTreeNodeCapacity) {
-            subdivide();
-        }
+    // Build all the children
+    node->fSplitPoint = SkIPoint::Make(node->fBounds.centerX(),
+                                       node->fBounds.centerY());
+    for(int index=0; index<kChildCount; ++index) {
+        node->fChildren[index] = fNodePool.acquire();
+    }
+    node->fChildren[0]->fBounds = SkIRect::MakeLTRB(
+        node->fBounds.fLeft,    node->fBounds.fTop,
+        node->fSplitPoint.fX,   node->fSplitPoint.fY);
+    node->fChildren[1]->fBounds = SkIRect::MakeLTRB(
+        node->fSplitPoint.fX,   node->fBounds.fTop,
+        node->fBounds.fRight,   node->fSplitPoint.fY);
+    node->fChildren[2]->fBounds = SkIRect::MakeLTRB(
+        node->fBounds.fLeft,    node->fSplitPoint.fY,
+        node->fSplitPoint.fX,   node->fBounds.fBottom);
+    node->fChildren[3]->fBounds = SkIRect::MakeLTRB(
+        node->fSplitPoint.fX,   node->fSplitPoint.fY,
+        node->fBounds.fRight,   node->fBounds.fBottom);
+    // reinsert all the entries of this node to allow child trickle
+    SkTInternalSList<Entry> entries;
+    entries.pushAll(&node->fEntries);
+    while(!entries.isEmpty()) {
+        insert(node, entries.pop());
+    }
+}
 
-        // If there is space in this quad tree, add the object here
-        // If this quadtree can't be subdivided, we have no choice but to add it here
-        if ((fData.count() <= kQuadTreeNodeCapacity) || !fCanSubdivide) {
-            if (fData.isEmpty()) {
-                fData.setReserve(kQuadTreeNodeCapacity);
-            }
-            fData.push(data);
-        } else if (!fTopLeft->doInsert(data) && !fTopRight->doInsert(data) &&
-                   !fBottomLeft->doInsert(data) && !fBottomRight->doInsert(data)) {
-            // Can't be pushed down to children ? keep it here
-            fData.push(data);
+void SkQuadTree::search(Node* node, const SkIRect& query,
+                        SkTDArray<void*>* results) const {
+    for (Entry* entry = node->fEntries.head(); NULL != entry;
+        entry = entry->getSListNext()) {
+        if (SkIRect::IntersectsNoEmptyCheck(entry->fBounds, query)) {
+            results->push(entry->fData);
         }
-
-        return true;
     }
-
-    bool hasChildren() const {
-        return (NULL != fTopLeft);
+    if (NULL == node->fChildren[0]) {
+        return;
+    }
+    // fast quadrant test
+    bool left = true;
+    bool right = true;
+    if (query.fRight < node->fSplitPoint.fX) {
+        right = false;
+    } else if(query.fLeft >= node->fSplitPoint.fX) {
+        left = false;
+    }
+    bool top = true;
+    bool bottom = true;
+    if (query.fBottom < node->fSplitPoint.fY) {
+        bottom = false;
+    } else if(query.fTop >= node->fSplitPoint.fY) {
+        top = false;
+    }
+    // search all the active quadrants
+    if (top && left) {
+        search(node->fChildren[0], query, results);
+    }
+    if (top && right) {
+        search(node->fChildren[1], query, results);
+    }
+    if (bottom && left) {
+        search(node->fChildren[2], query, results);
+    }
+    if (bottom && right) {
+        search(node->fChildren[3], query, results);
     }
-
-    // Arbitrary constant to indicate how many elements can be stored in this quad tree node
-    enum { kQuadTreeNodeCapacity = 4 };
-
-    // Bounds of this quad tree
-    SkIRect fBounds;
-
-    // Data in this quad tree node
-    SkTDArray<Data> fData;
-
-    // Children
-    QuadTreeNode* fTopLeft;
-    QuadTreeNode* fTopRight;
-    QuadTreeNode* fBottomLeft;
-    QuadTreeNode* fBottomRight;
-
-    // Whether or not this node can have children
-    bool fCanSubdivide;
-};
-
-///////////////////////////////////////////////////////////////////////////////////////////////////
-
-SkQuadTree* SkQuadTree::Create(const SkIRect& bounds) {
-    return new SkQuadTree(bounds);
 }
 
-SkQuadTree::SkQuadTree(const SkIRect& bounds)
-    : fCount(0)
-    , fRoot(SkNEW_ARGS(QuadTreeNode, (bounds))) {
-    SkASSERT((bounds.width() * bounds.height()) > 0);
+void SkQuadTree::clear(Node* node) {
+    // first clear the entries of this node
+    fEntryPool.releaseAll(&node->fEntries);
+    // recurse into and clear all child nodes
+    for(int index=0; index<kChildCount; ++index) {
+        Node* child = node->fChildren[index];
+        node->fChildren[index] = NULL;
+        if (NULL != child) {
+            clear(child);
+            fNodePool.release(child);
+        }
+    }
 }
 
-SkQuadTree::~SkQuadTree() {
-    SkDELETE(fRoot);
+int SkQuadTree::getDepth(Node* node) const {
+    int maxDepth = 0;
+    if (NULL != node->fChildren[0]) {
+        for(int index=0; index<kChildCount; ++index) {
+            maxDepth = SkMax32(maxDepth, getDepth(node->fChildren[index]));
+        }
+    }
+    return maxDepth + 1;
 }
 
 void SkQuadTree::insert(void* data, const SkIRect& bounds, bool) {
@@ -199,27 +168,54 @@ void SkQuadTree::insert(void* data, const SkIRect& bounds, bool) {
         SkASSERT(false);
         return;
     }
-
-    QuadTreeNode::Data quadTreeData(bounds, data);
-    fRoot->insert(quadTreeData);
-    ++fCount;
+    Entry* entry = fEntryPool.acquire();
+    entry->fData = data;
+    entry->fBounds = bounds;
+    ++fEntryCount;
+    if (fRoot->fEntries.isEmpty() && (NULL == fRoot->fChildren[0])) {
+        fDeferred.push(entry);
+    } else {
+        insert(fRoot, entry);
+    }
 }
 
 void SkQuadTree::search(const SkIRect& query, SkTDArray<void*>* results) {
+    SkASSERT(fDeferred.isEmpty());
     SkASSERT(NULL != results);
-    fRoot->queryRange(query, results);
+    if (SkIRect::Intersects(fRoot->fBounds, query)) {
+        search(fRoot, query, results);
+    }
 }
 
 void SkQuadTree::clear() {
-    fCount = 0;
-    fRoot->clear();
+    fEntryCount = 0;
+    clear(fRoot);
 }
 
 int SkQuadTree::getDepth() const {
-    return fRoot->getDepth();
+    return getDepth(fRoot);
 }
 
 void SkQuadTree::rewindInserts() {
     SkASSERT(fClient);
-    fRoot->rewindInserts(fClient);
+     // Currently only supports deferred inserts
+    SkASSERT(fRoot->fEntries.isEmpty() && fRoot->fChildren[0] == NULL);
+    SkTInternalSList<Entry> entries;
+    entries.pushAll(&fDeferred);
+    while(!entries.isEmpty()) {
+        Entry* entry = entries.pop();
+        if (fClient->shouldRewind(entry->fData)) {
+            entry->fData = NULL;
+            fEntryPool.release(entry);
+            --fEntryCount;
+        } else {
+            fDeferred.push(entry);
+        }
+    }
+}
+
+void SkQuadTree::flushDeferredInserts() {
+    while(!fDeferred.isEmpty()) {
+        insert(fRoot, fDeferred.pop());
+    }
 }