1 files changed, 264 insertions, 13 deletions

diff --git a/src/mongo/db/query/get_runner.cpp b/src/mongo/db/query/get_runner.cpp
index f230f4c95dc..36d5447b9c1 100644
--- a/src/mongo/db/query/get_runner.cpp
+++ b/src/mongo/db/query/get_runner.cpp
@@ -126,12 +126,19 @@ namespace mongo {
         return Status::OK();
     }
 
+    namespace {
+        // The body is below in the "count hack" section but getRunner calls it.
+        bool turnIxscanIntoCount(QuerySolution* soln);
+    }  // namespace
+
     /**
      * For a given query, get a runner.  The runner could be a SingleSolutionRunner, a
      * CachedQueryRunner, or a MultiPlanRunner, depending on the cache/query solver/etc.
      */
-    Status getRunner(Collection* collection, CanonicalQuery* rawCanonicalQuery,
-                     Runner** out, size_t plannerOptions) {
+    Status getRunner(Collection* collection,
+                     CanonicalQuery* rawCanonicalQuery,
+                     Runner** out,
+                     size_t plannerOptions) {
 
         verify(rawCanonicalQuery);
         auto_ptr<CanonicalQuery> canonicalQuery(rawCanonicalQuery);
@@ -239,7 +246,22 @@ namespace mongo {
             QuerySolution *qs, *backupQs;
             Status status = QueryPlanner::planFromCache(*canonicalQuery, plannerParams, *cs,
                                                         &qs, &backupQs);
+
             if (status.isOK()) {
+                if (plannerParams.options & QueryPlannerParams::PRIVATE_IS_COUNT) {
+                    if (turnIxscanIntoCount(qs)) {
+                        WorkingSet* ws;
+                        PlanStage* root;
+                        verify(StageBuilder::build(*qs, &root, &ws));
+                        *out = new SingleSolutionRunner(collection,
+                                                        canonicalQuery.release(), qs, root, ws);
+                        if (NULL != backupQs) {
+                            delete backupQs;
+                        }
+                        return Status::OK();
+                    }
+                }
+
                 WorkingSet* ws;
                 PlanStage* root;
                 verify(StageBuilder::build(*qs, &root, &ws));
@@ -270,13 +292,8 @@ namespace mongo {
                           " planner returned error: " + status.reason());
         }
 
-        /*
-        for (size_t i = 0; i < solutions.size(); ++i) {
-            QLOG() << "solution " << i << " is " << solutions[i]->toString() << endl;
-        }
-        */
-
-        // We cannot figure out how to answer the query.  Should this ever happen?
+        // We cannot figure out how to answer the query.  Perhaps it requires an index
+        // we do not have?
         if (0 == solutions.size()) {
             return Status(ErrorCodes::BadValue,
                           str::stream()
@@ -285,6 +302,31 @@ namespace mongo {
                           << " No query solutions");
         }
 
+        // See if one of our solutions is a fast count hack in disguise.
+        if (plannerParams.options & QueryPlannerParams::PRIVATE_IS_COUNT) {
+            for (size_t i = 0; i < solutions.size(); ++i) {
+                if (turnIxscanIntoCount(solutions[i])) {
+                    // Great, we can use solutions[i].  Clean up the other QuerySolution(s).
+                    for (size_t j = 0; j < solutions.size(); ++j) {
+                        if (j != i) {
+                            delete solutions[j];
+                        }
+                    }
+
+                    // We're not going to cache anything that's fast count.
+                    WorkingSet* ws;
+                    PlanStage* root;
+                    verify(StageBuilder::build(*solutions[i], &root, &ws));
+                    *out = new SingleSolutionRunner(collection,
+                                                    canonicalQuery.release(),
+                                                    solutions[i],
+                                                    root,
+                                                    ws);
+                    return Status::OK();
+                }
+            }
+        }
+
         if (1 == solutions.size()) {
             // Only one possible plan.  Run it.  Build the stages from the solution.
             WorkingSet* ws;
@@ -315,6 +357,219 @@ namespace mongo {
     }
 
     //
+    // Count hack
+    //
+
+    namespace {
+
+        /**
+         * Returns 'true' if the bounds 'bounds' can be represented as an interval between two the two
+         * values 'startKey' and 'endKey'.  Inclusivity of each bound is set through the relevant
+         * fooKeyInclusive parameter.
+         *
+         * Returns 'false' otherwise.
+         *
+         * XXX: unit test this.
+         */
+        bool isSingleInterval(const IndexBounds& bounds,
+                              BSONObj* startKey,
+                              bool* startKeyInclusive,
+                              BSONObj* endKey,
+                              bool* endKeyInclusive) {
+            // We build our start/end keys as we go.
+            BSONObjBuilder startBob;
+            BSONObjBuilder endBob;
+
+            // The start and end keys are inclusive unless we have a non-point interval, in which case
+            // we take the inclusivity from there.
+            *startKeyInclusive = true;
+            *endKeyInclusive = true;
+
+            size_t fieldNo = 0;
+
+            // First, we skip over point intervals.
+            for (; fieldNo < bounds.fields.size(); ++fieldNo) {
+                const OrderedIntervalList& oil = bounds.fields[fieldNo];
+                // A point interval requires just one interval...
+                if (1 != oil.intervals.size()) {
+                    break;
+                }
+                if (!oil.intervals[0].isPoint()) {
+                    break;
+                }
+                // Since it's a point, start == end.
+                startBob.append(oil.intervals[0].start);
+                endBob.append(oil.intervals[0].end);
+            }
+
+            if (fieldNo >= bounds.fields.size()) {
+                // All our intervals are points.  We count for all values of one field.
+                *startKey = startBob.obj();
+                *endKey = endBob.obj();
+                return true;
+            }
+
+            // After point intervals we can have exactly one non-point interval.
+            const OrderedIntervalList& nonPoint = bounds.fields[fieldNo];
+            if (1 != nonPoint.intervals.size()) {
+                return false;
+            }
+
+            // Add the non-point interval to our builder and set the inclusivity from it.
+            startBob.append(nonPoint.intervals[0].start);
+            *startKeyInclusive = nonPoint.intervals[0].startInclusive;
+            endBob.append(nonPoint.intervals[0].end);
+            *endKeyInclusive = nonPoint.intervals[0].endInclusive;
+
+            ++fieldNo;
+
+            // Get some "all values" intervals for comparison's sake.
+            // TODO: make static?
+            Interval minMax = IndexBoundsBuilder::allValues();
+            Interval maxMin = minMax;
+            maxMin.reverse();
+
+            // And after the non-point interval we can have any number of "all values" intervals.
+            for (; fieldNo < bounds.fields.size(); ++fieldNo) {
+                const OrderedIntervalList& oil = bounds.fields[fieldNo];
+                // "All Values" is just one point.
+                if (1 != oil.intervals.size()) {
+                    break;
+                }
+
+                // Must be min->max or max->min.
+                if (oil.intervals[0].equals(minMax)) {
+                    // As an example for the logic below, consider the index {a:1, b:1} and a count for
+                    // {a: {$gt: 2}}.  Our start key isn't inclusive (as it's $gt: 2) and looks like
+                    // {"":2} so far.  If we move to the key greater than {"":2, "": MaxKey} we will get
+                    // the first value of 'a' that is greater than 2.
+                    if (!*startKeyInclusive) {
+                        startBob.appendMaxKey("");
+                    }
+                    else {
+                        // In this case, consider the index {a:1, b:1} and a count for {a:{$gte: 2}}.
+                        // We want to look at all values where a is 2, so our start key is {"":2,
+                        // "":MinKey}.
+                        startBob.appendMinKey("");
+                    }
+
+                    // Same deal as above.  Consider the index {a:1, b:1} and a count for {a: {$lt: 2}}.
+                    // Our end key isn't inclusive as ($lt: 2) and looks like {"":2} so far.  We can't
+                    // look at any values where a is 2 so we have to stop at {"":2, "": MinKey} as
+                    // that's the smallest key where a is still 2.
+                    if (!*endKeyInclusive) {
+                        endBob.appendMinKey("");
+                    }
+                    else {
+                        endBob.appendMaxKey("");
+                    }
+                }
+                else if (oil.intervals[0].equals(maxMin)) {
+                    // The reasoning here is the same as above but with the directions reversed.
+                    if (!*startKeyInclusive) {
+                        startBob.appendMinKey("");
+                    }
+                    else {
+                        startBob.appendMaxKey("");
+                    }
+                    if (!*endKeyInclusive) {
+                        endBob.appendMaxKey("");
+                    }
+                    else {
+                        endBob.appendMinKey("");
+                    }
+                }
+                else {
+                    // No dice.
+                    break;
+                }
+            }
+
+            if (fieldNo >= bounds.fields.size()) {
+                *startKey = startBob.obj();
+                *endKey = endBob.obj();
+                return true;
+            }
+            else {
+                return false;
+            }
+        }
+
+        /**
+         * Returns 'true' if the provided solution 'soln' can be rewritten to use
+         * a fast counting stage.  Mutates the tree in 'soln->root'.
+         *
+         * Otherwise, returns 'false'.
+         */
+        bool turnIxscanIntoCount(QuerySolution* soln) {
+            QuerySolutionNode* root = soln->root.get();
+
+            // Root should be a fetch w/o any filters.
+            if (STAGE_FETCH != root->getType()) {
+                return false;
+            }
+
+            if (NULL != root->filter.get()) {
+                return false;
+            }
+
+            // Child should be an ixscan.
+            if (STAGE_IXSCAN != root->children[0]->getType()) {
+                return false;
+            }
+
+            IndexScanNode* isn = static_cast<IndexScanNode*>(root->children[0]);
+
+            // No filters allowed and side-stepping isSimpleRange for now.  TODO: do we ever see
+            // isSimpleRange here?  because we could well use it.  I just don't think we ever do see it.
+            if (NULL != isn->filter.get() || isn->bounds.isSimpleRange) {
+                return false;
+            }
+
+            // Make sure the bounds are OK.
+            BSONObj startKey;
+            bool startKeyInclusive;
+            BSONObj endKey;
+            bool endKeyInclusive;
+
+            if (!isSingleInterval(isn->bounds, &startKey, &startKeyInclusive, &endKey, &endKeyInclusive)) {
+                return false;
+            }
+
+            // Make the count node that we replace the fetch + ixscan with.
+            CountNode* cn = new CountNode();
+            cn->indexKeyPattern = isn->indexKeyPattern;
+            cn->startKey = startKey;
+            cn->startKeyInclusive = startKeyInclusive;
+            cn->endKey = endKey;
+            cn->endKeyInclusive = endKeyInclusive;
+            // Takes ownership of 'cn' and deletes the old root.
+            soln->root.reset(cn);
+            return true;
+        }
+
+    }  // namespace
+
+    Status getRunnerCount(Collection* collection,
+                          const BSONObj& query,
+                          const BSONObj& hintObj,
+                          Runner** out) {
+        verify(collection);
+
+        CanonicalQuery* cq;
+        uassertStatusOK(CanonicalQuery::canonicalize(collection->ns().ns(),
+                                                     query,
+                                                     BSONObj(),
+                                                     BSONObj(), 
+                                                     0,
+                                                     0,
+                                                     hintObj,
+                                                     &cq)); 
+
+        return getRunner(collection, cq, out, QueryPlannerParams::PRIVATE_IS_COUNT);
+    }
+
+    //
     // Distinct hack
     //
 
@@ -378,10 +633,6 @@ namespace mongo {
                              const BSONObj& query,
                              const string& field,
                              Runner** out) {
-
-        Database* db = cc().database();
-        verify(db);
-
         // This should'a been checked by the distinct command.
         verify(collection);