SERVER-32302 Introduce BSONObjStlIterator as a proper stl-style iterator

(cherry picked from commit 1497899f237dd60ad7313913c38c3f07fe168f2b)

4 files changed, 97 insertions, 28 deletions

diff --git a/src/mongo/bson/bsonelement.h b/src/mongo/bson/bsonelement.h
index 0c84387f2ad..7c427f38004 100644
--- a/src/mongo/bson/bsonelement.h
+++ b/src/mongo/bson/bsonelement.h
@@ -702,6 +702,7 @@ private:
     mutable int totalSize; /* caches the computed size */
 
     friend class BSONObjIterator;
+    friend class BSONObjStlIterator;
     friend class BSONObj;
     const BSONElement& chk(BSONType t) const {
         if (t != type()) {
diff --git a/src/mongo/bson/bsonobj.h b/src/mongo/bson/bsonobj.h
index 6898c63a840..7cd2a50870a 100644
--- a/src/mongo/bson/bsonobj.h
+++ b/src/mongo/bson/bsonobj.h
@@ -52,6 +52,8 @@
 
 namespace mongo {
 
+class BSONObjStlIterator;
+
 /**
    C++ representation of a "BSON" object -- that is, an extended JSON-style
    object in a binary representation.
@@ -516,7 +518,9 @@ public:
     void elems(std::list<BSONElement>&) const;
 
     friend class BSONObjIterator;
-    typedef BSONObjIterator iterator;
+    friend class BSONObjStlIterator;
+    typedef BSONObjStlIterator iterator;
+    typedef BSONObjStlIterator const_iterator;
 
     /**
      * These enable range-based for loops over BSONObjs:
@@ -525,8 +529,8 @@ public:
      *          ... // Do something with elem
      *      }
      */
-    BSONObjIterator begin() const;
-    BSONObjIterator end() const;
+    iterator begin() const;
+    iterator end() const;
 
     void appendSelfToBufBuilder(BufBuilder& b) const {
         verify(objsize());
@@ -580,18 +584,89 @@ struct BSONArray : BSONObj {
     explicit BSONArray(const BSONObj& obj) : BSONObj(obj) {}
 };
 
-/** iterator for a BSONObj
+/**
+ * An stl-compatible forward iterator over the elements of a BSONObj.
+ *
+ * The BSONObj must stay in scope for the duration of the iterator's execution.
+ */
+class BSONObjStlIterator {
+public:
+    using iterator_category = std::forward_iterator_tag;
+    using difference_type = ptrdiff_t;
+    using value_type = BSONElement;
+    using pointer = const BSONElement*;
+    using reference = const BSONElement&;
+
+    /**
+     * All default constructed iterators are equal to each other.
+     * They are in a dereferencable state, and return an EOO BSONElement.
+     * They must not be incremented.
+     */
+    BSONObjStlIterator() = default;
+
+    /**
+     * Constructs an iterator pointing to the first element in obj or EOO if it is empty.
+     */
+    explicit BSONObjStlIterator(const BSONObj& obj) : BSONObjStlIterator(obj.firstElement()) {}
 
-   Note each BSONObj ends with an EOO element: so you will get more() on an empty
-   object, although next().eoo() will be true.
+    /**
+     * Returns an iterator pointing to the EOO element in obj.
+     */
+    static BSONObjStlIterator endOf(const BSONObj& obj) {
+        auto eooElem = BSONElement();
+        eooElem.data = (obj.objdata() + obj.objsize() - 1);
+        dassert(eooElem.eoo());  // This is checked in the BSONObj constructor.
+        return BSONObjStlIterator(eooElem);
+    }
 
-   The BSONObj must stay in scope for the duration of the iterator's execution.
+    /** pre-increment */
+    BSONObjStlIterator& operator++() {
+        dassert(!_cur.eoo());
+        *this = BSONObjStlIterator(BSONElement(_cur.rawdata() + _curSize));
+        return *this;
+    }
 
-   todo: Finish making this an STL-compatible iterator.
-            Need iterator_catagory et al (maybe inherit from std::iterator).
-            Need operator->
-            operator* should return a const reference not a value.
-*/
+    /** post-increment */
+    BSONObjStlIterator operator++(int) {
+        BSONObjStlIterator oldPos = *this;
+        ++*this;
+        return oldPos;
+    }
+
+    const BSONElement& operator*() const {
+        return _cur;
+    }
+    const BSONElement* operator->() const {
+        return &_cur;
+    }
+
+    bool operator==(const BSONObjStlIterator& other) {
+        return _cur.rawdata() == other._cur.rawdata();
+    }
+    bool operator!=(const BSONObjStlIterator& other) {
+        return !(*this == other);
+    }
+
+private:
+    explicit BSONObjStlIterator(BSONElement elem)
+        : _cur(elem),
+          // Make sure the lazy size fields are filled in.
+          _curSize(_cur.size()) {}
+
+    BSONElement _cur;
+    int _curSize = 0;  // TODO consider removing if we make the size cache in BSONElement eager.
+};
+
+/**
+ * Non-STL iterator for a BSONObj
+ *
+ * For simple loops over BSONObj, do this instead: for (auto&& elem : obj) { ... }
+ *
+ * Note each BSONObj ends with an EOO element: so you will get moreWithEOO() on an empty
+ * object, although more() will be false and next().eoo() will be true.
+ *
+ * The BSONObj must stay in scope for the duration of the iterator's execution.
+ */
 class BSONObjIterator {
 public:
     /** Create an iterator for a BSON object.
@@ -611,12 +686,6 @@ public:
         _theend = end - 1;
     }
 
-    static BSONObjIterator endOf(const BSONObj& obj) {
-        BSONObjIterator end(obj);
-        end._pos = end._theend;
-        return end;
-    }
-
     /** @return true if more elements exist to be enumerated. */
     bool more() {
         return _pos < _theend;
@@ -733,11 +802,11 @@ public:
     BSONArrayIteratorSorted(const BSONArray& array);
 };
 
-inline BSONObjIterator BSONObj::begin() const {
-    return BSONObjIterator(*this);
+inline BSONObj::iterator BSONObj::begin() const {
+    return BSONObj::iterator(*this);
 }
-inline BSONObjIterator BSONObj::end() const {
-    return BSONObjIterator::endOf(*this);
+inline BSONObj::iterator BSONObj::end() const {
+    return BSONObj::iterator::endOf(*this);
 }
 
 /**
@@ -780,9 +849,7 @@ template <size_t N>
 inline void BSONObj::getFields(const std::array<StringData, N>& fieldNames,
                                std::array<BSONElement, N>* fields) const {
     std::bitset<N> foundFields;
-    auto iter = this->begin();
-    while (iter.more() && !foundFields.all()) {
-        auto el = iter.next();
+    for (auto&& el : *this) {
         auto fieldName = el.fieldNameStringData();
         for (std::size_t i = 0; i < N; ++i) {
             if (!foundFields.test(i) && (fieldNames[i] == fieldName)) {
@@ -791,6 +858,8 @@ inline void BSONObj::getFields(const std::array<StringData, N>& fieldNames,
                 break;
             }
         }
+        if (foundFields.all())
+            break;
     }
 }
 }  // namespace mongo
diff --git a/src/mongo/bson/mutable/mutable_bson_test.cpp b/src/mongo/bson/mutable/mutable_bson_test.cpp
index 8b832f13f6c..031a09b0128 100644
--- a/src/mongo/bson/mutable/mutable_bson_test.cpp
+++ b/src/mongo/bson/mutable/mutable_bson_test.cpp
@@ -1456,7 +1456,7 @@ TEST(Document, SetValueBSONElementFieldNameHandling) {
 
     static const char inJson2[] = "{ b : 5 }";
     mongo::BSONObj inObj2 = mongo::fromjson(inJson2);
-    mongo::BSONObjIterator iterator = inObj2.begin();
+    mongo::BSONObjIterator iterator(inObj2);
 
     ASSERT_TRUE(iterator.more());
     const mongo::BSONElement b = iterator.next();
diff --git a/src/mongo/bson/util/bson_check.h b/src/mongo/bson/util/bson_check.h
index edfff37b5a0..953279a751d 100644
--- a/src/mongo/bson/util/bson_check.h
+++ b/src/mongo/bson/util/bson_check.h
@@ -49,8 +49,7 @@ Status bsonCheckOnlyHasFieldsImpl(StringData objectName,
                                   const BSONObj& obj,
                                   const Condition& allowed) {
     StringMap<bool> seenFields;
-    for (BSONObj::iterator iter(obj); iter.more();) {
-        const BSONElement e = iter.next();
+    for (auto&& e : obj) {
         const auto name = e.fieldNameStringData();
 
         if (!allowed(name)) {