Adds restrictions on overlapping compound fields and removes the unnecessary invariant

3 files changed, 84 insertions, 106 deletions

diff --git a/src/mongo/db/index/wildcard_key_generator.cpp b/src/mongo/db/index/wildcard_key_generator.cpp
index 58ff1edb405..1a114cd4b06 100644
--- a/src/mongo/db/index/wildcard_key_generator.cpp
+++ b/src/mongo/db/index/wildcard_key_generator.cpp
@@ -34,6 +34,7 @@
 #include "mongo/db/exec/projection_executor.h"
 #include "mongo/db/exec/projection_executor_builder.h"
 #include "mongo/db/jsobj.h"
+#include "mongo/db/matcher/expression_algo.h"
 #include "mongo/db/query/collation/collation_index_key.h"
 #include "mongo/db/query/projection_parser.h"
 
@@ -61,6 +62,41 @@ void popPathComponent(BSONElement elem, bool enclosingObjIsArray, FieldRef* path
         pathToElem->removeLastPart();
     }
 }
+
+void validateWildcardIndexKeys(const std::vector<StringData>& nonWildcardFields,
+                               const StringData& wildcardField,
+                               const std::vector<StringData>& projectionInclusionSet,
+                               const std::set<StringData>& projectionExclusionSet) {
+    for (const auto& nonWildcardField : nonWildcardFields) {
+        if (wildcardField == "$**") {
+            if (projectionInclusionSet.size()) {
+                for (const auto& projInc : projectionInclusionSet) {
+                    uassert(
+                        9999901,
+                        "Compound wildcard index fields overlaps with the projected wildcard field",
+                        !(projInc == nonWildcardField ||
+                          expression::isPathPrefixOf(projInc, nonWildcardField) ||
+                          expression::isPathPrefixOf(nonWildcardField, projInc)));
+                }
+            } else {
+                uassert(9999902,
+                        "Wildcard index does not allow compound key fields with the toplevel "
+                        "wildcard field without exclusion projection on them",
+                        std::find_if(projectionExclusionSet.begin(),
+                                     projectionExclusionSet.end(),
+                                     [&](const auto& exclude) {
+                                         return exclude == nonWildcardField ||
+                                             expression::isPathPrefixOf(exclude, nonWildcardField);
+                                     }) != projectionExclusionSet.end());
+            }
+        }
+        uassert(9999903,
+                "Compound wildcard index does not allow fields overlapping with the wildcard field",
+                !(wildcardField == nonWildcardField ||
+                  expression::isPathPrefixOf(wildcardField, nonWildcardField) ||
+                  expression::isPathPrefixOf(nonWildcardField, wildcardField)));
+    }
+}
 }  // namespace
 
 constexpr StringData WildcardKeyGenerator::kSubtreeSuffix;
@@ -68,14 +104,18 @@ constexpr StringData WildcardKeyGenerator::kSubtreeSuffix;
 WildcardProjection WildcardKeyGenerator::createProjectionExecutor(BSONObj keyPattern,
                                                                   BSONObj pathProjection) {
     WildcardProjection* proj = nullptr;
+    std::vector<StringData> nonWildcardFields;
+    StringData wildcardField;
     for (const auto& elem : keyPattern) {
         if (auto keyStr = elem.fieldNameStringData();
             (keyStr == "$**") || keyStr.endsWith(".$**")) {
-            uassert(99999, "Wildcard index does not allow multiple wildcard compound keys", !proj);
+            uassert(
+                9999900, "Wildcard index does not allow multiple wildcard compound keys", !proj);
             // The _keyPattern is either { "$**": 1 } for all paths or { "path.$**": 1 } for a
             // single subtree. If we are indexing a single subtree, then we will project just that
             // path.
             auto suffixPos = keyStr.find(kSubtreeSuffix);
+            wildcardField = keyStr.substr(0, suffixPos);
 
             // If we're indexing a single subtree, we can't also specify a path projection.
             invariant(suffixPos == std::string::npos || pathProjection.isEmpty());
@@ -85,7 +125,7 @@ WildcardProjection WildcardKeyGenerator::createProjectionExecutor(BSONObj keyPat
             // projection is empty we default to {_id: 0}, since empty projections are illegal and
             // will be rejected when parsed.
             auto projSpec = (suffixPos != std::string::npos
-                                 ? BSON(keyStr.substr(0, suffixPos) << 1)
+                                 ? BSON(wildcardField << 1)
                                  : pathProjection.isEmpty() ? kDefaultProjection : pathProjection);
 
             // Construct a dummy ExpressionContext for ProjectionExecutor. It's OK to set the
@@ -97,8 +137,23 @@ WildcardProjection WildcardKeyGenerator::createProjectionExecutor(BSONObj keyPat
             auto projection = projection_ast::parse(expCtx, projSpec, policies);
             proj = new WildcardProjection{projection_executor::buildProjectionExecutor(
                 expCtx, &projection, policies, projection_executor::kDefaultBuilderParams)};
+        } else {
+            nonWildcardFields.push_back(elem.fieldNameStringData());
+        }
+    }
+    // The projections on wildcard are only used when wildcardField is "$**" and are mutually
+    // exclusive.
+    std::vector<StringData> projectionInclusionSet;
+    std::set<StringData> projectionExclusionSet;
+    for (const auto& proj : pathProjection) {
+        if (proj.numberInt() == 1) {
+            projectionInclusionSet.push_back(proj.fieldNameStringData());
+        } else {
+            projectionExclusionSet.insert(proj.fieldNameStringData());
         }
     }
+    validateWildcardIndexKeys(
+        nonWildcardFields, wildcardField, projectionInclusionSet, projectionExclusionSet);
     return std::move(*proj);
 }
 
diff --git a/src/mongo/db/index/wildcard_key_generator_test.cpp b/src/mongo/db/index/wildcard_key_generator_test.cpp
index c1771aaf770..aab887e125b 100644
--- a/src/mongo/db/index/wildcard_key_generator_test.cpp
+++ b/src/mongo/db/index/wildcard_key_generator_test.cpp
@@ -1175,15 +1175,15 @@ TEST_F(WildcardKeyGeneratorDottedFieldsTest, DoNotIndexDottedFieldsWithSimilarSu
 
 struct WildcardKeyGeneratorCompoundTest : public WildcardKeyGeneratorTest {};
 
-TEST_F(WildcardKeyGeneratorCompoundTest, ExtractTopLevelKeyCompound) {
+TEST_F(WildcardKeyGeneratorCompoundTest, TopLevelKeyInclusionCompound) {
     WildcardKeyGenerator keyGen{fromjson("{a: 1, '$**': 1}"),
-                                {},
+                                fromjson("{b: 1}"),
                                 nullptr,
                                 KeyString::Version::kLatestVersion,
                                 Ordering::make(BSONObj())};
-    auto inputDoc = fromjson("{a: 1}");
+    auto inputDoc = fromjson("{a: 1, b: 1}");
 
-    auto expectedKeys = makeKeySet({fromjson("{'': 1, '': 'a', '': 1}")});
+    auto expectedKeys = makeKeySet({fromjson("{'': 1, '': 'b', '': 1}")});
     auto expectedMultikeyPaths = makeKeySet();
 
     auto outputKeys = makeKeySet();
@@ -1194,28 +1194,7 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ExtractTopLevelKeyCompound) {
     ASSERT(assertKeysetsEqual(expectedMultikeyPaths, multikeyMetadataKeys));
 }
 
-TEST_F(WildcardKeyGeneratorCompoundTest, ExtractKeysFromNestedObjectCompound) {
-    WildcardKeyGenerator keyGen{fromjson("{a: 1, '$**': 1}"),
-                                {},
-                                nullptr,
-                                KeyString::Version::kLatestVersion,
-                                Ordering::make(BSONObj())};
-    auto inputDoc = fromjson("{a: {b: 'one', c: 2}}");
-
-    auto expectedKeys = makeKeySet({fromjson("{'': {b: 'one', c: 2}, '': 'a.b', '': 'one'}"),
-                                    fromjson("{'': {b: 'one', c: 2}, '': 'a.c', '': 2}")});
-
-    auto expectedMultikeyPaths = makeKeySet();
-
-    KeyStringSet outputKeys;
-    KeyStringSet multikeyMetadataKeys;
-    keyGen.generateKeys(allocator, inputDoc, &outputKeys, &multikeyMetadataKeys);
-
-    ASSERT(assertKeysetsEqual(expectedKeys, outputKeys));
-    ASSERT(assertKeysetsEqual(expectedMultikeyPaths, multikeyMetadataKeys));
-}
-
-TEST_F(WildcardKeyGeneratorCompoundTest, ExclusionProjectionWithCompound) {
+TEST_F(WildcardKeyGeneratorCompoundTest, TopLevelKeyExclusionCompound) {
     WildcardKeyGenerator keyGen{fromjson("{a: 1, '$**': 1}"),
                                 fromjson("{a: 0}"),
                                 nullptr,
@@ -1235,41 +1214,19 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ExclusionProjectionWithCompound) {
     ASSERT(assertKeysetsEqual(expectedMultikeyPaths, multikeyMetadataKeys));
 }
 
-TEST_F(WildcardKeyGeneratorCompoundTest, InclusionProjectionWithCompound) {
-    WildcardKeyGenerator keyGen{fromjson("{'a.b': 1, '$**': 1}"),
-                                fromjson("{'a.c': 1}"),
-                                nullptr,
-                                KeyString::Version::kLatestVersion,
-                                Ordering::make(BSONObj())};
-    auto inputDoc = fromjson("{a: {b: 'one', c: 2}}");
-
-    auto expectedKeys = makeKeySet({fromjson("{'': 'one', '': 'a.c', '': 2}")});
-
-    auto expectedMultikeyPaths = makeKeySet();
-
-    KeyStringSet outputKeys;
-    KeyStringSet multikeyMetadataKeys;
-    keyGen.generateKeys(allocator, inputDoc, &outputKeys, &multikeyMetadataKeys);
-
-    ASSERT(assertKeysetsEqual(expectedKeys, outputKeys));
-    ASSERT(assertKeysetsEqual(expectedMultikeyPaths, multikeyMetadataKeys));
-}
-
-TEST_F(WildcardKeyGeneratorCompoundTest, ShouldIndexNonNestedEmptyArrayAsUndefinedCompound) {
+TEST_F(WildcardKeyGeneratorCompoundTest, TopLevelKeyMultiInclusionCompound) {
     WildcardKeyGenerator keyGen{fromjson("{a: 1, '$**': 1}"),
-                                {},
+                                fromjson("{b: 1, d: 1}"),
                                 nullptr,
                                 KeyString::Version::kLatestVersion,
                                 Ordering::make(BSONObj())};
     auto inputDoc = fromjson("{ a: [], b: {c: []}, d: [[], {e: []}]}");
 
-    auto expectedKeys = makeKeySet({fromjson("{'': undefined, '': 'a', '': undefined}"),
-                                    fromjson("{'': undefined, '': 'b.c', '': undefined}"),
+    auto expectedKeys = makeKeySet({fromjson("{'': undefined, '': 'b.c', '': undefined}"),
                                     fromjson("{'': undefined, '': 'd', '': []}"),
                                     fromjson("{'': undefined, '': 'd.e', '': undefined}")});
     auto expectedMultikeyPaths =
-        makeKeySet({fromjson("{'': {$minKey: 1}, '': 1, '': 'a'}"),
-                    fromjson("{'': {$minKey: 1}, '': 1, '': 'b.c'}"),
+        makeKeySet({fromjson("{'': {$minKey: 1}, '': 1, '': 'b.c'}"),
                     fromjson("{'': {$minKey: 1}, '': 1, '': 'd'}"),
                     fromjson("{'': {$minKey: 1}, '': 1, '': 'd.e'}")},
                    RecordIdReservations::reservedIdFor(ReservationId::kWildcardMultikeyMetadataId));
@@ -1284,7 +1241,7 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ShouldIndexNonNestedEmptyArrayAsUndefin
 
 TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathWithSingleField) {
     WildcardKeyGenerator keyGen{fromjson("{'$**': 1, f: 1}"),
-                                {},
+                                fromjson("{f: 0}"),
                                 nullptr,
                                 KeyString::Version::kLatestVersion,
                                 Ordering::make(BSONObj())};
@@ -1296,8 +1253,7 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathWithSingleField) {
                                     fromjson("{'': 'a.c.d', '': 1, '': 1}"),
                                     fromjson("{'': 'a.c.d', '': 2, '': 1}"),
                                     fromjson("{'': 'e', '': 3, '': 1}"),
-                                    fromjson("{'': 'e', '': 5, '': 1}"),
-                                    fromjson("{'': 'f', '': 1, '': 1}")});
+                                    fromjson("{'': 'e', '': 5, '': 1}")});
 
     auto expectedMultikeyPaths =
         makeKeySet({fromjson("{'': 1, '': 'a', '': {$minKey: 1}}"),
@@ -1315,7 +1271,7 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathWithSingleField) {
 
 TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathWithArrayField) {
     WildcardKeyGenerator keyGen{fromjson("{e: 1, '$**': 1}"),
-                                {},
+                                fromjson("{e: 0}"),
                                 nullptr,
                                 KeyString::Version::kLatestVersion,
                                 Ordering::make(BSONObj())};
@@ -1326,22 +1282,17 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathWithArrayField) {
                                     fromjson("{'': 3, '': 'a.c', '': 1}"),
                                     fromjson("{'': 3, '': 'a.c.d', '': 1}"),
                                     fromjson("{'': 3, '': 'a.c.d', '': 2}"),
-                                    fromjson("{'': 3, '': 'e', '': 3}"),
-                                    fromjson("{'': 3, '': 'e', '': 5}"),
                                     fromjson("{'': 3, '': 'f', '': 1}"),
                                     fromjson("{'': 5, '': 'a.b', '': 1}"),
                                     fromjson("{'': 5, '': 'a.b', '': 2}"),
                                     fromjson("{'': 5, '': 'a.c', '': 1}"),
                                     fromjson("{'': 5, '': 'a.c.d', '': 1}"),
                                     fromjson("{'': 5, '': 'a.c.d', '': 2}"),
-                                    fromjson("{'': 5, '': 'e', '': 3}"),
-                                    fromjson("{'': 5, '': 'e', '': 5}"),
                                     fromjson("{'': 5, '': 'f', '': 1}")});
 
     auto expectedMultikeyPaths =
         makeKeySet({fromjson("{'': {$minKey: 1}, '': 1, '': 'a'}"),
-                    fromjson("{'': {$minKey: 1}, '': 1, '': 'a.c.d'}"),
-                    fromjson("{'': {$minKey: 1}, '': 1, '': 'e'}")},
+                    fromjson("{'': {$minKey: 1}, '': 1, '': 'a.c.d'}")},
                    RecordIdReservations::reservedIdFor(ReservationId::kWildcardMultikeyMetadataId));
 
     auto outputKeys = makeKeySet();
@@ -1354,33 +1305,21 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathWithArrayField) {
 
 TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathWithArrayElemField) {
     WildcardKeyGenerator keyGen{fromjson("{'a.b' : 1, '$**': 1}"),
-                                {},
+                                fromjson("{a: 0}"),
                                 nullptr,
                                 KeyString::Version::kLatestVersion,
                                 Ordering::make(BSONObj())};
     auto inputDoc = fromjson("{a: [{b: 1, c: 1}, {b: 2, c: {d: [1, 2]}}], e: [3, 5], f: 1}");
 
-    auto expectedKeys = makeKeySet({fromjson("{'': 1, '': 'a.b', '': 1}"),
-                                    fromjson("{'': 1, '': 'a.b', '': 2}"),
-                                    fromjson("{'': 1, '': 'a.c', '': 1}"),
-                                    fromjson("{'': 1, '': 'a.c.d', '': 1}"),
-                                    fromjson("{'': 1, '': 'a.c.d', '': 2}"),
-                                    fromjson("{'': 1, '': 'e', '': 3}"),
+    auto expectedKeys = makeKeySet({fromjson("{'': 1, '': 'e', '': 3}"),
                                     fromjson("{'': 1, '': 'e', '': 5}"),
                                     fromjson("{'': 1, '': 'f', '': 1}"),
-                                    fromjson("{'': 2, '': 'a.b', '': 1}"),
-                                    fromjson("{'': 2, '': 'a.b', '': 2}"),
-                                    fromjson("{'': 2, '': 'a.c', '': 1}"),
-                                    fromjson("{'': 2, '': 'a.c.d', '': 1}"),
-                                    fromjson("{'': 2, '': 'a.c.d', '': 2}"),
                                     fromjson("{'': 2, '': 'e', '': 3}"),
                                     fromjson("{'': 2, '': 'e', '': 5}"),
                                     fromjson("{'': 2, '': 'f', '': 1}")});
 
     auto expectedMultikeyPaths =
-        makeKeySet({fromjson("{'': {$minKey: 1}, '': 1, '': 'a'}"),
-                    fromjson("{'': {$minKey: 1}, '': 1, '': 'a.c.d'}"),
-                    fromjson("{'': {$minKey: 1}, '': 1, '': 'e'}")},
+        makeKeySet({fromjson("{'': {$minKey: 1}, '': 1, '': 'e'}")},
                    RecordIdReservations::reservedIdFor(ReservationId::kWildcardMultikeyMetadataId));
 
     auto outputKeys = makeKeySet();
@@ -1393,33 +1332,19 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathWithArrayElemField)
 
 TEST_F(WildcardKeyGeneratorCompoundTest, MultipleCompoundField) {
     WildcardKeyGenerator keyGen{fromjson("{'a.b': 1, '$**': 1, f: 1}"),
-                                {},
+                                fromjson("{a: 0, f: 0}"),
                                 nullptr,
                                 KeyString::Version::kLatestVersion,
                                 Ordering::make(BSONObj())};
     auto inputDoc = fromjson("{a: [{b: 1, c: 1}, {b: 2, c: {d: [1, 2]}}], e: [3, 5], f: 1}");
 
-    auto expectedKeys = makeKeySet({fromjson("{'': 1, '': 'a.b', '': 1, '': 1}"),
-                                    fromjson("{'': 1, '': 'a.b', '': 2, '': 1}"),
-                                    fromjson("{'': 1, '': 'a.c', '': 1, '': 1}"),
-                                    fromjson("{'': 1, '': 'a.c.d', '': 1, '': 1}"),
-                                    fromjson("{'': 1, '': 'a.c.d', '': 2, '': 1}"),
-                                    fromjson("{'': 1, '': 'e', '': 3, '': 1}"),
+    auto expectedKeys = makeKeySet({fromjson("{'': 1, '': 'e', '': 3, '': 1}"),
                                     fromjson("{'': 1, '': 'e', '': 5, '': 1}"),
-                                    fromjson("{'': 1, '': 'f', '': 1, '': 1}"),
-                                    fromjson("{'': 2, '': 'a.b', '': 1, '': 1}"),
-                                    fromjson("{'': 2, '': 'a.b', '': 2, '': 1}"),
-                                    fromjson("{'': 2, '': 'a.c', '': 1, '': 1}"),
-                                    fromjson("{'': 2, '': 'a.c.d', '': 1, '': 1}"),
-                                    fromjson("{'': 2, '': 'a.c.d', '': 2, '': 1}"),
                                     fromjson("{'': 2, '': 'e', '': 3, '': 1}"),
-                                    fromjson("{'': 2, '': 'e', '': 5, '': 1}"),
-                                    fromjson("{'': 2, '': 'f', '': 1, '': 1}")});
+                                    fromjson("{'': 2, '': 'e', '': 5, '': 1}")});
 
     auto expectedMultikeyPaths =
-        makeKeySet({fromjson("{'': {$minKey: 1}, '': 1, '': 'a', '': {$minKey: 1}}"),
-                    fromjson("{'': {$minKey: 1}, '': 1, '': 'a.c.d', '': {$minKey: 1}}"),
-                    fromjson("{'': {$minKey: 1}, '': 1, '': 'e', '': {$minKey: 1}}")},
+        makeKeySet({fromjson("{'': {$minKey: 1}, '': 1, '': 'e', '': {$minKey: 1}}")},
                    RecordIdReservations::reservedIdFor(ReservationId::kWildcardMultikeyMetadataId));
 
     auto outputKeys = makeKeySet();
@@ -1431,18 +1356,17 @@ TEST_F(WildcardKeyGeneratorCompoundTest, MultipleCompoundField) {
 }
 
 TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathAndDedupKeysCompound) {
-    WildcardKeyGenerator keyGen{fromjson("{'a.c': 1, '$**': 1}"),
+    WildcardKeyGenerator keyGen{fromjson("{'b.d': 1, 'a.$**': 1}"),
                                 {},
                                 nullptr,
                                 KeyString::Version::kLatestVersion,
                                 Ordering::make(BSONObj())};
-    auto inputDoc = fromjson("{a: [1, 2, {b: 'one', c: 2}, {c: 2, d: 3}, {d: 3}]}");
+    auto inputDoc = fromjson("{a: [1, 2, {b: 'one', c: 2}, {c: 2}], b: [{d: 3, e: 4}, {d: 3}]}");
 
-    auto expectedKeys = makeKeySet({fromjson("{'': 2, '': 'a', '': 1}"),
-                                    fromjson("{'': 2, '': 'a', '': 2}"),
-                                    fromjson("{'': 2, '': 'a.b', '': 'one'}"),
-                                    fromjson("{'': 2, '': 'a.c', '': 2}"),
-                                    fromjson("{'': 2, '': 'a.d', '': 3}")});
+    auto expectedKeys = makeKeySet({fromjson("{'': 3, '': 'a', '': 1}"),
+                                    fromjson("{'': 3, '': 'a', '': 2}"),
+                                    fromjson("{'': 3, '': 'a.b', '': 'one'}"),
+                                    fromjson("{'': 3, '': 'a.c', '': 2}")});
 
     auto expectedMultikeyPaths =
         makeKeySet({fromjson("{'': {$minKey: 1}, '': 1, '': 'a'}")},
@@ -1457,7 +1381,7 @@ TEST_F(WildcardKeyGeneratorCompoundTest, ExtractMultikeyPathAndDedupKeysCompound
 }
 
 TEST_F(WildcardKeyGeneratorCompoundTest, ExtractSubtreeWithArrayElemField) {
-    WildcardKeyGenerator keyGen{fromjson("{e : 1, 'a.$**': 1}"),
+    WildcardKeyGenerator keyGen{fromjson("{e: 1, 'a.$**': 1}"),
                                 {},
                                 nullptr,
                                 KeyString::Version::kLatestVersion,
diff --git a/src/mongo/db/query/wildcard_multikey_paths.cpp b/src/mongo/db/query/wildcard_multikey_paths.cpp
index 0677b5b2736..16571e9855d 100644
--- a/src/mongo/db/query/wildcard_multikey_paths.cpp
+++ b/src/mongo/db/query/wildcard_multikey_paths.cpp
@@ -68,7 +68,6 @@ static FieldRef extractMultikeyPathFromIndexKey(BSONObj keyPattern, const IndexK
 
             // Extract the path from the second piece of the key.
             const auto secondIndexElem = indexIter.next();
-            invariant(!indexIter.more());
             invariant(secondIndexElem.type() == BSONType::String);
             return FieldRef(secondIndexElem.valueStringData());
         } else {