SERVER-31962 Replace Variables::hasUserDefinedValue() with const val chk

6 files changed, 50 insertions, 22 deletions

diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index 9d909106423..bf8a890924f 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -600,7 +600,7 @@ void DocumentSourceLookUp::resolveLetVariables(const Document& localDoc, Variabl
 
     for (auto& letVar : _letVariables) {
         auto value = letVar.expression->evaluate(localDoc);
-        variables->setValue(letVar.id, value);
+        variables->setConstantValue(letVar.id, value);
     }
 }
 
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index c2c96cf6b35..1a51ec5ae94 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -1819,8 +1819,7 @@ intrusive_ptr<Expression> ExpressionFieldPath::optimize() {
         return ExpressionConstant::create(getExpressionContext(), Value());
     }
 
-    if (Variables::isUserDefinedVariable(_variable) &&
-        getExpressionContext()->variables.hasUserDefinedValue(_variable)) {
+    if (getExpressionContext()->variables.hasConstantValue(_variable)) {
         return ExpressionConstant::create(getExpressionContext(), evaluate(Document()));
     }
 
diff --git a/src/mongo/db/pipeline/expression_test.cpp b/src/mongo/db/pipeline/expression_test.cpp
index 6ada15ea651..b893229d6f0 100644
--- a/src/mongo/db/pipeline/expression_test.cpp
+++ b/src/mongo/db/pipeline/expression_test.cpp
@@ -2228,7 +2228,7 @@ TEST(FieldPath, NoOptimizationOnNormalPath) {
 TEST(FieldPath, OptimizeOnVariableWithConstantScalarValue) {
     intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
     auto varId = expCtx->variablesParseState.defineVariable("userVar");
-    expCtx->variables.setValue(varId, Value(123));
+    expCtx->variables.setConstantValue(varId, Value(123));
 
     auto expr = ExpressionFieldPath::parse(expCtx, "$$userVar", expCtx->variablesParseState);
     ASSERT_TRUE(dynamic_cast<ExpressionFieldPath*>(expr.get()));
@@ -2240,7 +2240,7 @@ TEST(FieldPath, OptimizeOnVariableWithConstantScalarValue) {
 TEST(FieldPath, OptimizeOnVariableWithConstantArrayValue) {
     intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
     auto varId = expCtx->variablesParseState.defineVariable("userVar");
-    expCtx->variables.setValue(varId, Value(BSON_ARRAY(1 << 2 << 3)));
+    expCtx->variables.setConstantValue(varId, Value(BSON_ARRAY(1 << 2 << 3)));
 
     auto expr = ExpressionFieldPath::parse(expCtx, "$$userVar", expCtx->variablesParseState);
     ASSERT_TRUE(dynamic_cast<ExpressionFieldPath*>(expr.get()));
@@ -2254,7 +2254,7 @@ TEST(FieldPath, OptimizeOnVariableWithConstantArrayValue) {
 TEST(FieldPath, OptimizeToEmptyArrayOnNumericalPathComponentAndConstantArrayValue) {
     intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
     auto varId = expCtx->variablesParseState.defineVariable("userVar");
-    expCtx->variables.setValue(varId, Value(BSON_ARRAY(1 << 2 << 3)));
+    expCtx->variables.setConstantValue(varId, Value(BSON_ARRAY(1 << 2 << 3)));
 
     auto expr = ExpressionFieldPath::parse(expCtx, "$$userVar.1", expCtx->variablesParseState);
     ASSERT_TRUE(dynamic_cast<ExpressionFieldPath*>(expr.get()));
@@ -2268,7 +2268,7 @@ TEST(FieldPath, OptimizeToEmptyArrayOnNumericalPathComponentAndConstantArrayValu
 TEST(FieldPath, OptimizeOnVariableWithConstantValueAndDottedPath) {
     intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
     auto varId = expCtx->variablesParseState.defineVariable("userVar");
-    expCtx->variables.setValue(varId, Value(Document{{"x", Document{{"y", 123}}}}));
+    expCtx->variables.setConstantValue(varId, Value(Document{{"x", Document{{"y", 123}}}}));
 
     auto expr = ExpressionFieldPath::parse(expCtx, "$$userVar.x.y", expCtx->variablesParseState);
     ASSERT_TRUE(dynamic_cast<ExpressionFieldPath*>(expr.get()));
@@ -2305,7 +2305,7 @@ TEST(FieldPath, NoOptimizationOnVariableWithMissingValue) {
 TEST(FieldPath, ScalarVariableWithDottedFieldPathOptimizesToConstantMissingValue) {
     intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
     auto varId = expCtx->variablesParseState.defineVariable("userVar");
-    expCtx->variables.setValue(varId, Value(123));
+    expCtx->variables.setConstantValue(varId, Value(123));
 
     auto expr = ExpressionFieldPath::parse(expCtx, "$$userVar.x.y", expCtx->variablesParseState);
     ASSERT_TRUE(dynamic_cast<ExpressionFieldPath*>(expr.get()));
diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp
index 31e1dd19ec0..20c6fba7576 100644
--- a/src/mongo/db/pipeline/pipeline_test.cpp
+++ b/src/mongo/db/pipeline/pipeline_test.cpp
@@ -624,10 +624,10 @@ TEST(PipelineOptimizationTest, LookupDoesNotAbsorbUnwindOnSubfieldOfAsButStillMo
 TEST(PipelineOptimizationTest, MatchShouldDuplicateItselfBeforeRedact) {
     string inputPipe = "[{$redact: '$$PRUNE'}, {$match: {a: 1, b:12}}]";
     string outputPipe =
-        "[{$match: {$and: [{a: {$eq: 1}}, {b: {$eq: 12}}]}}, {$redact: {$const: 'prune'}}, "
+        "[{$match: {$and: [{a: {$eq: 1}}, {b: {$eq: 12}}]}}, {$redact: '$$PRUNE'}, "
         "{$match: {$and: [{a: {$eq: 1}}, {b: {$eq: 12}}]}}]";
     string serializedPipe =
-        "[{$match: {a: 1, b: 12}}, {$redact: {$const: 'prune'}}, {$match: {a: 1, b: 12}}]";
+        "[{$match: {a: 1, b: 12}}, {$redact: '$$PRUNE'}, {$match: {a: 1, b: 12}}]";
     assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe);
 }
 
diff --git a/src/mongo/db/pipeline/variables.cpp b/src/mongo/db/pipeline/variables.cpp
index fc7983dbf5e..2c469cb39b4 100644
--- a/src/mongo/db/pipeline/variables.cpp
+++ b/src/mongo/db/pipeline/variables.cpp
@@ -98,15 +98,29 @@ void Variables::uassertValidNameForUserRead(StringData varName) {
     }
 }
 
-void Variables::setValue(Id id, const Value& value) {
+void Variables::setValue(Id id, const Value& value, bool isConstant) {
     uassert(17199, "can't use Variables::setValue to set a reserved builtin variable", id >= 0);
 
     const auto idAsSizeT = static_cast<size_t>(id);
     if (idAsSizeT >= _valueList.size()) {
         _valueList.resize(idAsSizeT + 1);
+    } else {
+        // If a value has already been set for 'id', and that value was marked as constant, then it
+        // is illegal to modify.
+        invariant(!_valueList[idAsSizeT].isConstant);
     }
 
-    _valueList[id] = value;
+    _valueList[idAsSizeT] = ValueAndState(value, isConstant);
+}
+
+void Variables::setValue(Variables::Id id, const Value& value) {
+    const bool isConstant = false;
+    setValue(id, value, isConstant);
+}
+
+void Variables::setConstantValue(Variables::Id id, const Value& value) {
+    const bool isConstant = true;
+    setValue(id, value, isConstant);
 }
 
 Value Variables::getUserDefinedValue(Variables::Id id) const {
@@ -115,7 +129,7 @@ Value Variables::getUserDefinedValue(Variables::Id id) const {
     uassert(40434,
             str::stream() << "Requesting Variables::getValue with an out of range id: " << id,
             static_cast<size_t>(id) < _valueList.size());
-    return _valueList[id];
+    return _valueList[id].value;
 }
 
 Value Variables::getValue(Id id, const Document& root) const {
diff --git a/src/mongo/db/pipeline/variables.h b/src/mongo/db/pipeline/variables.h
index d3fcc298d1c..703e013e9b0 100644
--- a/src/mongo/db/pipeline/variables.h
+++ b/src/mongo/db/pipeline/variables.h
@@ -38,7 +38,7 @@ namespace mongo {
 /**
  * The state used as input and working space for Expressions.
  */
-class Variables {
+class Variables final {
 public:
     // Each unique variable is assigned a unique id of this type. Negative ids are reserved for
     // system variables and non-negative ids are allocated for user variables.
@@ -83,6 +83,12 @@ public:
     void setValue(Variables::Id id, const Value& value);
 
     /**
+     * Same as 'setValue' but marks 'value' as being constant. It is illegal to change a value that
+     * has been marked constant.
+     */
+    void setConstantValue(Variables::Id id, const Value& value);
+
+    /**
      * Gets the value of a user-defined or system variable. If the 'id' provided represents the
      * special ROOT variable, then we return 'root' in Value form.
      */
@@ -95,13 +101,11 @@ public:
     Value getUserDefinedValue(Variables::Id id) const;
 
     /**
-     * Returns whether a value for 'id' has been stored in this Variables instance.
-     * TODO: This method does not distinguish between missing entries in _valueList and entries that
-     * have been explicitly set to missing.
+     * Returns whether a constant value for 'id' has been defined using setConstantValue().
      */
-    bool hasUserDefinedValue(Variables::Id id) const {
-        invariant(isUserDefinedVariable(id));
-        return _valueList.size() > static_cast<size_t>(id) && !getUserDefinedValue(id).missing();
+    bool hasConstantValue(Variables::Id id) const {
+
+        return _valueList.size() > static_cast<size_t>(id) && _valueList[id].isConstant;
     }
 
     /**
@@ -115,8 +119,19 @@ public:
     }
 
 private:
+    struct ValueAndState {
+        ValueAndState() = default;
+
+        ValueAndState(Value val, bool isConst) : value(std::move(val)), isConstant(isConst) {}
+
+        Value value;
+        bool isConstant = false;
+    };
+
+    void setValue(Id id, const Value& value, bool isConstant);
+
     IdGenerator _idGenerator;
-    std::vector<Value> _valueList;
+    std::vector<ValueAndState> _valueList;
 };
 
 /**
@@ -126,7 +141,7 @@ private:
  * and to propagate back to the original instance enough information to correctly construct a
  * Variables instance.
  */
-class VariablesParseState {
+class VariablesParseState final {
 public:
     explicit VariablesParseState(Variables::IdGenerator* variableIdGenerator)
         : _idGenerator(variableIdGenerator) {}