SERVER-42988 allow agg-style object syntax in find() projection

1 files changed, 21 insertions, 1 deletions

diff --git a/src/mongo/db/query/projection_ast.h b/src/mongo/db/query/projection_ast.h
index 8a279fbbf8c..3f85c85c419 100644
--- a/src/mongo/db/query/projection_ast.h
+++ b/src/mongo/db/query/projection_ast.h
@@ -74,11 +74,16 @@ public:
         return _children;
     }
 
+    ASTNode* child(size_t index) const {
+        invariant(index < _children.size());
+        return _children[index].get();
+    }
+
     const ASTNode* parent() const {
         return _parent;
     }
 
-    const bool isRoot() const {
+    bool isRoot() const {
         return !_parent;
     }
 
@@ -156,6 +161,21 @@ public:
         _fieldNames.push_back(fieldName.toString());
     }
 
+    /**
+     * Remove a node which is a direct child of this tree. Returns true if anything was removed,
+     * false otherwise.
+     */
+    bool removeChild(StringData fieldName) {
+        if (auto it = std::find(_fieldNames.begin(), _fieldNames.end(), fieldName);
+            it != _fieldNames.end()) {
+            _children.erase(_children.begin() + std::distance(_fieldNames.begin(), it));
+            _fieldNames.erase(it);
+            return true;
+        }
+
+        return false;
+    }
+
     const std::vector<std::string>& fieldNames() const {
         return _fieldNames;
     }