1 files changed, 92 insertions, 41 deletions
diff --git a/jstests/core/or_to_in.js b/jstests/core/or_to_in.js
index 3c9cacfbcdf..5d3c745dc95 100644
--- a/jstests/core/or_to_in.js
+++ b/jstests/core/or_to_in.js
@@ -25,19 +25,25 @@ function compareValues(v1, v2) {
 }
 
 // Check that 'expectedQuery' and 'actualQuery' have the same plans, and produce the same result.
-function assertEquivPlanAndResult(expectedQuery, actualQuery) {
+function assertEquivPlanAndResult(expectedQuery, actualQuery, supportWithCollation) {
     const expectedExplain = coll.find(expectedQuery).explain("queryPlanner");
     const actualExplain = coll.find(actualQuery).explain("queryPlanner");
 
     // The queries must be rewritten into the same form.
-    assert.docEq(expectedExplain.parsedQuery, actualExplain.parsedQuery);
+    assert.docEq(expectedExplain.queryPlanner.parsedQuery, actualExplain.queryPlanner.parsedQuery);
 
-    // Check if the test queries produce the same plans with collations
-    const expectedExplainColln =
+    // We are always running these queries to ensure a server crash is not triggered.
+    // TODO SERVER-72450: Add appropriate assertions for the output.
+    const expectedExplainCollation =
         coll.find(expectedQuery).sort({f1: 1}).collation({locale: 'en_US'}).explain("queryPlanner");
-    const actualExplainColln =
+    const actualExplainCollation =
         coll.find(actualQuery).sort({f1: 1}).collation({locale: 'en_US'}).explain("queryPlanner");
-    assert.docEq(expectedExplainColln.parsedQuery, actualExplainColln.parsedQuery);
+
+    if (supportWithCollation) {
+        // Check if the test queries produce the same plans with collations.
+        assert.docEq(expectedExplainCollation.queryPlanner.parsedQuery,
+                     actualExplainCollation.queryPlanner.parsedQuery);
+    }
 
     // Make sure both queries have the same access plan.
     const expectedPlan = getWinningPlan(expectedExplain.queryPlanner);
@@ -51,12 +57,13 @@ function assertEquivPlanAndResult(expectedQuery, actualQuery) {
     const actualRes = coll.find(actualQuery).toArray();
     assert(arrayEq(expectedRes, actualRes, false, compareValues),
            `expected=${expectedRes}, actual=${actualRes}`);
+
     // also with collation
-    const expectedResColln =
+    const expectedResCollation =
         coll.find(expectedQuery).sort({f1: 1}).collation({locale: 'en_US'}).toArray();
-    const actualResColln =
+    const actualResCollation =
         coll.find(actualQuery).sort({f1: 1}).collation({locale: 'en_US'}).toArray();
-    assert(arrayEq(expectedResColln, actualResColln, false, compareValues),
+    assert(arrayEq(expectedResCollation, actualResCollation, false, compareValues),
            `expected=${expectedRes}, actual=${actualRes}`);
 }
 
@@ -94,32 +101,71 @@ assert.commandWorked(coll.insert(data));
 
 // Pairs of queries where the first one is expressed via OR (which is supposed to be
 // rewritten as IN), and the second one is an equivalent query using IN.
+//
+// The third element of the array is optional, if present, implies that the rewrite is not
+// supported when there is a collation involved.
+//
+// TODO SERVER-72450: Remove or update this logic related to collation, and enforce stronger
+// assertions.
 const positiveTestQueries = [
-    [{$or: [{f1: 5}, {f1: 3}, {f1: 7}]}, {f1: {$in: [7, 3, 5]}}],
-    [{$or: [{f1: {$eq: 5}}, {f1: {$eq: 3}}, {f1: {$eq: 7}}]}, {f1: {$in: [7, 3, 5]}}],
-    [{$or: [{f1: 42}, {f1: NaN}, {f1: 99}]}, {f1: {$in: [42, NaN, 99]}}],
-    [{$or: [{f1: /^x/}, {f1: "ab"}]}, {f1: {$in: [/^x/, "ab"]}}],
-    [{$or: [{f1: /^x/}, {f1: "^a"}]}, {f1: {$in: [/^x/, "^a"]}}],
-    [{$or: [{f1: 42}, {f1: null}, {f1: 99}]}, {f1: {$in: [42, 99, null]}}],
-    [{$or: [{f1: 1}, {f2: 9}, {f1: 99}]}, {$or: [{f2: 9}, {f1: {$in: [1, 99]}}]}],
-    [{$or: [{f1: {$regex: /^x/}}, {f1: {$regex: /ab/}}]}, {f1: {$in: [/^x/, /ab/]}}],
-    [
-        {$and: [{$or: [{f1: 7}, {f1: 3}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
-        {$and: [{f1: {$in: [7, 3, 5]}}, {f1: {$in: [1, 2, 3]}}]}
-    ],
-    [
-        {$or: [{$or: [{f1: 7}, {f1: 3}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
-        {$or: [{f1: {$in: [7, 3, 5]}}, {f1: {$in: [1, 2, 3]}}]}
-    ],
-    [
-        {$or: [{$and: [{f1: 7}, {f2: 7}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
-        {$or: [{$and: [{f1: 7}, {f2: 7}, {f1: 5}]}, {f1: {$in: [1, 2, 3]}}]},
-    ],
-    [{$or: [{f2: [32, 52]}, {f2: [42, [13, 11]]}]}, {f2: {$in: [[32, 52], [42, [13, 11]]]}}],
-    [{$or: [{f2: 52}, {f2: 13}]}, {f2: {$in: [52, 13]}}],
-    [{$or: [{f2: [11]}, {f2: [23]}]}, {f2: {$in: [[11], [23]]}}],
-    [{$or: [{f1: 42}, {f1: null}]}, {f1: {$in: [42, null]}}],
-    [{$or: [{f1: "a"}, {f1: "b"}, {f1: /c/}]}, {f1: {$in: ["a", "b", /c/]}}],
+    {actualQuery: {$or: [{f1: 5}, {f1: 3}, {f1: 7}]}, expectedQuery: {f1: {$in: [7, 3, 5]}}},
+    {
+        actualQuery: {$or: [{f1: {$eq: 5}}, {f1: {$eq: 3}}, {f1: {$eq: 7}}]},
+        expectedQuery: {f1: {$in: [7, 3, 5]}}
+    },
+    {
+        actualQuery: {$or: [{f1: 42}, {f1: NaN}, {f1: 99}]},
+        expectedQuery: {f1: {$in: [42, NaN, 99]}}
+    },
+    {
+        actualQuery: {$or: [{f1: /^x/}, {f1: "ab"}]},
+        expectedQuery: {f1: {$in: [/^x/, "ab"]}},
+        cannotRewriteWithCollation: true
+    },
+    {
+        actualQuery: {$or: [{f1: /^x/}, {f1: "^a"}]},
+        expectedQuery: {f1: {$in: [/^x/, "^a"]}},
+        cannotRewriteWithCollation: true
+    },
+    {
+        actualQuery: {$or: [{f1: 42}, {f1: null}, {f1: 99}]},
+        expectedQuery: {f1: {$in: [42, 99, null]}}
+    },
+    {
+        actualQuery: {$or: [{f1: 1}, {f2: 9}, {f1: 99}]},
+        expectedQuery: {$or: [{f2: 9}, {f1: {$in: [1, 99]}}]}
+    },
+    {
+        actualQuery: {$or: [{f1: {$regex: /^x/}}, {f1: {$regex: /ab/}}]},
+        expectedQuery: {f1: {$in: [/^x/, /ab/]}}
+    },
+    {
+        actualQuery:
+            {$and: [{$or: [{f1: 7}, {f1: 3}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
+        expectedQuery: {$and: [{f1: {$in: [7, 3, 5]}}, {f1: {$in: [1, 2, 3]}}]}
+    },
+    {
+        actualQuery:
+            {$or: [{$or: [{f1: 7}, {f1: 3}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
+        expectedQuery: {$or: [{f1: {$in: [7, 3, 5]}}, {f1: {$in: [1, 2, 3]}}]}
+    },
+    {
+        actualQuery:
+            {$or: [{$and: [{f1: 7}, {f2: 7}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
+        expectedQuery: {$or: [{$and: [{f1: 7}, {f2: 7}, {f1: 5}]}, {f1: {$in: [1, 2, 3]}}]},
+    },
+    {
+        actualQuery: {$or: [{f2: [32, 52]}, {f2: [42, [13, 11]]}]},
+        expectedQuery: {f2: {$in: [[32, 52], [42, [13, 11]]]}}
+    },
+    {actualQuery: {$or: [{f2: 52}, {f2: 13}]}, expectedQuery: {f2: {$in: [52, 13]}}},
+    {actualQuery: {$or: [{f2: [11]}, {f2: [23]}]}, expectedQuery: {f2: {$in: [[11], [23]]}}},
+    {actualQuery: {$or: [{f1: 42}, {f1: null}]}, expectedQuery: {f1: {$in: [42, null]}}},
+    {
+        actualQuery: {$or: [{f1: "a"}, {f1: "b"}, {f1: /c/}]},
+        expectedQuery: {f1: {$in: ["a", "b", /c/]}},
+        cannotRewriteWithCollation: true
+    },
 ];
 
 // These $or queries should not be rewritten into $in because of different semantics.
@@ -133,22 +179,27 @@ for (const query of negativeTestQueries) {
     assertOrNotRewrittenToIn(query);
 }
 
-function testOrToIn(queries) {
+function testOrToIn(queries, usesCollation) {
     for (const queryPair of queries) {
-        assertEquivPlanAndResult(queryPair[0], queryPair[1]);
+        if (usesCollation && queryPair.cannotRewriteWithCollation) {
+            continue;
+        }
+        assertEquivPlanAndResult(
+            queryPair.actualQuery, queryPair.expectedQuery, !queryPair.cannotRewriteWithCollation);
     }
 }
 
-testOrToIn(positiveTestQueries);  // test without indexes
+testOrToIn(positiveTestQueries, false /* usesCollation */);  // test without indexes
 
 assert.commandWorked(coll.createIndex({f1: 1}));
 
-testOrToIn(positiveTestQueries);  // single index
+testOrToIn(positiveTestQueries, false /* usesCollation */);  // single index
 
 assert.commandWorked(coll.createIndex({f2: 1}));
 assert.commandWorked(coll.createIndex({f1: 1, f2: 1}));
 
-testOrToIn(positiveTestQueries);  // three indexes, requires multiplanning
+testOrToIn(positiveTestQueries,
+           false /* usesCollation */);  // three indexes, requires multiplanning
 
 // Test with a collection that has a collation, and that collation is the same as the query
 // collation
@@ -156,12 +207,12 @@ coll.drop();
 assert.commandWorked(db.createCollection("orToIn", {collation: {locale: 'en_US'}}));
 coll = db.orToIn;
 assert.commandWorked(coll.insert(data));
-testOrToIn(positiveTestQueries);
+testOrToIn(positiveTestQueries, true /* usesCollation */);
 // Test with a collection that has a collation, and that collation is different from the query
 // collation
 coll.drop();
 assert.commandWorked(db.createCollection("orToIn", {collation: {locale: 'de'}}));
 coll = db.orToIn;
 assert.commandWorked(coll.insert(data));
-testOrToIn(positiveTestQueries);
+testOrToIn(positiveTestQueries, true /* usesCollation */);
 }());