SERVER-40253 Implement count command in IDL

3 files changed, 205 insertions, 0 deletions

diff --git a/src/mongo/bson/bsonelement.cpp b/src/mongo/bson/bsonelement.cpp
index 605623fb54f..018b1f2e6e6 100644
--- a/src/mongo/bson/bsonelement.cpp
+++ b/src/mongo/bson/bsonelement.cpp
@@ -567,6 +567,84 @@ bool BSONElement::binaryEqualValues(const BSONElement& rhs) const {
     return (valueSize == 0) || (memcmp(value(), rhs.value(), valueSize) == 0);
 }
 
+StatusWith<long long> BSONElement::parseIntegerElementToNonNegativeLong() const {
+    auto number = parseIntegerElementToLong();
+    if (!number.isOK()) {
+        return number;
+    }
+
+    if (number.getValue() < 0) {
+        return Status(ErrorCodes::FailedToParse,
+                      str::stream() << "Expected a positive number in: " << toString(true, true));
+    }
+
+    return number;
+}
+
+StatusWith<long long> BSONElement::parseIntegerElementToLong() const {
+    if (!isNumber()) {
+        return Status(ErrorCodes::FailedToParse,
+                      str::stream() << "Expected a number in: " << toString(true, true));
+    }
+
+    long long number = 0;
+    if (type() == BSONType::NumberDouble) {
+        auto eDouble = numberDouble();
+
+        // NaN doubles are rejected.
+        if (std::isnan(eDouble)) {
+            return Status(ErrorCodes::FailedToParse,
+                          str::stream() << "Expected an integer, but found NaN in: "
+                                        << toString(true, true));
+        }
+
+        // No integral doubles that are too large to be represented as a 64 bit signed integer.
+        // We use 'kLongLongMaxAsDouble' because if we just did eDouble > 2^63-1, it would be
+        // compared against 2^63. eDouble=2^63 would not get caught that way.
+        if (eDouble >= kLongLongMaxPlusOneAsDouble ||
+            eDouble < std::numeric_limits<long long>::min()) {
+            return Status(ErrorCodes::FailedToParse,
+                          str::stream() << "Cannot represent as a 64-bit integer: "
+                                        << toString(true, true));
+        }
+
+        // This checks if elem is an integral double.
+        if (eDouble != static_cast<double>(static_cast<long long>(eDouble))) {
+            return Status(ErrorCodes::FailedToParse,
+                          str::stream() << "Expected an integer: " << toString(true, true));
+        }
+
+        number = numberLong();
+    } else if (type() == BSONType::NumberDecimal) {
+        uint32_t signalingFlags = Decimal128::kNoFlag;
+        number = numberDecimal().toLongExact(&signalingFlags);
+        if (signalingFlags != Decimal128::kNoFlag) {
+            return Status(ErrorCodes::FailedToParse,
+                          str::stream() << "Cannot represent as a 64-bit integer: "
+                                        << toString(true, true));
+        }
+    } else {
+        number = numberLong();
+    }
+
+    return number;
+}
+
+StatusWith<int> BSONElement::parseIntegerElementToInt() const {
+    auto parsedLong = parseIntegerElementToLong();
+    if (!parsedLong.isOK()) {
+        return parsedLong.getStatus();
+    }
+
+    auto valueLong = parsedLong.getValue();
+    if (valueLong < std::numeric_limits<int>::min() ||
+        valueLong > std::numeric_limits<int>::max()) {
+        return {ErrorCodes::FailedToParse,
+                str::stream() << "Cannot represent " << toString(true, true) << " in an int"};
+    }
+    return static_cast<int>(valueLong);
+}
+
 BSONObj BSONElement::embeddedObjectUserCheck() const {
     if (MONGO_likely(isABSONObj()))
         return BSONObj(value(), BSONObj::LargeSizeTrait{});
diff --git a/src/mongo/bson/bsonelement.h b/src/mongo/bson/bsonelement.h
index 8ccbd99936c..c5a08e8a9cc 100644
--- a/src/mongo/bson/bsonelement.h
+++ b/src/mongo/bson/bsonelement.h
@@ -359,6 +359,36 @@ public:
      */
     long long safeNumberLongForHash() const;
 
+    /**
+     * Parses a BSONElement of any numeric type into a positive long long, failing if the value
+     * is any of the following:
+     *
+     * - NaN.
+     * - Negative.
+     * - A floating point number which is not integral.
+     * - Too large to fit within a 64-bit signed integer.
+     */
+    StatusWith<long long> parseIntegerElementToNonNegativeLong() const;
+
+    /**
+     * Parses a BSONElement of any numeric type into a long long, failing if the value
+     * is any of the following:
+     *
+     * - NaN.
+     * - A floating point number which is not integral.
+     * - Too large in the positive or negative direction to fit within a 64-bit signed integer.
+     */
+    StatusWith<long long> parseIntegerElementToLong() const;
+
+    /**
+     * Parses a BSONElement of any numeric type into an integer, failing if the value is:
+     *
+     * - NaN
+     * - a non-integral number
+     * - too large in the positive or negative direction to fit in an int
+     */
+    StatusWith<int> parseIntegerElementToInt() const;
+
     /** Retrieve decimal value for the element safely. */
     Decimal128 numberDecimal() const;
 
diff --git a/src/mongo/bson/bsonelement_test.cpp b/src/mongo/bson/bsonelement_test.cpp
index 8901ab4339a..5c036ebeb23 100644
--- a/src/mongo/bson/bsonelement_test.cpp
+++ b/src/mongo/bson/bsonelement_test.cpp
@@ -208,5 +208,102 @@ TEST(BSONElement, SafeNumberLongNegativeBound) {
     ASSERT_EQ(obj["negativeInfinity"].safeNumberLong(), std::numeric_limits<long long>::lowest());
 }
 
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToNonNegativeLongRejectsNegative) {
+    BSONObj query = BSON("" << -2LL);
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToNonNegativeLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongAcceptsNegative) {
+    BSONObj query = BSON("" << -2LL);
+    auto result = query.firstElement().parseIntegerElementToLong();
+    ASSERT_OK(result.getStatus());
+    ASSERT_EQ(-2LL, result.getValue());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToNonNegativeLongRejectsTooLargeDouble) {
+    BSONObj query = BSON("" << BSONElement::kLongLongMaxPlusOneAsDouble);
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToNonNegativeLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongRejectsTooLargeDouble) {
+    BSONObj query = BSON("" << BSONElement::kLongLongMaxPlusOneAsDouble);
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongRejectsTooLargeNegativeDouble) {
+    BSONObj query = BSON("" << std::numeric_limits<double>::min());
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToNonNegativeLongRejectsString) {
+    BSONObj query = BSON(""
+                         << "1");
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToNonNegativeLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongRejectsString) {
+    BSONObj query = BSON(""
+                         << "1");
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToNonNegativeLongRejectsNonIntegralDouble) {
+    BSONObj query = BSON("" << 2.5);
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToNonNegativeLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongRejectsNonIntegralDouble) {
+    BSONObj query = BSON("" << 2.5);
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToNonNegativeLongRejectsNonIntegralDecimal) {
+    BSONObj query = BSON("" << Decimal128("2.5"));
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToNonNegativeLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongRejectsNonIntegralDecimal) {
+    BSONObj query = BSON("" << Decimal128("2.5"));
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToNonNegativeLongRejectsLargestDecimal) {
+    BSONObj query = BSON("" << Decimal128(Decimal128::kLargestPositive));
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToNonNegativeLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongRejectsLargestDecimal) {
+    BSONObj query = BSON("" << Decimal128(Decimal128::kLargestPositive));
+    ASSERT_NOT_OK(query.firstElement().parseIntegerElementToLong());
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToNonNegativeLongAcceptsZero) {
+    BSONObj query = BSON("" << 0);
+    auto result = query.firstElement().parseIntegerElementToNonNegativeLong();
+    ASSERT_OK(result.getStatus());
+    ASSERT_EQ(result.getValue(), 0LL);
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongAcceptsZero) {
+    BSONObj query = BSON("" << 0);
+    auto result = query.firstElement().parseIntegerElementToLong();
+    ASSERT_OK(result.getStatus());
+    ASSERT_EQ(result.getValue(), 0LL);
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToNonNegativeLongAcceptsThree) {
+    BSONObj query = BSON("" << 3.0);
+    auto result = query.firstElement().parseIntegerElementToNonNegativeLong();
+    ASSERT_OK(result.getStatus());
+    ASSERT_EQ(result.getValue(), 3LL);
+}
+
+TEST(BSONElementIntegerParseTest, ParseIntegerElementToLongAcceptsThree) {
+    BSONObj query = BSON("" << 3.0);
+    auto result = query.firstElement().parseIntegerElementToLong();
+    ASSERT_OK(result.getStatus());
+    ASSERT_EQ(result.getValue(), 3LL);
+}
+
 }  // namespace
 }  // namespace mongo