SERVER-28758 Add appendPart and removeLastPart operations to FieldRef

1 files changed, 67 insertions, 17 deletions

diff --git a/src/mongo/db/field_ref.cpp b/src/mongo/db/field_ref.cpp
index e3bc5305c0b..3a759febe3c 100644
--- a/src/mongo/db/field_ref.cpp
+++ b/src/mongo/db/field_ref.cpp
@@ -34,7 +34,9 @@
 
 namespace mongo {
 
-FieldRef::FieldRef() : _size(0) {}
+const size_t FieldRef::kReserveAhead;
+
+FieldRef::FieldRef() : _size(0), _cachedSize(0) {}
 
 FieldRef::FieldRef(StringData path) : _size(0) {
     parse(path);
@@ -74,9 +76,9 @@ void FieldRef::parse(StringData path) {
         // instead reach the break statement.
 
         if (cur != beg)
-            appendPart(StringData(&*beg, cur - beg));
+            appendParsedPart(StringData(&*beg, cur - beg));
         else
-            appendPart(StringData());
+            appendParsedPart(StringData());
 
         if (cur != end) {
             beg = ++cur;
@@ -90,25 +92,59 @@ void FieldRef::parse(StringData path) {
 void FieldRef::setPart(size_t i, StringData part) {
     dassert(i < _size);
 
-    if (_replacements.size() != _size) {
+    if (_replacements.empty()) {
         _replacements.resize(_size);
     }
 
     _replacements[i] = part.toString();
     if (i < kReserveAhead) {
-        _fixed[i] = _replacements[i];
+        _fixed[i] = boost::none;
     } else {
-        _variable[getIndex(i)] = _replacements[i];
+        _variable[getIndex(i)] = boost::none;
+    }
+}
+
+void FieldRef::appendPart(StringData part) {
+    if (_replacements.empty()) {
+        _replacements.resize(_size);
+    }
+
+    _replacements.push_back(part.toString());
+
+    if (_size < kReserveAhead) {
+        _fixed[_size] = boost::none;
+    } else {
+        _variable.push_back(boost::none);
+    }
+    _size++;
+}
+
+void FieldRef::removeLastPart() {
+    if (_size == 0) {
+        return;
+    }
+
+    if (!_replacements.empty()) {
+        _replacements.pop_back();
     }
+
+    if (_size > kReserveAhead) {
+        dassert(_variable.size() == _size - kReserveAhead);
+        _variable.pop_back();
+    }
+
+    _size--;
 }
 
-size_t FieldRef::appendPart(StringData part) {
+size_t FieldRef::appendParsedPart(StringData part) {
     if (_size < kReserveAhead) {
         _fixed[_size] = part;
     } else {
         _variable.push_back(part);
     }
-    return ++_size;
+    _size++;
+    _cachedSize++;
+    return _size;
 }
 
 void FieldRef::reserialize() const {
@@ -116,7 +152,7 @@ void FieldRef::reserialize() const {
     // Reserve some space in the string. We know we will have, at minimum, a character for
     // each component we are writing, and a dot for each component, less one. We don't want
     // to reserve more, since we don't want to forfeit the SSO if it is applicable.
-    nextDotted.reserve((_size * 2) - 1);
+    nextDotted.reserve((_size > 0) ? (_size * 2) - 1 : 0);
 
     // Concatenate the fields to a new string
     for (size_t i = 0; i != _size; ++i) {
@@ -129,13 +165,24 @@ void FieldRef::reserialize() const {
     // Make the new string our contents
     _dotted.swap(nextDotted);
 
+    // Before we reserialize, it's possible that _cachedSize != _size because parts were added or
+    // removed. This reserialization process reconciles the components in our cached string
+    // (_dotted) with the modified path.
+    _cachedSize = _size;
+
     // Fixup the parts to refer to the new string
     std::string::const_iterator where = _dotted.begin();
     const std::string::const_iterator end = _dotted.end();
     for (size_t i = 0; i != _size; ++i) {
-        StringData& part = (i < kReserveAhead) ? _fixed[i] : _variable[getIndex(i)];
-        const size_t size = part.size();
-        part = StringData(&*where, size);
+        boost::optional<StringData>& part =
+            (i < kReserveAhead) ? _fixed[i] : _variable[getIndex(i)];
+        const size_t size = part ? part.get().size() : _replacements[i].size();
+
+        // There is one case where we expect to see the "where" iterator to be at "end" here: we
+        // are at the last part of the FieldRef and that part is the empty string. In that case, we
+        // need to make sure we do not dereference the "where" iterator.
+        dassert(where != end || (size == 0 && i == _size - 1));
+        part = StringData((size > 0) ? &*where : nullptr, size);
         where += size;
         // skip over '.' unless we are at the end.
         if (where != end) {
@@ -151,10 +198,12 @@ void FieldRef::reserialize() const {
 StringData FieldRef::getPart(size_t i) const {
     dassert(i < _size);
 
-    if (i < kReserveAhead) {
-        return _fixed[i];
+    const boost::optional<StringData>& part =
+        (i < kReserveAhead) ? _fixed[i] : _variable[getIndex(i)];
+    if (part) {
+        return part.get();
     } else {
-        return _variable[getIndex(i)];
+        return StringData(_replacements[i]);
     }
 }
 
@@ -199,9 +248,9 @@ StringData FieldRef::dottedSubstring(size_t startPart, size_t endPart) const {
     if (_size == 0 || startPart >= endPart || endPart > numParts())
         return StringData();
 
-    if (!_replacements.empty())
+    if (!_replacements.empty() || _size != _cachedSize)
         reserialize();
-    dassert(_replacements.empty());
+    dassert(_replacements.empty() && _size == _cachedSize);
 
     StringData result(_dotted);
 
@@ -271,6 +320,7 @@ int FieldRef::compare(const FieldRef& other) const {
 
 void FieldRef::clear() {
     _size = 0;
+    _cachedSize = 0;
     _variable.clear();
     _dotted.clear();
     _replacements.clear();