SERVER-40826 Add getters to projection nodes

6 files changed, 61 insertions, 0 deletions

diff --git a/src/mongo/db/field_ref.cpp b/src/mongo/db/field_ref.cpp
index 35bc6a25a5c..c919a5df7c9 100644
--- a/src/mongo/db/field_ref.cpp
+++ b/src/mongo/db/field_ref.cpp
@@ -139,6 +139,16 @@ void FieldRef::removeLastPart() {
     _size--;
 }
 
+void FieldRef::removeFirstPart() {
+    if (_size == 0) {
+        return;
+    }
+    for (size_t i = 0; i + 1 < _size; ++i) {
+        setPart(i, getPart(i + 1));
+    }
+    removeLastPart();
+}
+
 size_t FieldRef::appendParsedPart(FieldRef::StringView part) {
     if (_size < kReserveAhead) {
         _fixed[_size] = part;
diff --git a/src/mongo/db/field_ref.h b/src/mongo/db/field_ref.h
index 736f843ec34..2fe3a7f375f 100644
--- a/src/mongo/db/field_ref.h
+++ b/src/mongo/db/field_ref.h
@@ -115,6 +115,12 @@ public:
     void removeLastPart();
 
     /**
+     * Removes the first part from the path, decreasing its size by 1. Has no effect on a FielRef
+     * with size 0.
+     */
+    void removeFirstPart();
+
+    /**
      * Returns the 'i-th' field part. Assumes i < size(). Behavior is undefined otherwise.
      */
     StringData getPart(size_t i) const;
diff --git a/src/mongo/db/field_ref_test.cpp b/src/mongo/db/field_ref_test.cpp
index f01020cbe0a..9a116ccf98b 100644
--- a/src/mongo/db/field_ref_test.cpp
+++ b/src/mongo/db/field_ref_test.cpp
@@ -872,5 +872,38 @@ TEST(NumericPathComponents, FieldsWithLeadingZeroesAreNotConsideredNumeric) {
     ASSERT(numericPathComponents == expectedComponents);
 }
 
+TEST(RemoveFirstPart, EmptyPathDoesNothing) {
+    FieldRef path;
+    path.removeFirstPart();
+    ASSERT_EQ(path.numParts(), 0U);
+}
+
+TEST(RemoveFirstPart, PathWithOneComponentBecomesEmpty) {
+    FieldRef path("first");
+    path.removeFirstPart();
+    ASSERT_EQ(path.numParts(), 0U);
+}
+
+TEST(RemoveFirstPart, PathWithTwoComponentsOnlyHoldsSecond) {
+    FieldRef path("remove.keep");
+    path.removeFirstPart();
+    ASSERT_EQ(path.numParts(), 1U);
+    ASSERT_EQ(path, FieldRef("keep"));
+}
+
+TEST(RemoveFirstPart, RemovingFirstPartFromLongPathMultipleTimes) {
+    FieldRef path("first.second.third.fourth.fifth.sixth.seventh.eigth.ninth.tenth");
+    path.removeFirstPart();
+    ASSERT_EQ(path, FieldRef("second.third.fourth.fifth.sixth.seventh.eigth.ninth.tenth"));
+    path.removeFirstPart();
+    ASSERT_EQ(path, FieldRef("third.fourth.fifth.sixth.seventh.eigth.ninth.tenth"));
+    path.removeFirstPart();
+    ASSERT_EQ(path, FieldRef("fourth.fifth.sixth.seventh.eigth.ninth.tenth"));
+    path.removeFirstPart();
+    ASSERT_EQ(path, FieldRef("fifth.sixth.seventh.eigth.ninth.tenth"));
+    path.removeFirstPart();
+    ASSERT_EQ(path, FieldRef("sixth.seventh.eigth.ninth.tenth"));
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_single_document_transformation.h b/src/mongo/db/pipeline/document_source_single_document_transformation.h
index 44a624013e3..f36287b0d37 100644
--- a/src/mongo/db/pipeline/document_source_single_document_transformation.h
+++ b/src/mongo/db/pipeline/document_source_single_document_transformation.h
@@ -81,6 +81,10 @@ public:
         return _parsedTransform->getType();
     }
 
+    const auto& getTransformer() const {
+        return *_parsedTransform;
+    }
+
     bool isSubsetOfProjection(const BSONObj& proj) const {
         return _parsedTransform->isSubsetOfProjection(proj);
     }
diff --git a/src/mongo/db/pipeline/parsed_exclusion_projection.h b/src/mongo/db/pipeline/parsed_exclusion_projection.h
index 4d0fb1c70cd..1e381877bc5 100644
--- a/src/mongo/db/pipeline/parsed_exclusion_projection.h
+++ b/src/mongo/db/pipeline/parsed_exclusion_projection.h
@@ -104,6 +104,10 @@ public:
         return TransformerType::kExclusionProjection;
     }
 
+    const ExclusionNode& getRoot() const {
+        return *_root;
+    }
+
     Document serializeTransformation(
         boost::optional<ExplainOptions::Verbosity> explain) const final;
 
diff --git a/src/mongo/db/pipeline/parsed_inclusion_projection.h b/src/mongo/db/pipeline/parsed_inclusion_projection.h
index d913bb2ab05..b1bce15c6fe 100644
--- a/src/mongo/db/pipeline/parsed_inclusion_projection.h
+++ b/src/mongo/db/pipeline/parsed_inclusion_projection.h
@@ -98,6 +98,10 @@ public:
         return TransformerType::kInclusionProjection;
     }
 
+    const InclusionNode& getRoot() const {
+        return *_root;
+    }
+
     /**
      * Parses the projection specification given by 'spec', populating internal data structures.
      */