enable within for indexed 2dsphere

7 files changed, 190 insertions, 44 deletions

diff --git a/jstests/geo_s2index.js b/jstests/geo_s2index.js
index 6fc6e87c653..0cca1ce0bbb 100755
--- a/jstests/geo_s2index.js
+++ b/jstests/geo_s2index.js
@@ -41,6 +41,9 @@ assert.eq(res.count(), 5);
 res = t.find({ "geo" : { "$geoIntersects" : { "$geometry" : somepoly} } })
 assert.eq(res.count(), 6);
 
+res = t.find({ "geo" : { "$within" : { "$geometry" : somepoly} } })
+assert.eq(res.count(), 6);
+
 res = t.find({ "geo" : { "$geoIntersects" : { "$geometry" : somepoly} } }).limit(1)
 assert.eq(res.itcount(), 1);
 
diff --git a/jstests/geo_s2within.js b/jstests/geo_s2within.js
new file mode 100644
index 00000000000..50585b85986
--- /dev/null
+++ b/jstests/geo_s2within.js
@@ -0,0 +1,36 @@
+// Test some cases that might be iffy with $within.
+t = db.geo_s2within
+t.drop()
+t.ensureIndex({geo: "2dsphere"})
+
+somepoly = { "type" : "Polygon",
+             "coordinates" : [ [ [40,5], [40,6], [41,6], [41,5], [40,5]]]}
+
+t.insert({geo: { "type" : "LineString", "coordinates": [ [ 40.1, 5.1], [40.2, 5.2]]}})
+// This is only partially contained within the polygon.
+t.insert({geo: { "type" : "LineString", "coordinates": [ [ 40.1, 5.1], [42, 7]]}})
+
+res = t.find({ "geo" : { "$within" : { "$geometry" : somepoly} } })
+assert.eq(res.count(), 1);
+
+t.drop()
+t.ensureIndex({geo: "2dsphere"})
+somepoly = { "type" : "Polygon",
+             "coordinates" : [ [ [40,5], [40,8], [43,8], [43,5], [40,5]],
+                               [ [41,6], [42,6], [42,7], [41,7], [41,6]]]}
+
+t.insert({geo:{ "type" : "Point", "coordinates": [ 40, 5 ] }})
+res = t.find({ "geo" : { "$within" : { "$geometry" : somepoly} } })
+assert.eq(res.count(), 1);
+// In the hole.  Shouldn't find it.
+t.insert({geo:{ "type" : "Point", "coordinates": [ 41.1, 6.1 ] }})
+res = t.find({ "geo" : { "$within" : { "$geometry" : somepoly} } })
+assert.eq(res.count(), 1);
+// Also in the hole.
+t.insert({geo: { "type" : "LineString", "coordinates": [ [ 41.1, 6.1], [41.2, 6.2]]}})
+res = t.find({ "geo" : { "$within" : { "$geometry" : somepoly} } })
+assert.eq(res.count(), 1);
+// Half-hole, half-not.  Shouldn't be $within.
+t.insert({geo: { "type" : "LineString", "coordinates": [ [ 41.5, 6.5], [42.5, 7.5]]}})
+res = t.find({ "geo" : { "$within" : { "$geometry" : somepoly} } })
+assert.eq(res.count(), 1);
diff --git a/src/mongo/db/geo/s2common.cpp b/src/mongo/db/geo/s2common.cpp
index 413c8d324db..4d1143e27e0 100644
--- a/src/mongo/db/geo/s2common.cpp
+++ b/src/mongo/db/geo/s2common.cpp
@@ -44,12 +44,12 @@ namespace mongo {
             // there's nothing there.
             considerCoarser = false;
             coverer->GetCovering(region, &cover);
-            warning() << "Trying to create BSON indexing obj w/too many regions = " << oldCoverSize
-                      << endl;
-            warning() << "Modifying coverer from (" << coverer->max_level() << "," << oldMaxLevel
-                      << ") to (" << coverer->min_level() << "," << coverer->max_level() << ")"
-                      << endl;
-            warning() << "New #regions = " << cover.size() << endl;
+            LOG(2) << "Trying to create BSON indexing obj w/too many regions = " << oldCoverSize
+                   << endl;
+            LOG(2) << "Modifying coverer from (" << coverer->max_level() << "," << oldMaxLevel
+                   << ") to (" << coverer->min_level() << "," << coverer->max_level() << ")"
+                   << endl;
+            LOG(2) << "New #regions = " << cover.size() << endl;
         }
 
         // Look at the cells we cover and all cells that are within our covering and
@@ -96,16 +96,48 @@ namespace mongo {
     string QueryGeometry::toString() const {
         stringstream ss;
         ss << "field = " << field;
+        ss << ", predicate = " << ((WITHIN == predicate) ? "within" : "intersect");
         if (NULL != cell.get()) {
             ss << ", cell";
         } else if (NULL != line.get()) {
-            ss << ", line = ";
+            ss << ", line";
         } else if (NULL != polygon.get()) {
-            ss << ", polygon = ";
+            ss << ", polygon";
         }
         return ss.str();
     }
 
+    bool QueryGeometry::satisfiesPredicate(const BSONObj &obj) {
+        verify(predicate == WITHIN || predicate == INTERSECT);
+
+        if (GeoParser::isPolygon(obj)) {
+            S2Polygon shape;
+            GeoParser::parsePolygon(obj, &shape);
+            if (WITHIN == predicate) {
+                return isWithin(shape);
+            } else {
+                return intersects(shape);
+            }
+        } else if (GeoParser::isLineString(obj)) {
+            S2Polyline shape;
+            GeoParser::parseLineString(obj, &shape);
+            if (WITHIN == predicate) {
+                return isWithin(shape);
+            } else {
+                return intersects(shape);
+            }
+        } else if (GeoParser::isPoint(obj)) {
+            S2Cell point;
+            GeoParser::parsePoint(obj, &point);
+            if (WITHIN == predicate) {
+                return isWithin(point);
+            } else {
+                return intersects(point);
+            }
+        }
+        return false;
+    }
+
     bool QueryGeometry::parseFrom(const BSONObj& obj) {
         if (GeoParser::isGeoJSONPolygon(obj)) {
             // We can't really pass these things around willy-nilly except by ptr.
@@ -123,8 +155,48 @@ namespace mongo {
         return true;
     }
 
+    // Is the geometry provided as an argument within our query geometry?
+    bool QueryGeometry::isWithin(const S2Cell &otherPoint) {
+        // Intersecting a point is containing a point.  Hooray!
+        return intersects(otherPoint);
+    }
+
+    bool QueryGeometry::isWithin(const S2Polyline& otherLine) {
+        if (NULL != cell) {
+            // Points don't contain lines.
+            return false;
+        } else if (NULL != line) {
+            // Doing line-in-line is scary.
+            return false;
+        } else {
+            // Kind of a mess.  We get a function for clipping the line to the
+            // polygon.  We do this and make sure the line is the same as the
+            // line we're clipping against.
+            vector<S2Polyline*> clipped;
+            polygon->IntersectWithPolyline(&otherLine, &clipped);
+            if (1 != clipped.size()) { return false; }
+            // If the line is entirely contained within the polygon, we should be
+            // getting it back verbatim, so really there should be no error.
+            bool ret = clipped[0]->NearlyCoversPolyline(otherLine, S1Angle::Degrees(1e-10));
+            for (size_t i = 0; i < clipped.size(); ++i) delete clipped[i];
+            return ret;
+        }
+    }
+
+    bool QueryGeometry::isWithin(const S2Polygon& otherPolygon) {
+        if (NULL != cell) {
+            // Points don't contain polygons.
+            return false;
+        } else if (NULL != line) {
+            // Lines don't contain polygons
+            return false;
+        } else {
+            return polygon->Contains(&otherPolygon);
+        }
+    }
+
     // Does this (QueryGeometry) intersect the provided data?
-    bool QueryGeometry::intersectsPoint(const S2Cell &otherPoint) {
+    bool QueryGeometry::intersects(const S2Cell &otherPoint) {
         if (NULL != cell) {
             return cell->MayIntersect(otherPoint);
         } else if (NULL != line) {
@@ -134,7 +206,7 @@ namespace mongo {
         }
     }
 
-    bool QueryGeometry::intersectsLine(const S2Polyline& otherLine) {
+    bool QueryGeometry::intersects(const S2Polyline& otherLine) {
         if (NULL != cell) {
             return otherLine.MayIntersect(*cell);
         } else if (NULL != line) {
@@ -148,10 +220,9 @@ namespace mongo {
             for (size_t i = 0; i < clipped.size(); ++i) delete clipped[i];
             return ret;
         }
-        return false;
     }
 
-    bool QueryGeometry::intersectsPolygon(const S2Polygon& otherPolygon) {
+    bool QueryGeometry::intersects(const S2Polygon& otherPolygon) {
         if (NULL != cell) {
             return otherPolygon.MayIntersect(*cell);
         } else if (NULL != line) {
@@ -168,9 +239,6 @@ namespace mongo {
     }
 
     S2Point QueryGeometry::getCentroid() const {
-        // TODO(hk): If the projection is planar this isn't valid.  Fix this
-        // when we actually use planar projections, or remove planar projections
-        // from the code.
         if (NULL != cell) {
             return cell->GetCenter();
         } else if (NULL != line) {
diff --git a/src/mongo/db/geo/s2common.h b/src/mongo/db/geo/s2common.h
index 7c713c44dab..b43da15de5b 100644
--- a/src/mongo/db/geo/s2common.h
+++ b/src/mongo/db/geo/s2common.h
@@ -36,7 +36,11 @@ namespace mongo {
 
     // Used for passing geo data from the newCursor entry point to the S2Cursor class.
     struct QueryGeometry {
-        QueryGeometry(const string& f) : field(f) {} //, cell(NULL), line(NULL), polygon(NULL) {}
+        QueryGeometry(const string& f) : field(f), predicate(INTERSECT) {}
+        enum Predicate {
+            WITHIN,
+            INTERSECT,
+        };
 
         // Name of the field in the query.
         string field;
@@ -46,14 +50,22 @@ namespace mongo {
         shared_ptr<S2Cell> cell;
         shared_ptr<S2Polyline> line;
         shared_ptr<S2Polygon> polygon;
+        Predicate predicate;
 
         string toString() const;
+
+        bool satisfiesPredicate(const BSONObj &obj);
         
         // Does this QueryGeometry intersect the provided data?  Sadly there is no common good way
         // to check this, so we do different things for all query/data pairs.
-        bool intersectsPoint(const S2Cell& otherPoint);
-        bool intersectsLine(const S2Polyline& otherLine);
-        bool intersectsPolygon(const S2Polygon& otherPolygon);
+        bool intersects(const S2Cell& otherPoint);
+        bool intersects(const S2Polyline& otherLine);
+        bool intersects(const S2Polygon& otherPolygon);
+        // And, within.
+        bool isWithin(const S2Cell& otherPoint);
+        bool isWithin(const S2Polyline& otherLine);
+        bool isWithin(const S2Polygon& otherPolygon);
+
         // One region is not NULL and this returns it.
         const S2Region& getRegion() const;
         // Get the centroid, boring if we're a point, interesting if we're not.
diff --git a/src/mongo/db/geo/s2cursor.cpp b/src/mongo/db/geo/s2cursor.cpp
index 3c308dba89f..17fac6e7b30 100644
--- a/src/mongo/db/geo/s2cursor.cpp
+++ b/src/mongo/db/geo/s2cursor.cpp
@@ -111,23 +111,10 @@ namespace mongo {
                 bool match = false;
 
                 for (BSONElementSet::iterator oi = geoFieldElements.begin();
-                     oi != geoFieldElements.end(); ++oi) {
+                     !match && (oi != geoFieldElements.end()); ++oi) {
                     if (!oi->isABSONObj()) { continue; }
                     const BSONObj &geoObj = oi->Obj();
-                    if (GeoParser::isPolygon(geoObj)) {
-                        S2Polygon shape;
-                        GeoParser::parsePolygon(geoObj, &shape);
-                        match = _fields[i].intersectsPolygon(shape);
-                    } else if (GeoParser::isLineString(geoObj)) {
-                        S2Polyline shape;
-                        GeoParser::parseLineString(geoObj, &shape);
-                        match = _fields[i].intersectsLine(shape);
-                    } else if (GeoParser::isPoint(geoObj)) {
-                        S2Cell point;
-                        GeoParser::parsePoint(geoObj, &point);
-                        match = _fields[i].intersectsPoint(point);
-                    }
-                    if (match) break;
+                    match = _fields[i].satisfiesPredicate(geoObj);
                 }
 
                 if (match) { ++geoFieldsMatched; }
diff --git a/src/mongo/db/geo/s2index.cpp b/src/mongo/db/geo/s2index.cpp
index 1b789aeddde..9f964027ebc 100644
--- a/src/mongo/db/geo/s2index.cpp
+++ b/src/mongo/db/geo/s2index.cpp
@@ -138,10 +138,10 @@ namespace mongo {
         }
 
         bool parseLegacy(const BSONObj &obj, QueryGeometry *out, bool *isNear, bool *intersect,
-                         double *maxDistance) const {
-            // Legacy intersect parsing: t.find({ loc : [0,0] })
+                         bool *within, double *maxDistance) const {
+            // Legacy within parsing: t.find({ loc : [0,0] })
             if (out->parseFrom(obj)) {
-                *isNear = true;
+                *within = true;
                 return true;
             }
 
@@ -168,10 +168,11 @@ namespace mongo {
         }
 
         bool parseQuery(const BSONObj &obj, QueryGeometry *out, bool *isNear, bool *intersect,
-                        double *maxDistance) const {
+                        bool *within, double *maxDistance) const {
             // pointA = { "type" : "Point", "coordinates": [ 40, 5 ] }
             // t.find({ "geo" : { "$intersect" : { "$geometry" : pointA} } })
             // t.find({ "geo" : { "$near" : { "$geometry" : pointA, $maxDistance : 20 }}})
+            // t.find({ "geo" : { "$within" : { "$geometry" : polygon } } })
             // where field.name is "geo"
             BSONElement e = obj.firstElement();
             if (!e.isABSONObj()) { return false; }
@@ -181,6 +182,8 @@ namespace mongo {
                 *intersect = true;
             } else if (BSONObj::opNEAR == matchType) {
                 *isNear = true;
+            } else if (BSONObj::opWITHIN == matchType) {
+                *within = true;
             } else {
                 return false;
             }
@@ -212,8 +215,11 @@ namespace mongo {
                                              int numWanted) const {
             vector<QueryGeometry> regions;
             double maxDistance = std::numeric_limits<double>::max();
+            // TODO(hk): Move this stuff into QueryGeometry and then double-check that
+            // the geometries in the query are all consistent.
             bool isNear = false;
             bool isIntersect = false;
+            bool isWithin = false;
 
             // Go through the fields that we index, and for each geo one, make a QueryGeometry
             // object for the S2Cursor class to do intersection testing/cover generating with.
@@ -227,17 +233,46 @@ namespace mongo {
                 BSONObj obj = e.Obj();
 
                 QueryGeometry geoQueryField(field.name);
-                if (parseLegacy(obj, &geoQueryField, &isNear, &isIntersect, &maxDistance)) {
+                if (parseLegacy(obj, &geoQueryField, &isNear, &isIntersect, &isWithin, &maxDistance)) {
                     regions.push_back(geoQueryField);
-                } else if (parseQuery(obj, &geoQueryField, &isNear, &isIntersect, &maxDistance)) {
+                } else if (parseQuery(obj, &geoQueryField, &isNear, &isIntersect, &isWithin, &maxDistance)) {
                     regions.push_back(geoQueryField);
                 } else {
                     uasserted(16535, "can't parse query for *2d geo search: " + obj.toString());
                 }
             }
 
-            if (isNear && isIntersect ) {
-                uasserted(16474, "Can't do both near and intersect, query: " +  query.toString());
+            int queryTypes = 0;
+            if (isNear) ++queryTypes;
+            if (isIntersect) {
+                for (size_t i = 0; i < regions.size(); ++i) {
+                    regions[i].predicate = QueryGeometry::INTERSECT;
+                }
+                ++queryTypes;
+            }
+            if (isWithin) {
+                for (size_t i = 0; i < regions.size(); ++i) {
+                    regions[i].predicate = QueryGeometry::WITHIN;
+                    // Why do we only deal with polygons?
+
+                    // 1. Finding things within a point is silly and only valid
+                    // for points and degenerate lines/polys.
+
+                    // 2. Finding points within a line is easy but that's called intersect.
+                    // Finding lines within a line is kind of tricky given what S2 gives us.
+                    // Doing line-within-line is a valid yet unsupported feature, though I wonder if we want to preserve
+                    // orientation for lines or allow (a,b),(c,d) to be within (c,d),(a,b).
+                    // Polygons aren't in lines.
+
+                    // 3. Finding <anything> in a polygon is what we support.
+
+                    uassert(16672, "$within is only supported with polygonal geometry",
+                            NULL != regions[i].polygon.get());
+                }
+                ++queryTypes;
+            }
+            if (1 != queryTypes) {
+                uasserted(16474, "Malformed query, require just one type of predicate: " +  query.toString());
             }
 
             // I copied this from 2d.cpp.  Guard against perversion.
@@ -258,7 +293,6 @@ namespace mongo {
                                                         _params, numWanted, maxDistance);
                 return shared_ptr<Cursor>(cursor);
             } else {
-                // Default to intersect.
                 S2Cursor *cursor = new S2Cursor(keyPattern(), getDetails(), filteredQuery, regions, _params,
                                                 numWanted);
                 return shared_ptr<Cursor>(cursor);
@@ -280,6 +314,7 @@ namespace mongo {
                 // getGtLtOp is horribly misnamed and really means get the operation.
                 switch (e.embeddedObject().firstElement().getGtLtOp()) {
                     case BSONObj::opNEAR:
+                    case BSONObj::opWITHIN:
                     case BSONObj::opGEO_INTERSECTS:
                         return OPTIMAL;
                     default:
@@ -314,7 +349,9 @@ namespace mongo {
                 vector<string> cells;
                 S2Polyline line;
                 S2Cell point;
-                // We only support GeoJSON polygons.
+                // We only support GeoJSON polygons.  Why?:
+                // 1. we don't automagically do WGS84/flat -> WGS84, and
+                // 2. the old polygon format must die.
                 if (GeoParser::isGeoJSONPolygon(obj)) {
                     S2Polygon polygon;
                     GeoParser::parseGeoJSONPolygon(obj, &polygon);
diff --git a/src/mongo/db/queryutil.cpp b/src/mongo/db/queryutil.cpp
index 5d84ec1cb33..af2c035bd65 100644
--- a/src/mongo/db/queryutil.cpp
+++ b/src/mongo/db/queryutil.cpp
@@ -447,6 +447,9 @@ namespace mongo {
         }
         case BSONObj::opWITHIN:
             _special.add("2d", SpecialIndices::NO_INDEX_REQUIRED);
+            // TODO(hk): make this work w/o an index too.  will require
+            // changing matcher.cpp to parse geojson stuff etc.
+            _special.add("2dsphere", SpecialIndices::INDEX_REQUIRED);
             break;
         case BSONObj::opNEAR:
             _special.add("2d", SpecialIndices::INDEX_REQUIRED);