SERVER-43396 use path tracking visitor when producing AST debug BSON

2 files changed, 35 insertions, 67 deletions

diff --git a/src/mongo/db/query/projection_ast_path_tracking_visitor.h b/src/mongo/db/query/projection_ast_path_tracking_visitor.h
index 2d989b88ae6..3909b8ad892 100644
--- a/src/mongo/db/query/projection_ast_path_tracking_visitor.h
+++ b/src/mongo/db/query/projection_ast_path_tracking_visitor.h
@@ -66,6 +66,10 @@ public:
         return _basePath;
     }
 
+    const std::string& childPath() const {
+        return _fieldNames.top().front();
+    }
+
     void setBasePath(boost::optional<FieldPath> path) {
         _basePath = std::move(path);
     }
@@ -196,6 +200,16 @@ private:
  *
  * The visitors specified in the 'preVisitors' and 'postVisitors' parameters will be visited in
  * the same order as they were added to the vector.
+ *
+ * A note about the order the visitors are executed in: There are two special "path tracking"
+ * visitors, which are responsible for maintaining state about which fields are being visited. The
+ * path tracking pre-visitor is responsible for setting up the field name state for child nodes.
+ * That is, when a ProjectionPathASTNode is encountered, the pre-visitor sets up the field name
+ * state for the ProjectionPathASTNode's first child. After a leaf node has been visited, the
+ * post-visitor is responsible for setting up the path state for its next sibling (if there is
+ * one). This means that the order the path tracking visitors are executed in is significant: They
+ * must both be run after the user-specified visitors.
+ *
  */
 template <class UserData = PathTrackingDummyDefaultType, bool IsConst = true>
 class PathTrackingWalker final {
@@ -207,7 +221,7 @@ public:
           _pathTrackingPostVisitor{context},
           _preVisitors{std::move(preVisitors)},
           _postVisitors{std::move(postVisitors)} {
-        _preVisitors.insert(_preVisitors.begin(), &_pathTrackingPreVisitor);
+        _preVisitors.push_back(&_pathTrackingPreVisitor);
         _postVisitors.push_back(&_pathTrackingPostVisitor);
     }
 
diff --git a/src/mongo/db/query/projection_ast_util.cpp b/src/mongo/db/query/projection_ast_util.cpp
index ce346c5b98e..ec2ec49560a 100644
--- a/src/mongo/db/query/projection_ast_util.cpp
+++ b/src/mongo/db/query/projection_ast_util.cpp
@@ -31,54 +31,42 @@
 
 #include "mongo/db/query/projection_ast_util.h"
 
+#include "mongo/db/query/projection_ast_path_tracking_visitor.h"
 #include "mongo/db/query/projection_ast_walker.h"
 
 namespace mongo::projection_ast {
 namespace {
 struct BSONVisitorContext {
-    std::stack<std::list<std::string>> fieldNames;
     std::stack<BSONObjBuilder> builders;
-
-    BSONObjBuilder& builder() {
-        return builders.top();
-    }
 };
 
 class BSONPreVisitor : public ProjectionASTConstVisitor {
 public:
-    BSONPreVisitor(BSONVisitorContext* context) : _context(context) {}
+    BSONPreVisitor(PathTrackingVisitorContext<BSONVisitorContext>* context)
+        : _context(context), _builders(context->data().builders) {}
 
     virtual void visit(const MatchExpressionASTNode* node) {
         static_cast<const MatchExpressionASTNode*>(node)->matchExpression()->serialize(
-            &_context->builder(), true);
-        _context->fieldNames.top().pop_front();
+            &_builders.top(), true);
     }
 
     virtual void visit(const ProjectionPathASTNode* node) {
         if (!node->parent()) {
             // No root of the tree, thus this node has no field name.
-            _context->builders.push(BSONObjBuilder());
+            _builders.push(BSONObjBuilder());
         } else {
-            _context->builders.push(_context->builder().subobjStart(getFieldName()));
+            _builders.push(_builders.top().subobjStart(getFieldName()));
         }
-
-        // Push all of the field names onto a new layer on the field name stack.
-        _context->fieldNames.push(
-            std::list<std::string>(node->fieldNames().begin(), node->fieldNames().end()));
     }
 
     virtual void visit(const ProjectionPositionalASTNode* node) {
         // ProjectionPositional always has the original query's match expression node as its
         // child. Serialize as: {"positional.projection.field.$": <original match expression>}.
-        _context->builders.push(_context->builder().subobjStart(getFieldName() + ".$"));
-        // Since match expressions serialize their own field name, this is not actually used. It's
-        // pushed since when a MatchExpressionASTNode is visited it expects a field name to have
-        // been put on the stack (just like every other node), and will pop it.
-        _context->fieldNames.push({"<dummy>"});
+        _context->data().builders.push(_builders.top().subobjStart(getFieldName() + ".$"));
     }
 
     virtual void visit(const ProjectionSliceASTNode* node) {
-        BSONObjBuilder sub(_context->builder().subobjStart(getFieldName()));
+        BSONObjBuilder sub(_builders.top().subobjStart(getFieldName()));
         if (node->skip()) {
             sub.appendArray("$slice", BSON_ARRAY(*node->skip() << node->limit()));
         } else {
@@ -91,23 +79,20 @@ public:
     }
 
     virtual void visit(const ExpressionASTNode* node) {
-        node->expression()->serialize(false).addToBsonObj(&_context->builder(), getFieldName());
+        node->expression()->serialize(false).addToBsonObj(&_builders.top(), getFieldName());
     }
 
     virtual void visit(const BooleanConstantASTNode* node) {
-        _context->builders.top().append(getFieldName(), node->value());
+        _builders.top().append(getFieldName(), node->value());
     }
 
 private:
     std::string getFieldName() {
-        invariant(!_context->fieldNames.empty());
-        invariant(!_context->fieldNames.top().empty());
-        auto ret = _context->fieldNames.top().front();
-        _context->fieldNames.top().pop_front();
-        return ret;
+        return _context->childPath();
     }
 
-    BSONVisitorContext* _context;
+    PathTrackingVisitorContext<BSONVisitorContext>* _context;
+    std::stack<BSONObjBuilder>& _builders;
 };
 
 class BSONPostVisitor : public ProjectionASTConstVisitor {
@@ -120,15 +105,10 @@ public:
             // Pop the BSONObjBuilder that was added in the pre visitor.
             _context->builders.pop();
         }
-
-        // Make sure all of the children were serialized.
-        invariant(_context->fieldNames.top().empty());
-        _context->fieldNames.pop();
     }
 
     virtual void visit(const ProjectionPositionalASTNode* node) {
         _context->builders.pop();
-        _context->fieldNames.pop();
     }
 
     virtual void visit(const MatchExpressionASTNode* node) {}
@@ -140,43 +120,17 @@ public:
 private:
     BSONVisitorContext* _context;
 };
-
-class BSONWalker {
-public:
-    BSONWalker() : _preVisitor(&_context), _postVisitor(&_context) {}
-
-    void preVisit(const ASTNode* node) {
-        node->acceptVisitor(&_preVisitor);
-    }
-
-    void postVisit(const ASTNode* node) {
-        node->acceptVisitor(&_postVisitor);
-    }
-
-    void inVisit(long count, const ASTNode* node) {
-        // No op.
-    }
-
-    BSONObjBuilder done() {
-        invariant(_context.fieldNames.empty());
-        invariant(_context.builders.size() == 1);
-
-        auto ret = std::move(_context.builders.top());
-        return ret;
-    }
-
-private:
-    BSONVisitorContext _context;
-    BSONPreVisitor _preVisitor;
-    BSONPostVisitor _postVisitor;
-};
 }  // namespace
 
 BSONObj astToDebugBSON(const ASTNode* root) {
-    BSONWalker walker;
+    PathTrackingVisitorContext<BSONVisitorContext> context;
+    BSONPreVisitor preVisitor{&context};
+    BSONPostVisitor postVisitor{&context.data()};
+    PathTrackingWalker walker{&context, {&preVisitor}, {&postVisitor}};
+
     projection_ast_walker::walk(&walker, root);
 
-    BSONObjBuilder bob = walker.done();
-    return bob.obj();
+    invariant(context.data().builders.size() == 1);
+    return context.data().builders.top().obj();
 }
 }  // namespace mongo::projection_ast