1 files changed, 98 insertions, 17 deletions

diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 9cce6d0b1e2..4ad28bf3d04 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -49,6 +49,7 @@
 #include "mongo/db/pipeline/expression_context.h"
 #include "mongo/db/pipeline/expression_visitor.h"
 #include "mongo/db/pipeline/field_path.h"
+#include "mongo/db/pipeline/monotonic_expression.h"
 #include "mongo/db/pipeline/variables.h"
 #include "mongo/db/query/allowed_contexts.h"
 #include "mongo/db/query/datetime/date_time_support.h"
@@ -170,7 +171,7 @@ public:
      */
     struct ComputedPaths {
         // Non-rename computed paths.
-        std::set<std::string> paths;
+        OrderedPathSet paths;
 
         // Mappings from the old name of a path before applying this expression, to the new one
         // after applying this expression.
@@ -339,6 +340,18 @@ public:
         return _expCtx;
     }
 
+    boost::optional<Variables::Id> getBoundaryVariableId() const {
+        return _boundaryVariableId;
+    }
+
+    bool isMonotonic(const FieldPath& sortedFieldPath) const {
+        return getMonotonicState(sortedFieldPath) != monotonic::State::NonMonotonic;
+    }
+
+    virtual monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const {
+        return monotonic::State::NonMonotonic;
+    }
+
 protected:
     using ExpressionVector = std::vector<boost::intrusive_ptr<Expression>>;
 
@@ -737,6 +750,10 @@ protected:
     void _doAddDependencies(DepsTracker* deps) const override;
 
 private:
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final {
+        return monotonic::State::Constant;
+    }
+
     Value _value;
 };
 
@@ -943,6 +960,11 @@ public:
     void acceptVisitor(ExpressionConstVisitor* visitor) const final {
         return visitor->visit(this);
     }
+
+private:
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final {
+        return monotonic::combineExpressions(sortedFieldPath, getChildren());
+    };
 };
 
 
@@ -1188,6 +1210,11 @@ public:
     void acceptVisitor(ExpressionConstVisitor* visitor) const final {
         return visitor->visit(this);
     }
+
+private:
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final {
+        return getChildren()[0]->getMonotonicState(sortedFieldPath);
+    }
 };
 
 
@@ -1662,6 +1689,8 @@ private:
 
     void _doAddDependencies(DepsTracker* deps) const final;
 
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final;
+
     // Starting time instant expression. Accepted types: Date_t, Timestamp, OID.
     boost::intrusive_ptr<Expression>& _startDate;
 
@@ -1832,6 +1861,8 @@ protected:
 
 
 private:
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final;
+
     /*
       Internal implementation of evaluate(), used recursively.
 
@@ -1923,6 +1954,11 @@ public:
     void acceptVisitor(ExpressionConstVisitor* visitor) const final {
         return visitor->visit(this);
     }
+
+private:
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final {
+        return getChildren()[0]->getMonotonicState(sortedFieldPath);
+    }
 };
 
 
@@ -2860,7 +2896,10 @@ public:
     }
 
     bool isCommutative() const final {
-        return true;
+        // Only commutative when performing binary string comparison. The first value entered when
+        // multiple collation-equal but binary-unequal values are added will dictate what is stored
+        // in the set.
+        return getExpressionContext()->getCollator() == nullptr;
     }
 
     void acceptVisitor(ExpressionMutableVisitor* visitor) final {
@@ -3276,6 +3315,9 @@ public:
     void acceptVisitor(ExpressionConstVisitor* visitor) const final {
         return visitor->visit(this);
     }
+
+private:
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final;
 };
 
 
@@ -3285,11 +3327,10 @@ public:
         std::pair<boost::intrusive_ptr<Expression>&, boost::intrusive_ptr<Expression>&>;
 
     ExpressionSwitch(ExpressionContext* const expCtx,
-                     std::vector<boost::intrusive_ptr<Expression>> children,
-                     std::vector<ExpressionPair> branches)
-        : Expression(expCtx, std::move(children)),
-          _default(_children.back()),
-          _branches(std::move(branches)) {}
+                     std::vector<boost::intrusive_ptr<Expression>> children)
+        : Expression(expCtx, std::move(children)) {
+        uassert(40068, "$switch requires at least one branch", numBranches() >= 1);
+    }
 
     Value evaluate(const Document& root, Variables* variables) const final;
     boost::intrusive_ptr<Expression> optimize() final;
@@ -3306,12 +3347,38 @@ public:
         return visitor->visit(this);
     }
 
+    /**
+     * Returns the number of cases in the switch expression. Each branch is made up of two
+     * expressions ('case' and 'then').
+     */
+    int numBranches() const {
+        return _children.size() / 2;
+    }
+
+    /**
+     * Returns a pair of expression pointers representing the 'case' and 'then' expressions for the
+     * i-th branch of the switch.
+     */
+    std::pair<const Expression*, const Expression*> getBranch(int i) const {
+        invariant(i >= 0);
+        invariant(i < numBranches());
+        return {_children[i * 2].get(), _children[i * 2 + 1].get()};
+    }
+
+    /**
+     * Returns the 'default' expression, or nullptr if there is no 'default'.
+     */
+    const Expression* defaultExpr() const {
+        return _children.back().get();
+    }
+
 protected:
     void _doAddDependencies(DepsTracker* deps) const final;
 
 private:
-    boost::intrusive_ptr<Expression>& _default;
-    std::vector<ExpressionPair> _branches;
+    // Helper for 'optimize()'. Deletes the 'case' and 'then' children associated with the i-th
+    // branch of the switch.
+    void deleteBranch(int i);
 };
 
 
@@ -3982,6 +4049,10 @@ protected:
                                           long long amount,
                                           const TimeZone& timezone) const = 0;
 
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final;
+    virtual monotonic::State combineMonotonicStateOfArguments(
+        monotonic::State startDataMonotonicState, monotonic::State amountMonotonicState) const = 0;
+
 private:
     // The expression representing the startDate argument.
     boost::intrusive_ptr<Expression>& _startDate;
@@ -4016,10 +4087,14 @@ public:
     }
 
 private:
-    virtual Value evaluateDateArithmetics(Date_t date,
-                                          TimeUnit unit,
-                                          long long amount,
-                                          const TimeZone& timezone) const override;
+    monotonic::State combineMonotonicStateOfArguments(
+        monotonic::State startDataMonotonicState,
+        monotonic::State amountMonotonicState) const final;
+
+    Value evaluateDateArithmetics(Date_t date,
+                                  TimeUnit unit,
+                                  long long amount,
+                                  const TimeZone& timezone) const final;
 };
 
 class ExpressionDateSubtract final : public ExpressionDateArithmetics {
@@ -4039,10 +4114,14 @@ public:
     }
 
 private:
-    virtual Value evaluateDateArithmetics(Date_t date,
-                                          TimeUnit unit,
-                                          long long amount,
-                                          const TimeZone& timezone) const override;
+    monotonic::State combineMonotonicStateOfArguments(
+        monotonic::State startDataMonotonicState,
+        monotonic::State amountMonotonicState) const final;
+
+    Value evaluateDateArithmetics(Date_t date,
+                                  TimeUnit unit,
+                                  long long amount,
+                                  const TimeZone& timezone) const final;
 };
 
 struct SubstituteFieldPathWalker {
@@ -4128,6 +4207,8 @@ private:
 
     void _doAddDependencies(DepsTracker* deps) const final;
 
+    monotonic::State getMonotonicState(const FieldPath& sortedFieldPath) const final;
+
     // Expression that evaluates to a date to truncate. Accepted BSON types: Date, bsonTimestamp,
     // jstOID.
     boost::intrusive_ptr<Expression>& _date;