SERVER-42905 Push down user-specified projections to PlanStage layer if possible

4 files changed, 165 insertions, 151 deletions

diff --git a/jstests/aggregation/bugs/server6192_server6193.js b/jstests/aggregation/bugs/server6192_server6193.js
index 9b74bdb3b4d..2194c19c3f9 100644
--- a/jstests/aggregation/bugs/server6192_server6193.js
+++ b/jstests/aggregation/bugs/server6192_server6193.js
@@ -7,31 +7,25 @@
 // This test makes assumptions about how the explain output will be formatted, so cannot be
 // transformed to be put inside a $facet stage.
 // @tags: [do_not_wrap_aggregations_in_facets,assumes_unsharded_collection]
+(function() {
+"use strict";
 
-var t = db.jstests_aggregation_server6192;
+load("jstests/libs/analyze_plan.js");  // For 'getPlanStage'.
+
+const t = db.jstests_aggregation_server6192;
 t.drop();
-t.save({x: true});
+assert.commandWorked(t.insert({x: true}));
 
 function assertOptimized(pipeline, v) {
-    var explained = t.runCommand("aggregate", {pipeline: pipeline, explain: true});
-
-    printjson({input: pipeline, output: explained});
-
-    assert("stages" in explained);
-    assert("$project" in explained.stages[1]);
-    var projectStage = explained.stages[1]["$project"];
-    assert.eq(projectStage.a["$const"], v, "ensure short-circuiting worked");
+    const explained = t.runCommand("aggregate", {pipeline: pipeline, explain: true});
+    const projectStage = getPlanStage(explained, "PROJECTION_DEFAULT");
+    assert.eq(projectStage.transformBy.a["$const"], v, "ensure short-circuiting worked", explained);
 }
 
 function assertNotOptimized(pipeline) {
-    var explained = t.runCommand("aggregate", {pipeline: pipeline, explain: true});
-
-    printjson({input: pipeline, output: explained});
-
-    assert("stages" in explained);
-    assert("$project" in explained.stages[1]);
-    var projectStage = explained.stages[1]["$project"];
-    assert(!("$const" in projectStage.a), "ensure no short-circuiting");
+    const explained = t.runCommand("aggregate", {pipeline: pipeline, explain: true});
+    const projectStage = getPlanStage(explained, "PROJECTION_DEFAULT");
+    assert(!("$const" in projectStage.transformBy.a), "ensure no short-circuiting");
 }
 
 // short-circuiting for $and
@@ -59,3 +53,4 @@ assertNotOptimized([{$project: {a: {$and: ['$x', '$x']}}}]);
 assertNotOptimized([{$project: {a: {$or: ['$x', '$x']}}}]);
 assertNotOptimized([{$project: {a: {$and: ['$x']}}}]);
 assertNotOptimized([{$project: {a: {$or: ['$x']}}}]);
+}());
diff --git a/jstests/aggregation/optimize_away_pipeline.js b/jstests/aggregation/optimize_away_pipeline.js
index d6a27573333..6119cc30d48 100644
--- a/jstests/aggregation/optimize_away_pipeline.js
+++ b/jstests/aggregation/optimize_away_pipeline.js
@@ -41,19 +41,12 @@ function assertPipelineUsesAggregation({
 } = {}) {
     const explainOutput = coll.explain().aggregate(pipeline, pipelineOptions);
 
-    assert(isAggregationPlan(explainOutput),
-           "Expected pipeline " + tojsononeline(pipeline) +
-               " to use an aggregation framework in the explain output: " + tojson(explainOutput));
-    assert(!isQueryPlan(explainOutput),
-           "Expected pipeline " + tojsononeline(pipeline) +
-               " *not* to use a query layer at the root level in the explain output: " +
-               tojson(explainOutput));
+    assert(isAggregationPlan(explainOutput), explainOutput);
+    assert(!isQueryPlan(explainOutput), explainOutput);
 
     if (optimizedAwayStages) {
         for (let stage of optimizedAwayStages) {
-            assert(!aggPlanHasStage(explainOutput, stage),
-                   "Expected pipeline " + tojsononeline(pipeline) + " to *not* include a " + stage +
-                       " stage in the explain output: " + tojson(explainOutput));
+            assert(!aggPlanHasStage(explainOutput, stage), explainOutput);
         }
     }
 
@@ -64,18 +57,12 @@ function assertPipelineUsesAggregation({
         cursor = getAggPlanStage(explainOutput, "$geoNearCursor").$geoNearCursor;
     }
 
-    assert(cursor,
-           "Expected pipeline " + tojsononeline(pipeline) + " to include a $cursor " +
-               " stage in the explain output: " + tojson(explainOutput));
-    assert(cursor.queryPlanner.optimizedPipeline === undefined,
-           "Expected pipeline " + tojsononeline(pipeline) + " to *not* include an " +
-               "'optimizedPipeline' field in the explain output: " + tojson(explainOutput));
+    assert(cursor, explainOutput);
+    assert(cursor.queryPlanner.optimizedPipeline === undefined, explainOutput);
 
     if (expectedStages) {
         for (let expectedStage of expectedStages) {
-            assert(aggPlanHasStage(explainOutput, expectedStage),
-                   "Expected pipeline " + tojsononeline(pipeline) + " to include a " +
-                       expectedStage + " stage in the explain output: " + tojson(explainOutput));
+            assert(aggPlanHasStage(explainOutput, expectedStage), explainOutput);
         }
     }
 
@@ -103,32 +90,20 @@ function assertPipelineDoesNotUseAggregation({
 } = {}) {
     const explainOutput = coll.explain().aggregate(pipeline, pipelineOptions);
 
-    assert(!isAggregationPlan(explainOutput),
-           "Expected pipeline " + tojsononeline(pipeline) +
-               " *not* to use an aggregation framework in the explain output: " +
-               tojson(explainOutput));
-    assert(isQueryPlan(explainOutput),
-           "Expected pipeline " + tojsononeline(pipeline) +
-               " to use a query layer at the root level in the explain output: " +
-               tojson(explainOutput));
+    assert(!isAggregationPlan(explainOutput), explainOutput);
+    assert(isQueryPlan(explainOutput), explainOutput);
     if (explainOutput.hasOwnProperty("shards")) {
         Object.keys(explainOutput.shards)
             .forEach((shard) =>
                          assert(explainOutput.shards[shard].queryPlanner.optimizedPipeline === true,
-                                "Expected pipeline " + tojsononeline(pipeline) + " to include an " +
-                                    "'optimizedPipeline' field in the explain output: " +
-                                    tojson(explainOutput)));
+                                explainOutput));
     } else {
-        assert(explainOutput.queryPlanner.optimizedPipeline === true,
-               "Expected pipeline " + tojsononeline(pipeline) + " to include an " +
-                   "'optimizedPipeline' field in the explain output: " + tojson(explainOutput));
+        assert(explainOutput.queryPlanner.optimizedPipeline === true, explainOutput);
     }
 
     if (expectedStages) {
         for (let expectedStage of expectedStages) {
-            assert(planHasStage(db, explainOutput, expectedStage),
-                   "Expected pipeline " + tojsononeline(pipeline) + " to include a " +
-                       expectedStage + " stage in the explain output: " + tojson(explainOutput));
+            assert(planHasStage(db, explainOutput, expectedStage), explainOutput);
         }
     }
 
@@ -200,45 +175,46 @@ assertPipelineDoesNotUseAggregation({
 });
 assert.commandWorked(coll.dropIndexes());
 
-// Pipelines which cannot be optimized away.
-
-// TODO SERVER-40254: Uncovered queries.
 assert.commandWorked(coll.insert({_id: 4, x: 40, a: {b: "ab1"}}));
-assertPipelineUsesAggregation({
+assertPipelineDoesNotUseAggregation({
     pipeline: [{$project: {x: 1, _id: 0}}],
-    expectedStages: ["COLLSCAN"],
+    expectedStages: ["COLLSCAN", "PROJECTION_SIMPLE"],
     expectedResult: [{x: 10}, {x: 20}, {x: 30}, {x: 40}]
 });
-assertPipelineUsesAggregation({
+assertPipelineDoesNotUseAggregation({
     pipeline: [{$match: {x: 20}}, {$project: {x: 1, _id: 0}}],
-    expectedStages: ["COLLSCAN"],
+    expectedStages: ["COLLSCAN", "PROJECTION_SIMPLE"],
     expectedResult: [{x: 20}]
 });
-assertPipelineUsesAggregation({
+assertPipelineDoesNotUseAggregation({
     pipeline: [{$project: {x: 1, "a.b": 1, _id: 0}}],
-    expectedStages: ["COLLSCAN"],
+    expectedStages: ["COLLSCAN", "PROJECTION_DEFAULT"],
     expectedResult: [{x: 10}, {x: 20}, {x: 30}, {x: 40, a: {b: "ab1"}}]
 });
-assertPipelineUsesAggregation({
+assertPipelineDoesNotUseAggregation({
     pipeline: [{$match: {x: 40}}, {$project: {"a.b": 1, _id: 0}}],
-    expectedStages: ["COLLSCAN"],
+    expectedStages: ["COLLSCAN", "PROJECTION_DEFAULT"],
     expectedResult: [{a: {b: "ab1"}}]
 });
+// We can collapse a $project stage if it has a complex pipeline expression.
+assertPipelineDoesNotUseAggregation({
+    pipeline: [{$project: {x: {$substr: ["$y", 0, 1]}, _id: 0}}],
+    expectedStages: ["COLLSCAN", "PROJECTION_DEFAULT"]
+});
+assertPipelineDoesNotUseAggregation({
+    pipeline: [{$match: {x: 20}}, {$project: {x: {$substr: ["$y", 0, 1]}, _id: 0}}],
+    expectedStages: ["COLLSCAN", "PROJECTION_DEFAULT"]
+});
 assert.commandWorked(coll.deleteOne({_id: 4}));
 
+// Pipelines which cannot be optimized away.
+
 // TODO SERVER-40909: $skip stage is not supported yet.
 assertPipelineUsesAggregation({
     pipeline: [{$match: {x: {$gte: 20}}}, {$skip: 1}],
     expectedStages: ["COLLSCAN"],
     expectedResult: [{_id: 3, x: 30}]
 });
-// We cannot collapse a $project stage if it has a complex pipeline expression.
-assertPipelineUsesAggregation(
-    {pipeline: [{$project: {x: {$substr: ["$y", 0, 1]}, _id: 0}}], expectedStages: ["COLLSCAN"]});
-assertPipelineUsesAggregation({
-    pipeline: [{$match: {x: 20}}, {$project: {x: {$substr: ["$y", 0, 1]}, _id: 0}}],
-    expectedStages: ["COLLSCAN"]
-});
 // We cannot optimize away a pipeline if there are stages which have no equivalent in the
 // find command.
 assertPipelineUsesAggregation({
@@ -304,13 +280,12 @@ assertPipelineDoesNotUseAggregation({
     expectedResult: [{x: 20}],
 });
 
-// $match, $project, and $limit cannot be optimized away when the projection is not covered. But the
-// $limit can be pushed down into the query layer.
-assertPipelineUsesAggregation({
+// $match, $project, and $limit can be optimized away.
+assertPipelineDoesNotUseAggregation({
     pipeline: [{$match: {x: {$gte: 20}}}, {$project: {_id: 0, x: 1, y: 1}}, {$limit: 1}],
-    expectedStages: ["IXSCAN", "FETCH", "LIMIT"],
+    expectedStages: ["IXSCAN", "FETCH", "LIMIT", "PROJECTION_SIMPLE"],
     expectedResult: [{x: 20}],
-    optimizedAwayStages: ["$limit"],
+    optimizedAwayStages: ["$limit", "$project"],
 });
 
 // $match, $project, $limit, $sort cannot be optimized away because the $limit comes before the
@@ -454,13 +429,11 @@ assertPipelineDoesNotUseAggregation({
     expectedResult: [{_id: 3, x: 30}],
 });
 
-// If there is a $project that can't result in a covered plan, however, then the pipeline cannot be
-// optimized away. But the $sort should still get pushed down into the PlanStage layer.
-assertPipelineUsesAggregation({
+// $match, $sort, $project, $limit can be optimized away.
+assertPipelineDoesNotUseAggregation({
     pipeline:
         [{$match: {x: {$gte: 20}}}, {$sort: {x: -1}}, {$project: {_id: 0, x: 1}}, {$limit: 2}],
-    expectedStages: ["COLLSCAN", "SORT"],
-    optimizedAwayStages: ["$match", "$sort", "$limit"],
+    expectedStages: ["COLLSCAN", "SORT", "PROJECTION_SIMPLE"],
     expectedResult: [{x: 30}, {x: 20}],
 });
 
@@ -479,6 +452,61 @@ assertPipelineDoesNotUseAggregation({
 });
 assert.commandWorked(coll.dropIndexes());
 
+// Test that even if we don't have a projection stage at the front of the pipeline but there is a
+// finite dependency set, a projection representing this dependency set is pushed down.
+pipeline = [{$group: {_id: "$a", b: {$sum: "$b"}}}];
+assertPipelineUsesAggregation({
+    pipeline: pipeline,
+    expectedStages: ["COLLSCAN", "PROJECTION_SIMPLE"],
+});
+explain = coll.explain().aggregate(pipeline);
+let projStage = getAggPlanStage(explain, "PROJECTION_SIMPLE");
+assert.neq(null, projStage, explain);
+assert.eq({a: 1, b: 1, _id: 0}, projStage.transformBy, explain);
+
+// Similar as above, but with $addFields stage at the front of the pipeline.
+pipeline = [{$addFields: {z: "abc"}}, {$group: {_id: "$a", b: {$sum: "$b"}}}];
+assertPipelineUsesAggregation({
+    pipeline: pipeline,
+    expectedStages: ["COLLSCAN", "PROJECTION_SIMPLE"],
+});
+explain = coll.explain().aggregate(pipeline);
+projStage = getAggPlanStage(explain, "PROJECTION_SIMPLE");
+assert.neq(null, projStage, explain);
+assert.eq({a: 1, b: 1, _id: 0}, projStage.transformBy, explain);
+
+// We generate a projection stage from dependency analysis, even if the pipeline begins with an
+// exclusion projection.
+pipeline = [{$project: {c: 0}}, {$group: {_id: "$a", b: {$sum: "$b"}}}];
+assertPipelineUsesAggregation({
+    pipeline: pipeline,
+    expectedStages: ["COLLSCAN", "PROJECTION_SIMPLE", "$project"],
+});
+explain = coll.explain().aggregate(pipeline);
+projStage = getAggPlanStage(explain, "PROJECTION_SIMPLE");
+assert.neq(null, projStage, explain);
+assert.eq({a: 1, b: 1, _id: 0}, projStage.transformBy, explain);
+
+// Similar as above, but with a field 'a' presented both in the finite dependency set, and in the
+// exclusion projection at the front of the pipeline.
+pipeline = [{$project: {a: 0}}, {$group: {_id: "$a", b: {$sum: "$b"}}}];
+assertPipelineUsesAggregation({
+    pipeline: pipeline,
+    expectedStages: ["COLLSCAN", "PROJECTION_SIMPLE", "$project"],
+});
+explain = coll.explain().aggregate(pipeline);
+projStage = getAggPlanStage(explain, "PROJECTION_SIMPLE");
+assert.neq(null, projStage, explain);
+assert.eq({a: 1, b: 1, _id: 0}, projStage.transformBy, explain);
+
+// Test that an exclusion projection at the front of the pipeline is not pushed down, if there no
+// finite dependency set.
+pipeline = [{$project: {x: 0}}];
+assertPipelineUsesAggregation({pipeline: pipeline, expectedStages: ["COLLSCAN"]});
+explain = coll.explain().aggregate(pipeline);
+assert(!planHasStage(db, explain, "PROJECTION_SIMPLE"), explain);
+assert(!planHasStage(db, explain, "PROJECTION_DEFAULT"), explain);
+
 // getMore cases.
 
 // Test getMore on a collection with an optimized away pipeline.
diff --git a/jstests/aggregation/sources/project/remove_redundant_projects.js b/jstests/aggregation/sources/project/remove_redundant_projects.js
index ea620be2c2a..e686db8291f 100644
--- a/jstests/aggregation/sources/project/remove_redundant_projects.js
+++ b/jstests/aggregation/sources/project/remove_redundant_projects.js
@@ -41,10 +41,10 @@ function assertResultsMatch({
         let result;
 
         if (pipelineOptimizedAway) {
-            assert(isQueryPlan(explain));
+            assert(isQueryPlan(explain), explain);
             result = explain.queryPlanner.winningPlan;
         } else {
-            assert(isAggregationPlan(explain));
+            assert(isAggregationPlan(explain), explain);
             result = explain.stages[0].$cursor.queryPlanner.winningPlan;
         }
 
@@ -52,13 +52,14 @@ function assertResultsMatch({
         assert.eq(expectProjectToCoalesce,
                   planHasStage(db, result, "PROJECTION_DEFAULT") ||
                       planHasStage(db, result, "PROJECTION_COVERED") ||
-                      planHasStage(db, result, "PROJECTION_SIMPLE"));
+                      planHasStage(db, result, "PROJECTION_SIMPLE"),
+                  explain);
 
         if (!pipelineOptimizedAway) {
             // Check that $project was removed from pipeline and pushed to the query system.
             explain.stages.forEach(function(stage) {
                 if (stage.hasOwnProperty("$project"))
-                    assert.neq(removedProjectStage, stage["$project"]);
+                    assert.neq(removedProjectStage, stage["$project"], explain);
             });
         }
     }
@@ -75,7 +76,6 @@ function assertResultsMatch({
 assertResultsMatch({
     pipeline: [{$project: {_id: 0, a: 1}}],
     expectProjectToCoalesce: true,
-    removedProjectStage: {_id: 0, a: 1},
     pipelineOptimizedAway: true
 });
 assertResultsMatch({
@@ -86,7 +86,6 @@ assertResultsMatch({
 assertResultsMatch({
     pipeline: [{$sort: {a: -1}}, {$project: {_id: 0, a: 1}}],
     expectProjectToCoalesce: true,
-    removedProjectStage: {_id: 0, a: 1},
     pipelineOptimizedAway: true
 });
 assertResultsMatch({
@@ -101,32 +100,48 @@ assertResultsMatch({
 assertResultsMatch({
     pipeline: [{$project: {_id: 0, c: {d: 1}}}],
     expectProjectToCoalesce: true,
-    removedProjectStage: {_id: 0, c: {d: 1}},
     pipelineOptimizedAway: true
 });
 
-// Test that projections with renamed fields are not removed from the pipeline, however an
-// inclusion projection is still pushed to the query system.
-assertResultsMatch({pipeline: [{$project: {_id: 0, f: "$a"}}], expectProjectToCoalesce: true});
-assertResultsMatch(
-    {pipeline: [{$project: {_id: 0, a: 1, f: "$a"}}], expectProjectToCoalesce: true});
+// Test that projections with renamed fields are removed from the pipeline.
+assertResultsMatch({
+    pipeline: [{$project: {_id: 0, f: "$a"}}],
+    expectProjectToCoalesce: true,
+    pipelineOptimizedAway: true
+});
+assertResultsMatch({
+    pipeline: [{$project: {_id: 0, a: 1, f: "$a"}}],
+    expectProjectToCoalesce: true,
+    pipelineOptimizedAway: true
+});
 
-// Test that uncovered projections include the $project stage in the pipeline.
-assertResultsMatch(
-    {pipeline: [{$sort: {a: 1}}, {$project: {_id: 1, b: 1}}], expectProjectToCoalesce: false});
+// Test that uncovered projections are removed from the pipeline.
+assertResultsMatch({
+    pipeline: [{$sort: {a: 1}}, {$project: {_id: 1, b: 1}}],
+    expectProjectToCoalesce: true,
+    pipelineOptimizedAway: true
+});
 assertResultsMatch({
     pipeline: [{$sort: {a: 1}}, {$group: {_id: "$_id", arr: {$push: "$a"}}}, {$project: {arr: 1}}],
-    expectProjectToCoalesce: false
+    expectProjectToCoalesce: true
 });
 
-// Test that projections with computed fields are kept in the pipeline.
-assertResultsMatch(
-    {pipeline: [{$project: {computedField: {$sum: "$a"}}}], expectProjectToCoalesce: false});
-assertResultsMatch({pipeline: [{$project: {a: ["$a", "$b"]}}], expectProjectToCoalesce: false});
+// Test that projections with computed fields are removed from the pipeline.
+assertResultsMatch({
+    pipeline: [{$project: {computedField: {$sum: "$a"}}}],
+    expectProjectToCoalesce: true,
+    pipelineOptimizedAway: true
+});
+assertResultsMatch({
+    pipeline: [{$project: {a: ["$a", "$b"]}}],
+    expectProjectToCoalesce: true,
+    pipelineOptimizedAway: true
+});
 assertResultsMatch({
     pipeline:
         [{$project: {e: {$filter: {input: "$e", as: "item", cond: {"$eq": ["$$item", "elem0"]}}}}}],
-    expectProjectToCoalesce: false
+    expectProjectToCoalesce: true,
+    pipelineOptimizedAway: true
 });
 
 // Test that only the first projection is removed from the pipeline.
@@ -150,5 +165,6 @@ assertResultsMatch({
     expectProjectToCoalesce: true,
     index: indexSpec,
     pipelineOptimizedAway: true,
+    removedProjectStage: {'_id.a': 1},
 });
 }());
diff --git a/jstests/aggregation/use_query_projection.js b/jstests/aggregation/use_query_projection.js
index 248ad19bbf6..030a786f807 100644
--- a/jstests/aggregation/use_query_projection.js
+++ b/jstests/aggregation/use_query_projection.js
@@ -19,31 +19,20 @@ for (let i = 0; i < 100; ++i) {
 }
 assert.commandWorked(bulk.execute());
 
-function assertQueryCoversProjection({pipeline = [], pipelineOptimizedAway = true} = {}) {
-    const explainOutput = coll.explain().aggregate(pipeline);
+function assertQueryCoversProjection(
+    {pipeline = [], pipelineOptimizedAway = true, options = {}} = {}) {
+    const explainOutput = coll.explain().aggregate(pipeline, options);
 
     if (pipelineOptimizedAway) {
         assert(isQueryPlan(explainOutput), explainOutput);
-        assert(
-            !planHasStage(db, explainOutput, "FETCH"),
-            "Expected pipeline " + tojsononeline(pipeline) +
-                " *not* to include a FETCH stage in the explain output: " + tojson(explainOutput));
-        assert(planHasStage(db, explainOutput, "IXSCAN"),
-               "Expected pipeline " + tojsononeline(pipeline) +
-                   " to include an index scan in the explain output: " + tojson(explainOutput));
+        assert(!planHasStage(db, explainOutput, "FETCH"), explainOutput);
+        assert(planHasStage(db, explainOutput, "IXSCAN"), explainOutput);
     } else {
         assert(isAggregationPlan(explainOutput), explainOutput);
-        assert(
-            !aggPlanHasStage(explainOutput, "FETCH"),
-            "Expected pipeline " + tojsononeline(pipeline) +
-                " *not* to include a FETCH stage in the explain output: " + tojson(explainOutput));
-        assert(aggPlanHasStage(explainOutput, "IXSCAN"),
-               "Expected pipeline " + tojsononeline(pipeline) +
-                   " to include an index scan in the explain output: " + tojson(explainOutput));
+        assert(!aggPlanHasStage(explainOutput, "FETCH"), explainOutput);
+        assert(aggPlanHasStage(explainOutput, "IXSCAN"), explainOutput);
     }
-    assert(!hasRejectedPlans(explainOutput),
-           "Expected pipeline " + tojsononeline(pipeline) +
-               " not to have any rejected plans in the explain output: " + tojson(explainOutput));
+    assert(!hasRejectedPlans(explainOutput), explainOutput);
     return explainOutput;
 }
 
@@ -51,27 +40,15 @@ function assertQueryDoesNotCoverProjection({pipeline = [], pipelineOptimizedAway
     const explainOutput = coll.explain().aggregate(pipeline);
 
     if (pipelineOptimizedAway) {
-        assert(isQueryPlan(explainOutput));
+        assert(isQueryPlan(explainOutput), explainOutput);
         assert(planHasStage(db, explainOutput, "FETCH") || aggPlanHasStage("COLLSCAN"),
-               "Expected pipeline " + tojsononeline(pipeline) +
-                   " to include a FETCH or COLLSCAN stage in the explain output: " +
-                   tojson(explainOutput));
-        assert(
-            !hasRejectedPlans(explainOutput),
-            "Expected pipeline " + tojsononeline(pipeline) +
-                " not to have any rejected plans in the explain output: " + tojson(explainOutput));
+               explainOutput);
     } else {
-        assert(isAggregationPlan(explainOutput));
+        assert(isAggregationPlan(explainOutput), explainOutput);
         assert(aggPlanHasStage(explainOutput, "FETCH") || aggPlanHasStage("COLLSCAN"),
-               "Expected pipeline " + tojsononeline(pipeline) +
-                   " to include a FETCH or COLLSCAN stage in the explain output: " +
-                   tojson(explainOutput));
-        assert(
-            !hasRejectedPlans(explainOutput),
-            "Expected pipeline " + tojsononeline(pipeline) +
-                " not to have any rejected plans in the explain output: " + tojson(explainOutput));
+               explainOutput);
     }
-
+    assert(!hasRejectedPlans(explainOutput), explainOutput);
     return explainOutput;
 }
 
@@ -85,12 +62,13 @@ assertQueryCoversProjection(
     {pipeline: [{$match: {x: "string"}}, {$project: {_id: 0, x: 1, a: 1}}]});
 assertQueryCoversProjection(
     {pipeline: [{$match: {x: "string"}}, {$project: {_id: 1, x: 1, a: 1}}]});
-assertQueryCoversProjection(
-    {pipeline: [{$match: {_id: 0, x: "string"}}, {$project: {_id: 1, x: 1, a: 1}}]});
+assertQueryCoversProjection({
+    pipeline: [{$match: {_id: 0, x: "string"}}, {$project: {_id: 1, x: 1, a: 1}}],
+    options: {hint: {x: 1, a: -1, _id: 1}}
+});
 
 // Test that a pipeline requiring a field that is not in the index cannot use a covered plan.
-assertQueryDoesNotCoverProjection(
-    {pipeline: [{$match: {x: "string"}}, {$project: {notThere: 1}}], pipelineOptimizedAway: false});
+assertQueryDoesNotCoverProjection({pipeline: [{$match: {x: "string"}}, {$project: {notThere: 1}}]});
 
 // Test that a covered plan is the only plan considered, even if another plan would be equally
 // selective. Add an equally selective index, then rely on assertQueryCoversProjection() to
@@ -102,17 +80,14 @@ assertQueryCoversProjection({
         {$sort: {x: 1, a: 1}},  // Note: not indexable, but doesn't add any additional dependencies.
         {$project: {_id: 1, x: 1, a: 1}},
     ],
+    options: {hint: {x: 1, a: -1, _id: 1}}
 });
 
 // Test that a multikey index will prevent a covered plan.
 assert.commandWorked(coll.dropIndex({x: 1}));  // Make sure there is only one plan considered.
 assert.commandWorked(coll.insert({x: ["an", "array!"]}));
-assertQueryDoesNotCoverProjection({
-    pipeline: [{$match: {x: "string"}}, {$project: {_id: 1, x: 1}}],
-    pipelineOptimizedAway: false
-});
-assertQueryDoesNotCoverProjection({
-    pipeline: [{$match: {x: "string"}}, {$project: {_id: 1, x: 1, a: 1}}],
-    pipelineOptimizedAway: false
-});
+assertQueryDoesNotCoverProjection(
+    {pipeline: [{$match: {x: "string"}}, {$project: {_id: 1, x: 1}}]});
+assertQueryDoesNotCoverProjection(
+    {pipeline: [{$match: {x: "string"}}, {$project: {_id: 1, x: 1, a: 1}}]});
 }());