SERVER-15562 Search neighbors for density estimation with 2d index

1 files changed, 63 insertions, 0 deletions

diff --git a/src/mongo/db/geo/hash.cpp b/src/mongo/db/geo/hash.cpp
index ea481e14e14..6daf1e2744b 100644
--- a/src/mongo/db/geo/hash.cpp
+++ b/src/mongo/db/geo/hash.cpp
@@ -496,6 +496,69 @@ namespace mongo {
         return GeoHash(_hash, _bits - 1);
     }
 
+
+    void GeoHash::appendVertexNeighbors(unsigned level, vector<GeoHash>* output) const {
+        invariant(level >= 0 && level < _bits);
+
+        // Parent at the given level.
+        GeoHash parentHash = parent(level);
+        output->push_back(parentHash);
+
+        // Generate the neighbors of parent that are closest to me.
+        unsigned px, py, parentBits;
+        parentHash.unhash(&px, &py);
+        parentBits = parentHash.getBits();
+
+        // No Neighbors for the top level.
+        if (parentBits == 0U) return;
+
+        // Position in parent
+        // Y
+        // ^
+        // |  01, 11
+        // |  00, 10
+        // +----------> X
+        // We can guarantee _bits > 0.
+        long long posInParent = (_hash >> (64 - 2 * (parentBits + 1))) & 3LL;
+
+        // 1 bit at parent's level, the least significant bit of parent.
+        unsigned parentMask = 1U << (32 - parentBits);
+
+        // Along X Axis
+        if ((posInParent & 2LL) == 0LL) {
+            // Left side of parent, X - 1
+            if (!parentHash.atMinX()) output->push_back(GeoHash(px - parentMask, py, parentBits));
+        } else {
+            // Right side of parent, X + 1
+            if (!parentHash.atMaxX()) output->push_back(GeoHash(px + parentMask, py, parentBits));
+        }
+
+        // Along Y Axis
+        if ((posInParent & 1LL) == 0LL) {
+            // Bottom of parent, Y - 1
+            if (!parentHash.atMinY()) output->push_back(GeoHash(px, py - parentMask, parentBits));
+        } else {
+            // Top of parent, Y + 1
+            if (!parentHash.atMaxY()) output->push_back(GeoHash(px, py + parentMask, parentBits));
+        }
+
+        // Four corners
+        if (posInParent == 0LL) {
+            if (!parentHash.atMinX() && !parentHash.atMinY())
+                output->push_back(GeoHash(px - parentMask, py - parentMask, parentBits));
+        } else if (posInParent == 1LL) {
+            if (!parentHash.atMinX() && !parentHash.atMaxY())
+                output->push_back(GeoHash(px - parentMask, py + parentMask, parentBits));
+        } else if (posInParent == 2LL) {
+            if (!parentHash.atMaxX() && !parentHash.atMinY())
+                output->push_back(GeoHash(px + parentMask, py - parentMask, parentBits));
+        } else {
+            // PosInParent == 3LL
+            if (!parentHash.atMaxX() && !parentHash.atMaxY())
+                output->push_back(GeoHash(px + parentMask, py + parentMask, parentBits));
+        }
+    }
+
     static BSONField<int> bitsField("bits", 26);
     static BSONField<double> maxField("max", 180.0);
     static BSONField<double> minField("min", -180.0);