SERVER-66379 $or to $in conversion flawed

(cherry picked from commit 951197b3834d509d7903efa4a66c0e9f79235262)
(cherry picked from commit b499c915d1c21f063548c96d5d547dcfcda55754)

1 files changed, 37 insertions, 3 deletions

diff --git a/jstests/core/or_to_in.js b/jstests/core/or_to_in.js
index 2dfb893776f..8d080256d05 100644
--- a/jstests/core/or_to_in.js
+++ b/jstests/core/or_to_in.js
@@ -13,7 +13,7 @@
 load("jstests/aggregation/extras/utils.js");  // For assertArrayEq.
 load("jstests/libs/analyze_plan.js");
 
-const coll = db.orToIn;
+var coll = db.orToIn;
 coll.drop();
 
 function compareValues(v1, v2) {
@@ -30,6 +30,14 @@ function assertEquivPlanAndResult(expectedQuery, actualQuery) {
     const actualExplain = coll.find(actualQuery).explain("queryPlanner");
     // The queries must be rewritten into the same form.
     assert.docEq(expectedExplain.parsedQuery, actualExplain.parsedQuery);
+
+    // Check if the test queries produce the same plans with collations
+    const expectedExplainColln =
+        coll.find(expectedQuery).sort({f1: 1}).collation({locale: 'en_US'}).explain("queryPlanner");
+    const actualExplainColln =
+        coll.find(actualQuery).sort({f1: 1}).collation({locale: 'en_US'}).explain("queryPlanner");
+    assert.docEq(expectedExplainColln.parsedQuery, actualExplainColln.parsedQuery);
+
     // Make sure both queries have the same access plan.
     const expectedPlan = getWinningPlan(expectedExplain.queryPlanner);
     const actualPlan = getWinningPlan(actualExplain.queryPlanner);
@@ -39,6 +47,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 =
+        coll.find(expectedQuery).sort({f1: 1}).collation({locale: 'en_US'}).toArray();
+    const actualResColln =
+        coll.find(actualQuery).sort({f1: 1}).collation({locale: 'en_US'}).toArray();
+    assert(arrayEq(expectedResColln, actualResColln, false, compareValues),
+           `expected=${expectedRes}, actual=${actualRes}`);
 }
 
 // Make sure that certain $or expressions are not rewritten to $in.
@@ -56,7 +71,7 @@ function assertOrNotRewrittenToIn(query) {
     }
 }
 
-assert.commandWorked(coll.insert([
+const data = [
     {_id: 0, f1: 3, f2: 7},
     {_id: 1, f1: 1, f2: [32, 42, 52, [11]]},
     {_id: 2, f1: 2, f2: 9},
@@ -69,7 +84,9 @@ assert.commandWorked(coll.insert([
     {_id: 9, f1: NaN, f2: NaN},
     {_id: 10, f1: 1, f2: [32, 52]},
     {_id: 11, f1: 1, f2: [42, [13, 11]]}
-]));
+];
+
+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.
@@ -98,6 +115,8 @@ const positiveTestQueries = [
     [{$or: [{f2: 52}, {f2: 13}]}, {f2: {$in: [52, 13]}}],
     [{$or: [{f2: [11]}, {f2: [23]}]}, {f2: {$in: [[11], [23]]}}],
     {$or: [{f1: 42}, {f1: null}]},
+    [{$or: [{f1: 42}, {f1: null}]}, {f1: {$in: [42, null]}}],
+    [{$or: [{f1: "a"}, {f1: "b"}, {f1: /c/}]}, {f1: {$in: ["a", "b", /c/]}}],
 ];
 
 // These $or queries should not be rewritten into $in because of different semantics.
@@ -127,4 +146,19 @@ assert.commandWorked(coll.createIndex({f2: 1}));
 assert.commandWorked(coll.createIndex({f1: 1, f2: 1}));
 
 testOrToIn(positiveTestQueries);  // three indexes, requires multiplanning
+
+// Test with a collection that has a collation, and that collation is the same as the query
+// collation
+coll.drop();
+assert.commandWorked(db.createCollection("orToIn", {collation: {locale: 'en_US'}}));
+coll = db.orToIn;
+assert.commandWorked(coll.insert(data));
+testOrToIn(positiveTestQueries);
+// 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);
 }());