SERVER-32302 Compute BSONElement sizes eagerly

4 files changed, 32 insertions, 155 deletions

diff --git a/src/mongo/bson/bsonelement.cpp b/src/mongo/bson/bsonelement.cpp
index 3a4e754fa31..e5c009308bf 100644
--- a/src/mongo/bson/bsonelement.cpp
+++ b/src/mongo/bson/bsonelement.cpp
@@ -590,99 +590,7 @@ BSONElement BSONElement::operator[](StringData field) const {
     return o[field];
 }
 
-int BSONElement::size(int maxLen) const {
-    if (totalSize >= 0)
-        return totalSize;
-
-    int remain = maxLen - fieldNameSize() - 1;
-
-    int x = 0;
-    switch (type()) {
-        case EOO:
-        case Undefined:
-        case jstNULL:
-        case MaxKey:
-        case MinKey:
-            break;
-        case mongo::Bool:
-            x = 1;
-            break;
-        case NumberInt:
-            x = 4;
-            break;
-        case bsonTimestamp:
-        case mongo::Date:
-        case NumberDouble:
-        case NumberLong:
-            x = 8;
-            break;
-        case NumberDecimal:
-            x = 16;
-            break;
-        case jstOID:
-            x = OID::kOIDSize;
-            break;
-        case Symbol:
-        case Code:
-        case mongo::String:
-            massert(
-                10313, "Insufficient bytes to calculate element size", maxLen == -1 || remain > 3);
-            x = valuestrsize() + 4;
-            break;
-        case CodeWScope:
-            massert(
-                10314, "Insufficient bytes to calculate element size", maxLen == -1 || remain > 3);
-            x = objsize();
-            break;
-
-        case DBRef:
-            massert(
-                10315, "Insufficient bytes to calculate element size", maxLen == -1 || remain > 3);
-            x = valuestrsize() + 4 + 12;
-            break;
-        case Object:
-        case mongo::Array:
-            massert(
-                10316, "Insufficient bytes to calculate element size", maxLen == -1 || remain > 3);
-            x = objsize();
-            break;
-        case BinData:
-            massert(
-                10317, "Insufficient bytes to calculate element size", maxLen == -1 || remain > 3);
-            x = valuestrsize() + 4 + 1 /*subtype*/;
-            break;
-        case RegEx: {
-            const char* p = value();
-            size_t len1 = (maxLen == -1) ? strlen(p) : strnlen(p, remain);
-            massert(10318, "Invalid regex string", maxLen == -1 || len1 < size_t(remain));
-            p = p + len1 + 1;
-            size_t len2;
-            if (maxLen == -1)
-                len2 = strlen(p);
-            else {
-                size_t x = remain - len1 - 1;
-                verify(x <= 0x7fffffff);
-                len2 = strnlen(p, x);
-                massert(10319, "Invalid regex options string", len2 < x);
-            }
-            x = (int)(len1 + 1 + len2 + 1);
-        } break;
-        default: {
-            StringBuilder ss;
-            ss << "BSONElement: bad type " << (int)type();
-            std::string msg = ss.str();
-            massert(13655, msg.c_str(), false);
-        }
-    }
-    totalSize = x + fieldNameSize() + 1;  // BSONType
-
-    return totalSize;
-}
-
-int BSONElement::size() const {
-    if (totalSize >= 0)
-        return totalSize;
-
+int BSONElement::computeSize() const {
     int x = 0;
     switch (type()) {
         case EOO:
@@ -740,9 +648,7 @@ int BSONElement::size() const {
             massert(10320, msg.c_str(), false);
         }
     }
-    totalSize = x + fieldNameSize() + 1;  // BSONType
-
-    return totalSize;
+    return x + fieldNameSize() + 1;  // BSONType
 }
 
 std::string BSONElement::toString(bool includeFieldName, bool full) const {
diff --git a/src/mongo/bson/bsonelement.h b/src/mongo/bson/bsonelement.h
index 7c427f38004..6ad90b3ba02 100644
--- a/src/mongo/bson/bsonelement.h
+++ b/src/mongo/bson/bsonelement.h
@@ -237,11 +237,12 @@ public:
         return type() == EOO;
     }
 
-    /** Size of the element.
-        @param maxLen If maxLen is specified, don't scan more than maxLen bytes to calculate size.
-    */
-    int size(int maxLen) const;
-    int size() const;
+    /**
+     * Size of the element.
+     */
+    int size() const {
+        return totalSize;
+    }
 
     /** Wrap this element up as a singleton object. */
     BSONObj wrap() const;
@@ -263,8 +264,6 @@ public:
      * NOTE: size includes the NULL terminator.
      */
     int fieldNameSize() const {
-        if (fieldNameSize_ == -1)
-            fieldNameSize_ = (int)strlen(fieldName()) + 1;
         return fieldNameSize_;
     }
 
@@ -653,27 +652,13 @@ public:
     }
 
     // @param maxLen don't scan more than maxLen bytes
-    explicit BSONElement(const char* d, int maxLen) : data(d) {
-        if (eoo()) {
-            totalSize = 1;
-            fieldNameSize_ = 0;
-        } else {
-            totalSize = -1;
-            fieldNameSize_ = -1;
-            if (maxLen != -1) {
-                size_t size = strnlen(fieldName(), maxLen - 1);
-                uassert(10333, "Invalid field name", size < size_t(maxLen - 1));
-                fieldNameSize_ = size + 1;
-            }
-        }
-    }
-
     explicit BSONElement(const char* d) : data(d) {
-        fieldNameSize_ = -1;
-        totalSize = -1;
         if (eoo()) {
             fieldNameSize_ = 0;
             totalSize = 1;
+        } else {
+            fieldNameSize_ = strlen(d + 1 /*skip type*/) + 1 /*include NUL byte*/;
+            totalSize = computeSize();
         }
     }
 
@@ -684,11 +669,19 @@ public:
      *  represent an EOO. You may pass -1 to indicate that you don't actually know the
      *  size.
      */
-    BSONElement(const char* d, int fieldNameSize, FieldNameSizeTag)
-        : data(d),
-          fieldNameSize_(fieldNameSize)  // internal size includes null terminator
-          ,
-          totalSize(-1) {}
+    BSONElement(const char* d, int fieldNameSize, FieldNameSizeTag) : data(d) {
+        if (eoo()) {
+            fieldNameSize_ = 0;
+            totalSize = 1;
+        } else {
+            if (fieldNameSize == -1) {
+                fieldNameSize_ = strlen(d + 1 /*skip type*/) + 1 /*include NUL byte*/;
+            } else {
+                fieldNameSize_ = fieldNameSize;
+            }
+            totalSize = computeSize();
+        }
+    }
 
     std::string _asCode() const;
 
@@ -697,9 +690,8 @@ public:
 
 private:
     const char* data;
-    mutable int fieldNameSize_;  // cached value
-
-    mutable int totalSize; /* caches the computed size */
+    int fieldNameSize_;  // internal size includes null terminator
+    int totalSize;
 
     friend class BSONObjIterator;
     friend class BSONObjStlIterator;
@@ -715,6 +707,9 @@ private:
         }
         return *this;
     }
+
+    // Only called from constructors.
+    int computeSize() const;
 };
 
 inline bool BSONElement::trueValue() const {
diff --git a/src/mongo/bson/bsonobj.cpp b/src/mongo/bson/bsonobj.cpp
index 5ec83659f63..6f357308727 100644
--- a/src/mongo/bson/bsonobj.cpp
+++ b/src/mongo/bson/bsonobj.cpp
@@ -608,7 +608,7 @@ void BSONObj::toString(
     bool first = true;
     while (1) {
         massert(10327, "Object does not end with EOO", i.moreWithEOO());
-        BSONElement e = i.next(true);
+        BSONElement e = i.next();
         massert(10328, "Invalid element size", e.size() > 0);
         massert(10329, "Element too large", e.size() < (1 << 30));
         int offset = (int)(e.rawdata() - this->objdata());
diff --git a/src/mongo/bson/bsonobj.h b/src/mongo/bson/bsonobj.h
index 7cd2a50870a..c83f1abf368 100644
--- a/src/mongo/bson/bsonobj.h
+++ b/src/mongo/bson/bsonobj.h
@@ -622,7 +622,7 @@ public:
     /** pre-increment */
     BSONObjStlIterator& operator++() {
         dassert(!_cur.eoo());
-        *this = BSONObjStlIterator(BSONElement(_cur.rawdata() + _curSize));
+        *this = BSONObjStlIterator(BSONElement(_cur.rawdata() + _cur.size()));
         return *this;
     }
 
@@ -648,13 +648,9 @@ public:
     }
 
 private:
-    explicit BSONObjStlIterator(BSONElement elem)
-        : _cur(elem),
-          // Make sure the lazy size fields are filled in.
-          _curSize(_cur.size()) {}
+    explicit BSONObjStlIterator(BSONElement elem) : _cur(elem) {}
 
     BSONElement _cur;
-    int _curSize = 0;  // TODO consider removing if we make the size cache in BSONElement eager.
 };
 
 /**
@@ -697,26 +693,6 @@ public:
         return _pos <= _theend;
     }
 
-    /**
-     * @return the next element in the object. For the final element, element.eoo() will be true.
-     */
-    BSONElement next(bool checkEnd) {
-        verify(_pos <= _theend);
-
-        int maxLen = -1;
-        if (checkEnd) {
-            maxLen = _theend + 1 - _pos;
-            verify(maxLen > 0);
-        }
-
-        BSONElement e(_pos, maxLen);
-        int esize = e.size(maxLen);
-        massert(16446, "BSONElement has bad size", esize > 0);
-        _pos += esize;
-
-        return e;
-    }
-
     BSONElement next() {
         verify(_pos <= _theend);
         BSONElement e(_pos);