SERVER-28757 Create BSONElementViewMatchableDocument

1 files changed, 45 insertions, 0 deletions

diff --git a/src/mongo/db/matcher/matchable.h b/src/mongo/db/matcher/matchable.h
index 66d1a417af0..4a9329f858c 100644
--- a/src/mongo/db/matcher/matchable.h
+++ b/src/mongo/db/matcher/matchable.h
@@ -31,6 +31,7 @@
 #pragma once
 
 #include "mongo/bson/bsonobj.h"
+#include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/db/field_ref.h"
 #include "mongo/db/matcher/path.h"
 
@@ -104,4 +105,48 @@ private:
     mutable BSONElementIterator _iterator;
     mutable bool _iteratorUsed;
 };
+
+/**
+ * A MatchableDocument interface for viewing a BSONElement as if it were wrapped in the top-level
+ * field of any given path. For example, given the object obj={a: [5]}, we can create a view over
+ * the element obj["a"]["0"]. An iterator over this view with path "i" would behave identically to
+ * an iterator over {i: 5} with path "i".
+ */
+class BSONElementViewMatchableDocument : public MatchableDocument {
+public:
+    BSONElementViewMatchableDocument(BSONElement elem) : _elem(elem), _iteratorUsed(false) {}
+
+    BSONObj toBSON() const override {
+        return BSON("" << _elem);
+    }
+
+    /**
+     * Creates an iterator over '_elem' as if '_elem' were wrapped in the first field of 'path'.
+     * 'path' must have at least one field.
+     */
+    ElementIterator* allocateIterator(const ElementPath* path) const override {
+        // Skip the first field in the path when traversing the document.
+        const size_t suffixIndex = 1;
+
+        if (_iteratorUsed) {
+            return new BSONElementIterator(path, suffixIndex, _elem);
+        }
+        _iteratorUsed = true;
+        _iterator.reset(path, suffixIndex, _elem);
+        return &_iterator;
+    }
+
+    void releaseIterator(ElementIterator* iterator) const override {
+        if (iterator == &_iterator) {
+            _iteratorUsed = false;
+        } else {
+            delete iterator;
+        }
+    }
+
+private:
+    BSONElement _elem;
+    mutable BSONElementIterator _iterator;
+    mutable bool _iteratorUsed;
+};
 }