SERVER-30985 Implement numChildren() in InternalSchemaObjectMatchExpression

20 files changed, 280 insertions, 25 deletions

diff --git a/src/mongo/db/matcher/expression.h b/src/mongo/db/matcher/expression.h
index 49c032232ef..a7e89a59c9e 100644
--- a/src/mongo/db/matcher/expression.h
+++ b/src/mongo/db/matcher/expression.h
@@ -133,30 +133,24 @@ public:
     }
 
     /**
-     * How many children does the node have?  Most nodes are leaves so the default impl. is for
-     * a leaf.
+     * Returns the number of child MatchExpression nodes contained by this node. It is expected that
+     * a node that does not support child nodes will return 0.
      */
-    virtual size_t numChildren() const {
-        return 0;
-    }
+    virtual size_t numChildren() const = 0;
 
     /**
-     * Get the i-th child.
+     * Returns the child of the current node at zero-based position 'index'. 'index' must be within
+     * the range of [0, numChildren()).
      */
-    virtual MatchExpression* getChild(size_t i) const {
-        return NULL;
-    }
+    virtual MatchExpression* getChild(size_t index) const = 0;
 
     /**
-     * Returns the underlying vector storing the children of a logical node. Note that this is not
-     * guaranteed to return all children. It can be used to modify the children of logical nodes
-     * like AND/OR, but it cannot be used to traverse the MatchExpression tree. Traversing the
-     * MatchExpression tree should instead be achieved using numChildren() and getChild(), which are
-     * guaranteed to be accurate.
+     * For MatchExpression nodes that can participate in tree restructuring (like AND/OR), returns a
+     * non-const vector of MatchExpression* child nodes.
+     * Do not use to traverse the MatchExpression tree. Use numChildren() and getChild(), which
+     * provide access to all nodes.
      */
-    virtual std::vector<MatchExpression*>* getChildVector() {
-        return NULL;
-    }
+    virtual std::vector<MatchExpression*>* getChildVector() = 0;
 
     /**
      * Get the path of the leaf.  Returns StringData() if there is no path (node is logical).
@@ -232,7 +226,7 @@ public:
         return _tagData.get();
     }
     virtual void resetTag() {
-        setTag(NULL);
+        setTag(nullptr);
         for (size_t i = 0; i < numChildren(); ++i) {
             getChild(i)->resetTag();
         }
diff --git a/src/mongo/db/matcher/expression_always_boolean.h b/src/mongo/db/matcher/expression_always_boolean.h
index d54e4ac424d..1625dcf7394 100644
--- a/src/mongo/db/matcher/expression_always_boolean.h
+++ b/src/mongo/db/matcher/expression_always_boolean.h
@@ -69,6 +69,18 @@ public:
         return MatchCategory::kOther;
     }
 
+    size_t numChildren() const override {
+        return 0;
+    }
+
+    MatchExpression* getChild(size_t i) const override {
+        MONGO_UNREACHABLE;
+    }
+
+    std::vector<MatchExpression*>* getChildVector() override {
+        return nullptr;
+    }
+
 private:
     bool _value;
 };
diff --git a/src/mongo/db/matcher/expression_arity.h b/src/mongo/db/matcher/expression_arity.h
index 07196a56900..dc43269a91a 100644
--- a/src/mongo/db/matcher/expression_arity.h
+++ b/src/mongo/db/matcher/expression_arity.h
@@ -79,6 +79,19 @@ public:
             [](const auto& expr1, const auto& expr2) { return expr1->equivalent(expr2.get()); });
     }
 
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
+    size_t numChildren() const final {
+        return nargs;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        invariant(i < nargs);
+        return _expressions[i].get();
+    }
+
     /**
      * Takes ownership of the MatchExpressions in 'expressions'.
      */
diff --git a/src/mongo/db/matcher/expression_array.h b/src/mongo/db/matcher/expression_array.h
index 0eb32b024d2..1ea6df0a10a 100644
--- a/src/mongo/db/matcher/expression_array.h
+++ b/src/mongo/db/matcher/expression_array.h
@@ -90,6 +90,10 @@ public:
 
     virtual void serialize(BSONObjBuilder* out) const;
 
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
     virtual size_t numChildren() const {
         return 1;
     }
@@ -170,6 +174,18 @@ public:
         return std::move(e);
     }
 
+    size_t numChildren() const override {
+        return 0;
+    }
+
+    MatchExpression* getChild(size_t i) const override {
+        return nullptr;
+    }
+
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
     virtual bool matchesArray(const BSONObj& anArray, MatchDetails* details) const;
 
     virtual void debugString(StringBuilder& debug, int level) const;
diff --git a/src/mongo/db/matcher/expression_expr.h b/src/mongo/db/matcher/expression_expr.h
index 1177b7ef360..7818d75a339 100644
--- a/src/mongo/db/matcher/expression_expr.h
+++ b/src/mongo/db/matcher/expression_expr.h
@@ -70,6 +70,18 @@ public:
         return MatchCategory::kOther;
     }
 
+    size_t numChildren() const final {
+        return 0;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        MONGO_UNREACHABLE;
+    }
+
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
 private:
     void _doAddDependencies(DepsTracker* deps) const final {
         if (_expression) {
diff --git a/src/mongo/db/matcher/expression_leaf.h b/src/mongo/db/matcher/expression_leaf.h
index eb573debe5c..86106f2b653 100644
--- a/src/mongo/db/matcher/expression_leaf.h
+++ b/src/mongo/db/matcher/expression_leaf.h
@@ -51,6 +51,18 @@ public:
 
     virtual ~LeafMatchExpression() {}
 
+    size_t numChildren() const override {
+        return 0;
+    }
+
+    MatchExpression* getChild(size_t i) const override {
+        MONGO_UNREACHABLE;
+    }
+
+    std::vector<MatchExpression*>* getChildVector() override {
+        return nullptr;
+    }
+
     bool shouldExpandLeafArray() const override {
         return true;
     }
diff --git a/src/mongo/db/matcher/expression_tree.h b/src/mongo/db/matcher/expression_tree.h
index 3c2925ccc1e..5bfd86820ee 100644
--- a/src/mongo/db/matcher/expression_tree.h
+++ b/src/mongo/db/matcher/expression_tree.h
@@ -177,7 +177,7 @@ public:
     virtual void serialize(BSONObjBuilder* out) const;
 };
 
-class NotMatchExpression : public MatchExpression {
+class NotMatchExpression final : public MatchExpression {
 public:
     NotMatchExpression() : MatchExpression(NOT) {}
     NotMatchExpression(MatchExpression* e) : MatchExpression(NOT), _exp(e) {}
@@ -212,14 +212,18 @@ public:
 
     bool equivalent(const MatchExpression* other) const;
 
-    virtual size_t numChildren() const {
+    size_t numChildren() const final {
         return 1;
     }
 
-    virtual MatchExpression* getChild(size_t i) const {
+    MatchExpression* getChild(size_t i) const final {
         return _exp.get();
     }
 
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
     MatchExpression* releaseChild(void) {
         return _exp.release();
     }
diff --git a/src/mongo/db/matcher/expression_where_base.h b/src/mongo/db/matcher/expression_where_base.h
index 10204ec07cb..ec1d2d97915 100644
--- a/src/mongo/db/matcher/expression_where_base.h
+++ b/src/mongo/db/matcher/expression_where_base.h
@@ -44,9 +44,17 @@ public:
 
     WhereMatchExpressionBase(WhereParams params);
 
-    //
-    // Methods inherited from MatchExpression.
-    //
+    size_t numChildren() const final {
+        return 0;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        MONGO_UNREACHABLE;
+    }
+
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
 
     bool matchesSingleElement(const BSONElement& e, MatchDetails* details = nullptr) const final {
         return false;
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h b/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h
index 3f7e79fda8f..dd17ea89fca 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h
@@ -72,6 +72,19 @@ public:
 
     bool equivalent(const MatchExpression* other) const final;
 
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
+    size_t numChildren() const final {
+        return 1;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        invariant(i == 0);
+        return _query.get();
+    }
+
 private:
     long long _index;
     std::unique_ptr<MatchExpression> _query;
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index_test.cpp b/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index_test.cpp
index 87c0896193a..f21d80676f8 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index_test.cpp
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index_test.cpp
@@ -32,6 +32,7 @@
 #include "mongo/db/json.h"
 #include "mongo/db/matcher/expression_parser.h"
 #include "mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h"
+#include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
@@ -92,5 +93,24 @@ TEST(InternalSchemaAllElemMatchFromIndexMatchExpression, MatchedQueriesWithDotte
         expr.getValue()->matchesBSON(BSON("a" << BSON("b" << BSON_ARRAY(1 << 2 << 3 << 4)))));
 }
 
+TEST(InternalSchemaAllElemMatchFromIndexMatchExpression, HasSingleChild) {
+    auto query = fromjson("{'a.b': {$_internalSchemaAllElemMatchFromIndex: [2, {a: {$lt: 5}}]}}");
+    auto objMatch = MatchExpressionParser::parse(query, nullptr);
+    ASSERT_OK(objMatch.getStatus());
+
+    ASSERT_EQ(objMatch.getValue()->numChildren(), 1U);
+    ASSERT(objMatch.getValue()->getChild(0));
+}
+
+DEATH_TEST(InternalSchemaAllElemMatchFromIndexMatchExpression,
+           GetChildFailsIndexGreaterThanOne,
+           "Invariant failure i == 0") {
+    auto query = fromjson("{'a.b': {$_internalSchemaAllElemMatchFromIndex: [2, {a: {$lt: 5}}]}}");
+    auto objMatch = MatchExpressionParser::parse(query, nullptr);
+    ASSERT_OK(objMatch.getStatus());
+
+    objMatch.getValue()->getChild(1);
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h b/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h
index 48503c46b0e..76945fd4983 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h
@@ -140,6 +140,24 @@ public:
 
     std::unique_ptr<MatchExpression> shallowClone() const final;
 
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
+    size_t numChildren() const final {
+        return _patternProperties.size() + 1;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        invariant(i < numChildren());
+
+        if (i == 0) {
+            return _otherwise->getFilter();
+        }
+
+        return _patternProperties[i - 1].second->getFilter();
+    }
+
 private:
     /**
      * Helper function for matches() and matchesSingleElement().
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties_test.cpp b/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties_test.cpp
index 5572831cb8e..01dac629ea1 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties_test.cpp
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties_test.cpp
@@ -31,6 +31,7 @@
 #include "mongo/bson/json.h"
 #include "mongo/db/matcher/expression_parser.h"
 #include "mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h"
+#include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
@@ -115,4 +116,28 @@ TEST(InternalSchemaAllowedPropertiesMatchExpression, EquivalentToClone) {
     clone = expr.getValue()->shallowClone();
     ASSERT_TRUE(expr.getValue()->equivalent(clone.get()));
 }
+
+TEST(InternalSchemaAllowedPropertiesMatchExpression, HasCorrectNumberOfChilden) {
+    auto query = fromjson(
+        "{$_internalSchemaAllowedProperties: {properties: ['a'], namePlaceholder: 'i',"
+        "patternProperties: [{regex: /a/, expression: {i: 1}}], otherwise: {i: 7}}}");
+    auto objMatch = MatchExpressionParser::parse(query, nullptr);
+    ASSERT_OK(objMatch.getStatus());
+
+    ASSERT_EQ(objMatch.getValue()->numChildren(), 2U);
+    ASSERT(objMatch.getValue()->getChild(0));
+}
+
+DEATH_TEST(InternalSchemaAllowedPropertiesMatchExpression,
+           GetChildFailsOnIndexLargerThanChildSet,
+           "i < numChildren()") {
+    auto query = fromjson(
+        "{$_internalSchemaAllowedProperties: {properties: ['a'], namePlaceholder: 'i',"
+        "patternProperties: [{regex: /a/, expression: {i: 1}}], otherwise: {i: 7}}}");
+    auto objMatch = MatchExpressionParser::parse(query, nullptr);
+    ASSERT_OK(objMatch.getStatus());
+
+    const size_t numChildren = 2;
+    objMatch.getValue()->getChild(numChildren);
+}
 }  // namespace mongo
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index.h b/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index.h
index c3536e7deaa..acf83d71d50 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index.h
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index.h
@@ -75,6 +75,19 @@ public:
 
     std::unique_ptr<MatchExpression> shallowClone() const final;
 
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
+    size_t numChildren() const final {
+        return 1;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        invariant(i == 0);
+        return _expression->getFilter();
+    }
+
 private:
     long long _index = 0;
     std::unique_ptr<ExpressionWithPlaceholder> _expression;
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index_test.cpp b/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index_test.cpp
index 87c1d847beb..fcceda2f0d3 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index_test.cpp
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index_test.cpp
@@ -33,6 +33,7 @@
 #include "mongo/db/matcher/expression_with_placeholder.h"
 #include "mongo/db/matcher/schema/expression_internal_schema_match_array_index.h"
 #include "mongo/db/query/collation/collator_interface_mock.h"
+#include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
@@ -108,5 +109,28 @@ TEST(InternalSchemaMatchArrayIndexMatchExpression, EquivalentToClone) {
     auto clone = expr.getValue()->shallowClone();
     ASSERT_TRUE(expr.getValue()->equivalent(clone.get()));
 }
+
+TEST(InternalSchemaMatchArrayIndexMatchExpression, HasSingleChild) {
+    auto query = fromjson(
+        "{foo: {$_internalSchemaMatchArrayIndex:"
+        "{index: 0, namePlaceholder: 'i', expression: {i: {$type: 'number'}}}}}");
+    auto objMatch = MatchExpressionParser::parse(query, nullptr);
+    ASSERT_OK(objMatch.getStatus());
+
+    ASSERT_EQ(objMatch.getValue()->numChildren(), 1U);
+    ASSERT(objMatch.getValue()->getChild(0));
+}
+
+DEATH_TEST(InternalSchemaMatchArrayIndexMatchExpression,
+           GetChildFailsIndexGreaterThanZero,
+           "Invariant failure i == 0") {
+    auto query = fromjson(
+        "{foo: {$_internalSchemaMatchArrayIndex:"
+        "{index: 0, namePlaceholder: 'i', expression: {i: {$type: 'number'}}}}}");
+    auto objMatch = MatchExpressionParser::parse(query, nullptr);
+    ASSERT_OK(objMatch.getStatus());
+
+    objMatch.getValue()->getChild(1);
+}
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_num_array_items.h b/src/mongo/db/matcher/schema/expression_internal_schema_num_array_items.h
index 0acedc4ee67..86e1281c0fd 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_num_array_items.h
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_num_array_items.h
@@ -54,6 +54,18 @@ public:
 
     bool equivalent(const MatchExpression* other) const final;
 
+    size_t numChildren() const final {
+        return 0;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        MONGO_UNREACHABLE;
+    }
+
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
 protected:
     long long numItems() const {
         return _numItems;
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_num_properties.h b/src/mongo/db/matcher/schema/expression_internal_schema_num_properties.h
index c31b0aaf5ac..67f3a6e880b 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_num_properties.h
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_num_properties.h
@@ -49,6 +49,18 @@ public:
         return Status::OK();
     }
 
+    size_t numChildren() const final {
+        return 0;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        MONGO_UNREACHABLE;
+    }
+
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
     void debugString(StringBuilder& debug, int level) const final;
 
     void serialize(BSONObjBuilder* out) const final;
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_object_match.h b/src/mongo/db/matcher/schema/expression_internal_schema_object_match.h
index 4257b5b05a5..53f387752b7 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_object_match.h
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_object_match.h
@@ -53,6 +53,15 @@ public:
 
     bool equivalent(const MatchExpression* other) const final;
 
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
+    size_t numChildren() const final {
+        invariant(_sub);
+        return 1;
+    }
+
     MatchExpression* getChild(size_t i) const final {
         // 'i' must be 0 since there's always exactly one child.
         invariant(i == 0);
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_object_match_test.cpp b/src/mongo/db/matcher/schema/expression_internal_schema_object_match_test.cpp
index d195b901af6..22137895c3d 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_object_match_test.cpp
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_object_match_test.cpp
@@ -32,6 +32,7 @@
 #include "mongo/db/matcher/matcher.h"
 #include "mongo/db/matcher/schema/expression_internal_schema_object_match.h"
 #include "mongo/db/query/collation/collator_interface_mock.h"
+#include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
@@ -195,5 +196,30 @@ TEST(InternalSchemaObjectMatchExpression, RejectsArraysContainingMatchingSubObje
     ASSERT_FALSE(objMatch.getValue()->matchesBSON(fromjson("{a: [{b: 1}, {b: 2}]}")));
 }
 
+TEST(InternalSchemaObjectMatchExpression, HasSingleChild) {
+    auto query = fromjson(
+        "    {a: {$_internalSchemaObjectMatch: {"
+        "        c: {$eq: 3}"
+        "    }}}");
+    auto objMatch = MatchExpressionParser::parse(query, nullptr);
+    ASSERT_OK(objMatch.getStatus());
+
+    ASSERT_EQ(objMatch.getValue()->numChildren(), 1U);
+    ASSERT(objMatch.getValue()->getChild(0));
+}
+
+DEATH_TEST(InternalSchemaObjectMatchExpression,
+           GetChildFailsIndexGreaterThanZero,
+           "Invariant failure i == 0") {
+    auto query = fromjson(
+        "    {a: {$_internalSchemaObjectMatch: {"
+        "        c: {$eq: 3}"
+        "    }}}");
+    auto objMatch = MatchExpressionParser::parse(query, nullptr);
+    ASSERT_OK(objMatch.getStatus());
+
+    objMatch.getValue()->getChild(1);
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_unique_items.h b/src/mongo/db/matcher/schema/expression_internal_schema_unique_items.h
index c1bfc208a91..570cf448f07 100644
--- a/src/mongo/db/matcher/schema/expression_internal_schema_unique_items.h
+++ b/src/mongo/db/matcher/schema/expression_internal_schema_unique_items.h
@@ -49,6 +49,18 @@ public:
         return setPath(path);
     }
 
+    size_t numChildren() const final {
+        return 0;
+    }
+
+    MatchExpression* getChild(size_t i) const final {
+        MONGO_UNREACHABLE;
+    }
+
+    std::vector<MatchExpression*>* getChildVector() final {
+        return nullptr;
+    }
+
     bool matchesArray(const BSONObj& array, MatchDetails*) const final {
         auto set = _comparator.makeBSONEltSet();
         for (auto&& elem : array) {
diff --git a/src/mongo/db/pipeline/document_source_match_test.cpp b/src/mongo/db/pipeline/document_source_match_test.cpp
index 77f31252caa..c1e004d9e26 100644
--- a/src/mongo/db/pipeline/document_source_match_test.cpp
+++ b/src/mongo/db/pipeline/document_source_match_test.cpp
@@ -302,7 +302,7 @@ TEST_F(DocumentSourceMatchTest,
     auto match = DocumentSourceMatch::create(query, getExpCtx());
     DepsTracker dependencies;
     ASSERT_EQUALS(DocumentSource::SEE_NEXT, match->getDependencies(&dependencies));
-    ASSERT_EQUALS(0U, dependencies.fields.size());
+    ASSERT_EQUALS(1U, dependencies.fields.size());
     ASSERT_EQUALS(true, dependencies.needWholeDocument);
     ASSERT_EQUALS(false, dependencies.getNeedTextScore());
 }