SERVER-40830 Move 'ReplaceRootTransformation' to header file

1 files changed, 48 insertions, 0 deletions

diff --git a/src/mongo/db/pipeline/document_source_replace_root.h b/src/mongo/db/pipeline/document_source_replace_root.h
index 0abb4dcaff1..290919615bc 100644
--- a/src/mongo/db/pipeline/document_source_replace_root.h
+++ b/src/mongo/db/pipeline/document_source_replace_root.h
@@ -29,10 +29,58 @@
 
 #pragma once
 
+#include "mongo/db/pipeline/document.h"
 #include "mongo/db/pipeline/document_source_single_document_transformation.h"
+#include "mongo/db/pipeline/expression.h"
 
 namespace mongo {
 
+/**
+ * This class implements the transformation logic for the $replaceRoot and $replaceWith stages.
+ */
+class ReplaceRootTransformation final : public TransformerInterface {
+public:
+    ReplaceRootTransformation(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                              boost::intrusive_ptr<Expression> newRootExpression)
+        : _expCtx(expCtx), _newRoot(std::move(newRootExpression)) {}
+
+    TransformerType getType() const final {
+        return TransformerType::kReplaceRoot;
+    }
+
+    Document applyTransformation(const Document& input) final;
+
+    // Optimize the newRoot expression.
+    void optimize() final {
+        _newRoot->optimize();
+    }
+
+    Document serializeTransformation(
+        boost::optional<ExplainOptions::Verbosity> explain) const final {
+        return Document{{"newRoot", _newRoot->serialize(static_cast<bool>(explain))}};
+    }
+
+    DepsTracker::State addDependencies(DepsTracker* deps) const final {
+        _newRoot->addDependencies(deps);
+        // This stage will replace the entire document with a new document, so any existing fields
+        // will be replaced and cannot be required as dependencies.
+        return DepsTracker::State::EXHAUSTIVE_FIELDS;
+    }
+
+    DocumentSource::GetModPathsReturn getModifiedPaths() const final {
+        // Replaces the entire root, so all paths are modified.
+        return {DocumentSource::GetModPathsReturn::Type::kAllPaths, std::set<std::string>{}, {}};
+    }
+
+    const boost::intrusive_ptr<Expression>& getExpression() const {
+        return _newRoot;
+    }
+
+private:
+    const boost::intrusive_ptr<ExpressionContext> _expCtx;
+    boost::intrusive_ptr<Expression> _newRoot;
+};
+
 /*
  * $replaceRoot takes an object containing only an expression in the newRoot field, and replaces
  * each incoming document with the result of evaluating that expression. Throws an error if the