SERVER-30189 Reduce calls to allocator for large $in expressions

(cherry picked from commit 50a1afafc816097ed57804ff7033dffd85dbe160)
(cherry picked from commit ef408a20e6fd71ffc97f3602dd65ef5b03de6a45)

6 files changed, 48 insertions, 43 deletions

diff --git a/src/mongo/bson/bsonobj.h b/src/mongo/bson/bsonobj.h
index 56785df3a62..18e918190a4 100644
--- a/src/mongo/bson/bsonobj.h
+++ b/src/mongo/bson/bsonobj.h
@@ -30,6 +30,7 @@
 #pragma once
 
 #include <bitset>
+#include <boost/container/flat_set.hpp>
 #include <list>
 #include <set>
 #include <string>
@@ -51,6 +52,7 @@
 namespace mongo {
 
 typedef std::set<BSONElement, BSONElementCmpWithoutField> BSONElementSet;
+typedef boost::container::flat_set<BSONElement, BSONElementCmpWithoutField> BSONElementFlatSet;
 typedef std::multiset<BSONElement, BSONElementCmpWithoutField> BSONElementMSet;
 
 /**
diff --git a/src/mongo/db/matcher/expression_leaf.cpp b/src/mongo/db/matcher/expression_leaf.cpp
index 5b658640578..dce381e8a27 100644
--- a/src/mongo/db/matcher/expression_leaf.cpp
+++ b/src/mongo/db/matcher/expression_leaf.cpp
@@ -534,22 +534,24 @@ ArrayFilterEntries::~ArrayFilterEntries() {
     _regexes.clear();
 }
 
-Status ArrayFilterEntries::addEquality(const BSONElement& e) {
-    if (e.type() == RegEx)
-        return Status(ErrorCodes::BadValue, "ArrayFilterEntries equality cannot be a regex");
+Status ArrayFilterEntries::setEqualities(std::vector<BSONElement> equalities) {
+    for (auto&& equality : equalities) {
+        if (equality.type() == BSONType::RegEx) {
+            return Status(ErrorCodes::BadValue, "ArrayFilterEntries equality cannot be a regex");
+        }
 
-    if (e.type() == Undefined) {
-        return Status(ErrorCodes::BadValue, "ArrayFilterEntries equality cannot be undefined");
-    }
+        if (equality.type() == BSONType::Undefined) {
+            return Status(ErrorCodes::BadValue, "ArrayFilterEntries equality cannot be undefined");
+        }
 
-    if (e.type() == jstNULL) {
-        _hasNull = true;
+        if (equality.type() == BSONType::jstNULL) {
+            _hasNull = true;
+        } else if (equality.type() == BSONType::Array && equality.Obj().isEmpty()) {
+            _hasEmptyArray = true;
+        }
     }
 
-    if (e.type() == Array && e.Obj().isEmpty())
-        _hasEmptyArray = true;
-
-    _equalities.insert(e);
+    _equalities = BSONElementFlatSet(equalities.begin(), equalities.end());
     return Status::OK();
 }
 
@@ -582,8 +584,8 @@ void ArrayFilterEntries::copyTo(ArrayFilterEntries& toFillIn) const {
 
 void ArrayFilterEntries::debugString(StringBuilder& debug) const {
     debug << "[ ";
-    for (BSONElementSet::const_iterator it = _equalities.begin(); it != _equalities.end(); ++it) {
-        debug << it->toString(false) << " ";
+    for (auto&& equality : _equalities) {
+        debug << equality.toString(false) << " ";
     }
     for (size_t i = 0; i < _regexes.size(); ++i) {
         _regexes[i]->shortDebugString(debug);
@@ -593,8 +595,8 @@ void ArrayFilterEntries::debugString(StringBuilder& debug) const {
 }
 
 void ArrayFilterEntries::toBSON(BSONArrayBuilder* out) const {
-    for (BSONElementSet::const_iterator it = _equalities.begin(); it != _equalities.end(); ++it) {
-        out->append(*it);
+    for (auto&& equality : _equalities) {
+        out->append(equality);
     }
     for (size_t i = 0; i < _regexes.size(); ++i) {
         BSONObjBuilder regexBob;
diff --git a/src/mongo/db/matcher/expression_leaf.h b/src/mongo/db/matcher/expression_leaf.h
index 2b382768788..2a830244dfa 100644
--- a/src/mongo/db/matcher/expression_leaf.h
+++ b/src/mongo/db/matcher/expression_leaf.h
@@ -301,10 +301,10 @@ public:
     ArrayFilterEntries();
     ~ArrayFilterEntries();
 
-    Status addEquality(const BSONElement& e);
+    Status setEqualities(std::vector<BSONElement> equalities);
     Status addRegex(RegexMatchExpression* expr);
 
-    const BSONElementSet& equalities() const {
+    const BSONElementFlatSet& equalities() const {
         return _equalities;
     }
     bool contains(const BSONElement& elem) const {
@@ -342,7 +342,7 @@ public:
 private:
     bool _hasNull;  // if _equalities has a jstNULL element in it
     bool _hasEmptyArray;
-    BSONElementSet _equalities;
+    BSONElementFlatSet _equalities;
     std::vector<RegexMatchExpression*> _regexes;
 };
 
diff --git a/src/mongo/db/matcher/expression_leaf_test.cpp b/src/mongo/db/matcher/expression_leaf_test.cpp
index d4d01eb7325..79624a409b5 100644
--- a/src/mongo/db/matcher/expression_leaf_test.cpp
+++ b/src/mongo/db/matcher/expression_leaf_test.cpp
@@ -1458,7 +1458,8 @@ TEST(InMatchExpression, MatchesElementSingle) {
     BSONObj match = BSON("a" << 1);
     BSONObj notMatch = BSON("a" << 2);
     InMatchExpression in;
-    in.getArrayFilterEntries()->addEquality(operand.firstElement());
+    std::vector<BSONElement> equalities{operand.firstElement()};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
     ASSERT(in.matchesSingleElement(match["a"]));
     ASSERT(!in.matchesSingleElement(notMatch["a"]));
 }
@@ -1476,10 +1477,8 @@ TEST(InMatchExpression, MatchesEmpty) {
 TEST(InMatchExpression, MatchesElementMultiple) {
     BSONObj operand = BSON_ARRAY(1 << "r" << true << 1);
     InMatchExpression in;
-    in.getArrayFilterEntries()->addEquality(operand[0]);
-    in.getArrayFilterEntries()->addEquality(operand[1]);
-    in.getArrayFilterEntries()->addEquality(operand[2]);
-    in.getArrayFilterEntries()->addEquality(operand[3]);
+    std::vector<BSONElement> equalities{operand[0], operand[1], operand[2], operand[3]};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 
     BSONObj matchFirst = BSON("a" << 1);
     BSONObj matchSecond = BSON("a"
@@ -1497,7 +1496,8 @@ TEST(InMatchExpression, MatchesScalar) {
     BSONObj operand = BSON_ARRAY(5);
     InMatchExpression in;
     in.init("a");
-    in.getArrayFilterEntries()->addEquality(operand.firstElement());
+    std::vector<BSONElement> equalities{operand.firstElement()};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 
     ASSERT(in.matchesBSON(BSON("a" << 5.0), NULL));
     ASSERT(!in.matchesBSON(BSON("a" << 4), NULL));
@@ -1507,7 +1507,8 @@ TEST(InMatchExpression, MatchesArrayValue) {
     BSONObj operand = BSON_ARRAY(5);
     InMatchExpression in;
     in.init("a");
-    in.getArrayFilterEntries()->addEquality(operand.firstElement());
+    std::vector<BSONElement> equalities{operand.firstElement()};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 
     ASSERT(in.matchesBSON(BSON("a" << BSON_ARRAY(5.0 << 6)), NULL));
     ASSERT(!in.matchesBSON(BSON("a" << BSON_ARRAY(6 << 7)), NULL));
@@ -1519,7 +1520,8 @@ TEST(InMatchExpression, MatchesNull) {
 
     InMatchExpression in;
     in.init("a");
-    in.getArrayFilterEntries()->addEquality(operand.firstElement());
+    std::vector<BSONElement> equalities{operand.firstElement()};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 
     ASSERT(in.matchesBSON(BSONObj(), NULL));
     ASSERT(in.matchesBSON(BSON("a" << BSONNULL), NULL));
@@ -1533,15 +1535,16 @@ TEST(InMatchExpression, MatchesUndefined) {
 
     InMatchExpression in;
     in.init("a");
-    Status s = in.getArrayFilterEntries()->addEquality(operand.firstElement());
-    ASSERT_NOT_OK(s);
+    std::vector<BSONElement> equalities{operand.firstElement()};
+    ASSERT_NOT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 }
 
 TEST(InMatchExpression, MatchesMinKey) {
     BSONObj operand = BSON_ARRAY(MinKey);
     InMatchExpression in;
     in.init("a");
-    in.getArrayFilterEntries()->addEquality(operand.firstElement());
+    std::vector<BSONElement> equalities{operand.firstElement()};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 
     ASSERT(in.matchesBSON(BSON("a" << MinKey), NULL));
     ASSERT(!in.matchesBSON(BSON("a" << MaxKey), NULL));
@@ -1552,7 +1555,8 @@ TEST(InMatchExpression, MatchesMaxKey) {
     BSONObj operand = BSON_ARRAY(MaxKey);
     InMatchExpression in;
     in.init("a");
-    in.getArrayFilterEntries()->addEquality(operand.firstElement());
+    std::vector<BSONElement> equalities{operand.firstElement()};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 
     ASSERT(in.matchesBSON(BSON("a" << MaxKey), NULL));
     ASSERT(!in.matchesBSON(BSON("a" << MinKey), NULL));
@@ -1563,9 +1567,8 @@ TEST(InMatchExpression, MatchesFullArray) {
     BSONObj operand = BSON_ARRAY(BSON_ARRAY(1 << 2) << 4 << 5);
     InMatchExpression in;
     in.init("a");
-    in.getArrayFilterEntries()->addEquality(operand[0]);
-    in.getArrayFilterEntries()->addEquality(operand[1]);
-    in.getArrayFilterEntries()->addEquality(operand[2]);
+    std::vector<BSONElement> equalities{operand[0], operand[1], operand[2]};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 
     ASSERT(in.matchesBSON(BSON("a" << BSON_ARRAY(1 << 2)), NULL));
     ASSERT(!in.matchesBSON(BSON("a" << BSON_ARRAY(1 << 2 << 3)), NULL));
@@ -1577,8 +1580,8 @@ TEST(InMatchExpression, ElemMatchKey) {
     BSONObj operand = BSON_ARRAY(5 << 2);
     InMatchExpression in;
     in.init("a");
-    in.getArrayFilterEntries()->addEquality(operand[0]);
-    in.getArrayFilterEntries()->addEquality(operand[1]);
+    std::vector<BSONElement> equalities{operand[0], operand[1]};
+    ASSERT_OK(in.getArrayFilterEntries()->setEqualities(std::move(equalities)));
 
     MatchDetails details;
     details.requestElemMatchKey();
diff --git a/src/mongo/db/matcher/expression_parser.cpp b/src/mongo/db/matcher/expression_parser.cpp
index 732211ccb44..ae7a4cabd66 100644
--- a/src/mongo/db/matcher/expression_parser.cpp
+++ b/src/mongo/db/matcher/expression_parser.cpp
@@ -588,6 +588,7 @@ StatusWithMatchExpression MatchExpressionParser::_parseRegexDocument(const char*
 
 Status MatchExpressionParser::_parseArrayFilterEntries(ArrayFilterEntries* entries,
                                                        const BSONObj& theArray) {
+    std::vector<BSONElement> equalities;
     BSONObjIterator i(theArray);
     while (i.more()) {
         BSONElement e = i.next();
@@ -606,12 +607,10 @@ Status MatchExpressionParser::_parseArrayFilterEntries(ArrayFilterEntries* entri
             if (!s.isOK())
                 return s;
         } else {
-            Status s = entries->addEquality(e);
-            if (!s.isOK())
-                return s;
+            equalities.push_back(e);
         }
     }
-    return Status::OK();
+    return entries->setEqualities(std::move(equalities));
 }
 
 StatusWithMatchExpression MatchExpressionParser::_parseType(const char* name,
diff --git a/src/mongo/db/query/index_bounds_builder.cpp b/src/mongo/db/query/index_bounds_builder.cpp
index 364f9577adb..69778c14dd1 100644
--- a/src/mongo/db/query/index_bounds_builder.cpp
+++ b/src/mongo/db/query/index_bounds_builder.cpp
@@ -522,9 +522,8 @@ void IndexBoundsBuilder::translate(const MatchExpression* expr,
         // Create our various intervals.
 
         IndexBoundsBuilder::BoundsTightness tightness;
-        for (BSONElementSet::iterator it = afr.equalities().begin(); it != afr.equalities().end();
-             ++it) {
-            translateEquality(*it, isHashed, oilOut, &tightness);
+        for (auto&& equality : afr.equalities()) {
+            translateEquality(equality, isHashed, oilOut, &tightness);
             if (tightness != IndexBoundsBuilder::EXACT) {
                 *tightnessOut = tightness;
             }