SERVER-11692 step 2: use VariablesParseState at parse time

Now we are assigning Ids but still using names to get/set variables at run time.  The main
functional change at this point is that we now error at parse time rather than run time if a user
accesses an undefined variable.

3 files changed, 146 insertions, 53 deletions

diff --git a/src/mongo/db/pipeline/document_source_redact.cpp b/src/mongo/db/pipeline/document_source_redact.cpp
index d30cb58d2db..19bc8d3a61a 100644
--- a/src/mongo/db/pipeline/document_source_redact.cpp
+++ b/src/mongo/db/pipeline/document_source_redact.cpp
@@ -152,6 +152,11 @@ namespace mongo {
         Expression::ObjectCtx oCtx(0);
 
         VariablesParseState vps;
+        // TODO save and use the returned ids somehow
+        vps.defineVariable("CURRENT"); // will differ from ROOT in this DocumentSource
+        vps.defineVariable("DESCEND");
+        vps.defineVariable("PRUNE");
+        vps.defineVariable("KEEP");
         intrusive_ptr<Expression> expression = Expression::parseObject(elem.Obj(), &oCtx, vps);
 
         return new DocumentSourceRedact(expCtx, expression);
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index e8c96b8eedd..3772e6dad9e 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -110,6 +110,27 @@ namespace mongo {
         }
     }
 
+    Variables::Id VariablesParseState::defineVariable(const StringData& name) {
+        // caller should have validated before hand by using Variables::uassertValidNameForUserWrite
+        massert(17275, "Can't redefine ROOT",
+                name != "ROOT");
+
+        Variables::Id id = (*_nextId)++;
+        _variables[name] = id;
+        return id;
+    }
+
+    Variables::Id VariablesParseState::getVariable(const StringData& name) const {
+        StringMap<Variables::Id>::const_iterator it = _variables.find(name);
+        if (it != _variables.end())
+            return it->second;
+
+        uassert(17276, str::stream() << "Use of undefinded variable: " << name,
+                name == "ROOT" || name == "CURRENT");
+
+        return Variables::ROOT_ID;
+    }
+
     /* --------------------------- Expression ------------------------------ */
 
     Expression::ObjectCtx::ObjectCtx(int theOptions)
@@ -1119,9 +1140,9 @@ namespace {
 
     /* --------------------- ExpressionFieldPath --------------------------- */
 
-    // this is the old deprecated version
+    // this is the old deprecated version only used by tests not using variables
     intrusive_ptr<ExpressionFieldPath> ExpressionFieldPath::create(const string& fieldPath) {
-        return new ExpressionFieldPath("CURRENT." + fieldPath);
+        return new ExpressionFieldPath("CURRENT." + fieldPath, Variables::ROOT_ID);
     }
 
     // this is the new version that supports every syntax
@@ -1138,18 +1159,20 @@ namespace {
         if (raw[1] == '$') {
             const StringData rawSD = raw;
             const StringData fieldPath = rawSD.substr(2); // strip off $$
-            const StringData varName = fieldPath.substr(0, fieldPath.find('.')-1);
+            const StringData varName = fieldPath.substr(0, fieldPath.find('.'));
             Variables::uassertValidNameForUserRead(varName);
-            return new ExpressionFieldPath(fieldPath.toString());
+            return new ExpressionFieldPath(fieldPath.toString(), vps.getVariable(varName));
         }
         else {
-            return new ExpressionFieldPath("CURRENT." + raw.substr(1)); // strip the "$" prefix
+            return new ExpressionFieldPath("CURRENT." + raw.substr(1), // strip the "$" prefix
+                                           vps.getVariable("CURRENT"));
         }
     }
 
 
-    ExpressionFieldPath::ExpressionFieldPath(const string& theFieldPath)
+    ExpressionFieldPath::ExpressionFieldPath(const string& theFieldPath, Variables::Id variable)
         : _fieldPath(theFieldPath)
+        , _variable(variable)
         , _baseVar(_fieldPath.getFieldName(0) == "CURRENT" ? CURRENT :
                    _fieldPath.getFieldName(0) == "ROOT" ?    ROOT :
                                                              OTHER)
@@ -1243,7 +1266,7 @@ namespace {
     REGISTER_EXPRESSION("$let", ExpressionLet::parse);
     intrusive_ptr<Expression> ExpressionLet::parse(
             BSONElement expr,
-            const VariablesParseState& vps) {
+            const VariablesParseState& vpsIn) {
 
         verify(str::equals(expr.fieldName(), "$let"));
 
@@ -1251,22 +1274,14 @@ namespace {
                 expr.type() == Object);
         const BSONObj args = expr.embeddedObject();
 
-        // used for input validation
-        bool haveVars = false;
-        bool haveIn = false;
-
-        VariableMap vars;
-        intrusive_ptr<Expression> subExpression;
+        // varsElem must be parsed before inElem regardless of BSON order.
+        BSONElement varsElem;
+        BSONElement inElem;
         BSONForEach(arg, args) {
             if (str::equals(arg.fieldName(), "vars")) {
-                haveVars = true;
-                BSONForEach(variable, arg.embeddedObjectUserCheck()) {
-                    Variables::uassertValidNameForUserWrite(variable.fieldName());
-                    vars[variable.fieldName()] = parseOperand(variable, vps);
-                }
+                varsElem = arg;
             } else if (str::equals(arg.fieldName(), "in")) {
-                haveIn = true;
-                subExpression = parseOperand(arg, vps);
+                inElem = arg;
             } else {
                 uasserted(16875, str::stream()
                         << "Unrecognized parameter to $let: " << arg.fieldName());
@@ -1274,9 +1289,24 @@ namespace {
         }
 
         uassert(16876, "Missing 'vars' parameter to $let",
-                haveVars);
+                !varsElem.eoo());
         uassert(16877, "Missing 'in' parameter to $let",
-                haveIn);
+                !inElem.eoo());
+
+        // parse "vars"
+        VariablesParseState vpsSub(vpsIn); // vpsSub gets our vars, vpsIn doesn't.
+        VariableMap vars;
+        BSONForEach(varElem, varsElem.embeddedObjectUserCheck()) {
+            const string varName = varElem.fieldName();
+            Variables::uassertValidNameForUserWrite(varName);
+            Variables::Id id = vpsSub.defineVariable(varName);
+
+            vars[id] = NameAndExpression(varName,
+                                         parseOperand(varElem, vpsIn)); // only has outer vars
+        }
+
+        // parse "in"
+        intrusive_ptr<Expression> subExpression = parseOperand(inElem, vpsSub); // has our vars
 
         return new ExpressionLet(vars, subExpression);
     }
@@ -1293,7 +1323,7 @@ namespace {
         }
 
         for (VariableMap::iterator it=_variables.begin(), end=_variables.end(); it != end; ++it) {
-            it->second = it->second->optimize();
+            it->second.expression = it->second.expression->optimize();
         }
 
         // TODO be smarter with constant "variables"
@@ -1306,7 +1336,7 @@ namespace {
         MutableDocument vars;
         for (VariableMap::const_iterator it=_variables.begin(), end=_variables.end();
                 it != end; ++it) {
-            vars[it->first] = it->second->serialize(explain);
+            vars[it->second.name] = it->second.expression->serialize(explain);
         }
 
         return Value(DOC("$let" << DOC("vars" << vars.freeze()
@@ -1320,13 +1350,13 @@ namespace {
         for (VariableMap::const_iterator it=_variables.begin(), end=_variables.end();
                 it != end; ++it) {
 
-            const Value newVar = it->second->evaluateInternal(originalVars);
+            const Value newVar = it->second.expression->evaluateInternal(originalVars);
 
             // Can't set ROOT (checked in parse())
-            if (it->first == "CURRENT") {
+            if (it->second.name == "CURRENT") {
                 newVars.current = newVar;
             } else {
-                newRest[it->first] = newVar;
+                newRest[it->second.name] = newVar;
             }
         }
 
@@ -1337,7 +1367,7 @@ namespace {
     void ExpressionLet::addDependencies(set<string>& deps, vector<string>* path) const {
         for (VariableMap::const_iterator it=_variables.begin(), end=_variables.end();
                 it != end; ++it) {
-            it->second->addDependencies(deps);
+            it->second.expression->addDependencies(deps);
         }
 
         // TODO be smarter when CURRENT is a bound variable
@@ -1350,34 +1380,25 @@ namespace {
     REGISTER_EXPRESSION("$map", ExpressionMap::parse);
     intrusive_ptr<Expression> ExpressionMap::parse(
             BSONElement expr,
-            const VariablesParseState& vps) {
+            const VariablesParseState& vpsIn) {
 
         verify(str::equals(expr.fieldName(), "$map"));
 
         uassert(16878, "$map only supports an object as it's argument",
                 expr.type() == Object);
 
-        // used for input validation
-        bool haveInput = false;
-        bool haveAs = false;
-        bool haveIn = false;
-
-        string varName;
-        intrusive_ptr<Expression> input;
-        intrusive_ptr<Expression> in;
-
+        // "in" must be parsed after "as" regardless of BSON order
+        BSONElement inputElem;
+        BSONElement asElem;
+        BSONElement inElem;
         const BSONObj args = expr.embeddedObject();
         BSONForEach(arg, args) {
             if (str::equals(arg.fieldName(), "input")) {
-                haveInput = true;
-                input = parseOperand(arg, vps);
+                inputElem = arg;
             } else if (str::equals(arg.fieldName(), "as")) {
-                haveAs = true;
-                varName = arg.str();
-                Variables::uassertValidNameForUserWrite(varName);
+                asElem = arg;
             } else if (str::equals(arg.fieldName(), "in")) {
-                haveIn = true;
-                in = parseOperand(arg, vps);
+                inElem = arg;
             } else {
                 uasserted(16879, str::stream()
                         << "Unrecognized parameter to $map: " << arg.fieldName());
@@ -1385,19 +1406,33 @@ namespace {
         }
 
         uassert(16880, "Missing 'input' parameter to $map",
-                haveInput);
+                !inputElem.eoo());
         uassert(16881, "Missing 'as' parameter to $map",
-                haveAs);
+                !asElem.eoo());
         uassert(16882, "Missing 'in' parameter to $map",
-                haveIn);
+                !inElem.eoo());
+
+        // parse "input"
+        intrusive_ptr<Expression> input = parseOperand(inputElem, vpsIn); // only has outer vars
+
+        // parse "as"
+        VariablesParseState vpsSub(vpsIn); // vpsSub gets our vars, vpsIn doesn't.
+        string varName = asElem.str();
+        Variables::uassertValidNameForUserWrite(varName);
+        Variables::Id varId = vpsSub.defineVariable(varName);
+
+        // parse "in"
+        intrusive_ptr<Expression> in = parseOperand(inElem, vpsSub); // has access to map variable
 
-        return new ExpressionMap(varName, input, in);
+        return new ExpressionMap(varName, varId, input, in);
     }
 
     ExpressionMap::ExpressionMap(const string& varName,
+                                 Variables::Id varId,
                                  intrusive_ptr<Expression> input,
                                  intrusive_ptr<Expression> each)
         : _varName(varName)
+        , _varId(varId)
         , _input(input)
         , _each(each)
     {}
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 0eb17fde433..ee02c7a945d 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -35,6 +35,7 @@
 #include "mongo/db/pipeline/value.h"
 #include "mongo/util/intrusive_counter.h"
 #include "mongo/util/mongoutils/str.h"
+#include "mongo/util/string_map.h"
 
 namespace mongo {
 
@@ -47,6 +48,11 @@ namespace mongo {
     /// The state used as input to Expressions
     class Variables {
     public:
+        /**
+         * Each unique variable is assigned a unique id of this type
+         */
+        typedef size_t Id;
+
         Variables() {}
 
         explicit Variables(const Document& rootAndCurrent)
@@ -63,12 +69,44 @@ namespace mongo {
         static void uassertValidNameForUserWrite(StringData varName);
         static void uassertValidNameForUserRead(StringData varName);
 
+        static const Id ROOT_ID = Id(-1);
+
         Value root;
         Value current;
         Document rest;
     };
 
+    /**
+     * This class represents the Variables that are defined in an Expression tree.
+     *
+     * All copies from a given instance share enough information to ensure unique Ids are assigned
+     * and to propagate back to the original instance enough information to correctly construct a
+     * Variables instance.
+     */
     class VariablesParseState {
+    public:
+        VariablesParseState() : _nextId(new Variables::Id(0)) {}
+
+        /**
+         * Assigns a named variable a unique Id. This differs from all other variables, even
+         * others with the same name.
+         *
+         * The special variables ROOT and CURRENT are always implicitly defined with CURRENT
+         * equivalent to ROOT. If CURRENT is explicitly defined by a call to this function, it
+         * breaks that equivalence.
+         *
+         * NOTE: Name validation is responsibility of caller.
+         */
+        Variables::Id defineVariable(const StringData& name);
+
+        /**
+         * Returns the current Id for a variable. uasserts if the variable isn't defined.
+         */
+        Variables::Id getVariable(const StringData& name) const;
+
+    private:
+        StringMap<Variables::Id> _variables;
+        shared_ptr<Variables::Id> _nextId;
     };
 
     class Expression :
@@ -478,7 +516,7 @@ namespace mongo {
         const FieldPath& getFieldPath() const { return _fieldPath; }
 
     private:
-        ExpressionFieldPath(const string &fieldPath);
+        ExpressionFieldPath(const string& fieldPath, Variables::Id variable);
 
         /*
           Internal implementation of evaluateInternal(), used recursively.
@@ -506,7 +544,8 @@ namespace mongo {
         };
 
         const FieldPath _fieldPath;
-        const BaseVar _baseVar;
+        const Variables::Id _variable;
+        const BaseVar _baseVar; // TODO remove
     };
 
 
@@ -538,10 +577,22 @@ namespace mongo {
             BSONElement expr,
             const VariablesParseState& vps);
 
-        typedef map<string, intrusive_ptr<Expression> > VariableMap;
+        struct NameAndExpression {
+            NameAndExpression() {}
+            NameAndExpression(string name, intrusive_ptr<Expression> expression)
+                : name(name)
+                , expression(expression)
+            {}
+
+            string name;
+            intrusive_ptr<Expression> expression;
+        };
+
+        typedef map<Variables::Id, NameAndExpression> VariableMap;
 
     private:
-        ExpressionLet(const VariableMap& vars, intrusive_ptr<Expression> subExpression);
+        ExpressionLet(const VariableMap& vars,
+                      intrusive_ptr<Expression> subExpression);
 
         VariableMap _variables;
         intrusive_ptr<Expression> _subExpression;
@@ -561,10 +612,12 @@ namespace mongo {
 
     private:
         ExpressionMap(const string& varName, // name of variable to set
+                      Variables::Id varId, // id of variable to set
                       intrusive_ptr<Expression> input, // yields array to iterate
                       intrusive_ptr<Expression> each); // yields results to be added to output array
 
         string _varName;
+        Variables::Id _varId;
         intrusive_ptr<Expression> _input;
         intrusive_ptr<Expression> _each;
     };